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

2 files changed, 30 insertions, 7 deletions

diff --git a/src/0x.js.ts b/src/0x.js.ts
index 7932559fb..96b290439 100644
--- a/src/0x.js.ts
+++ b/src/0x.js.ts
@@ -13,16 +13,16 @@ import compareVersions = require('compare-versions');
 import {ExchangeWrapper} from './contract_wrappers/exchange_wrapper';
 import {TokenRegistryWrapper} from './contract_wrappers/token_registry_wrapper';
 import {ecSignatureSchema} from './schemas/ec_signature_schema';
-import {ERC20Wrapper} from './contract_wrappers/erc20_wrapper';
+import {TokenWrapper} from './contract_wrappers/token_wrapper';
 import {SolidityTypes, ECSignature, ZeroExError} from './types';
 
 const MAX_DIGITS_IN_UNSIGNED_256_INT = 78;
 
 export class ZeroEx {
-    public web3Wrapper: Web3Wrapper;
     public exchange: ExchangeWrapper;
     public tokenRegistry: TokenRegistryWrapper;
-    public erc20: ERC20Wrapper;
+    public erc20: TokenWrapper;
+    private web3Wrapper: Web3Wrapper;
     /**
      * Computes the orderHash given the order parameters and returns it as a hex encoded string.
      */
@@ -137,7 +137,16 @@ export class ZeroEx {
         this.web3Wrapper = new Web3Wrapper(web3);
         this.exchange = new ExchangeWrapper(this.web3Wrapper);
         this.tokenRegistry = new TokenRegistryWrapper(this.web3Wrapper);
-        this.erc20 = new ERC20Wrapper(this.web3Wrapper);
+        this.erc20 = new TokenWrapper(this.web3Wrapper);
+    }
+    /**
+     * Sets a new provider for the web3 instance used by 0x.js
+     */
+    public setProvider(provider: Web3.Provider) {
+        this.web3Wrapper.setProvider(provider);
+        this.exchange.invalidateContractInstance();
+        this.tokenRegistry.invalidateContractInstance();
+        this.erc20.invalidateContractInstances();
     }
     /**
      * Signs an orderHash and returns it's elliptic curve signature
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;
+    }
 }