rename erc20Wrapper to tokenWrapper to be more generic in case we end up supporting another spec

1 files changed, 17 insertions, 3 deletions

diff --git a/src/contract_wrappers/erc20_wrapper.ts b/src/contract_wrappers/token_wrapper.ts
index 60b1887db..bd815554a 100644
--- a/src/contract_wrappers/erc20_wrapper.ts
+++ b/src/contract_wrappers/token_wrapper.ts
@@ -6,9 +6,14 @@ import {ContractWrapper} from './contract_wrapper';
 import * as TokenArtifacts from '../artifacts/Token.json';
 import {ERC20Contract} from '../types';
 
-export class ERC20Wrapper extends ContractWrapper {
+export class TokenWrapper extends ContractWrapper {
+    private tokenContractsByAddress: {[address: string]: ERC20Contract};
     constructor(web3Wrapper: Web3Wrapper) {
         super(web3Wrapper);
+        this.tokenContractsByAddress = {};
+    }
+    public invalidateContractInstances() {
+        this.tokenContractsByAddress = {};
     }
     /**
      * Returns an owner's ERC20 token balance
@@ -17,12 +22,21 @@ export class ERC20Wrapper extends ContractWrapper {
         assert.isETHAddressHex('ownerAddress', ownerAddress);
         assert.isETHAddressHex('tokenAddress', tokenAddress);
 
-        const contractInstance = await this.instantiateContractIfExistsAsync((TokenArtifacts as any), tokenAddress);
-        const tokenContract = contractInstance as ERC20Contract;
+        const tokenContract = await this.getTokenContractAsync(tokenAddress);
         let balance = await tokenContract.balanceOf.call(ownerAddress);
         // The BigNumber instance returned by Web3 is of a much older version then our own, we therefore
         // should always re-instantiate the returned BigNumber after retrieval.
         balance = _.isUndefined(balance) ? new BigNumber(0) : new BigNumber(balance);
         return balance;
     }
+    private async getTokenContractAsync(tokenAddress: string): Promise<ERC20Contract> {
+        let tokenContract = this.tokenContractsByAddress[tokenAddress];
+        if (!_.isUndefined(tokenContract)) {
+            return tokenContract;
+        }
+        const contractInstance = await this.instantiateContractIfExistsAsync((TokenArtifacts as any), tokenAddress);
+        tokenContract = contractInstance as ERC20Contract;
+        this.tokenContractsByAddress[tokenAddress] = tokenContract;
+        return tokenContract;
+    }
 }