Allow overriding default contract addresses (#50)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vana"
-version = "0.32.0"
+version = "0.33.0"
 description = ""
 authors = ["Tim Nunamaker <[email protected]>", "Volodymyr Isai <[email protected]>", "Kahtaf Alam <[email protected]>"]
 readme = "README.md"

vana/__init__.py

@@ -15,7 +15,7 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-__version__ = "0.32.0"
+__version__ = "0.33.0"
 
 import rich
 

vana/client.py

@@ -2,9 +2,8 @@
 import copy
 import json
 import os
-from typing import Optional
-
 import vana
+from typing import Optional
 from vana.chain_data import Proof, ProofData
 from vana.contracts import contracts
 from vana.utils.web3 import as_wad
@@ -14,9 +13,32 @@ class Client:
     @staticmethod
     def config() -> "config":
         parser = argparse.ArgumentParser()
-        vana.ChainManager.add_args(parser)
+        vana.Client.add_args(parser)
         return vana.Config(parser, args=[])
 
+    @classmethod
+    def add_args(cls, parser: argparse.ArgumentParser, prefix: Optional[str] = None):
+        prefix_str = "" if prefix is None else f"{prefix}."
+        try:
+            parser.add_argument(
+                "--" + prefix_str + "client.tee_pool_contract_address",
+                default=os.getenv("TEE_POOL_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the TEE Pool Contract.""")
+            parser.add_argument(
+                "--" + prefix_str + "client.data_registry_contract_address",
+                default=os.getenv("DATA_REGISTRY_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the Data Registry Contract.""")
+            parser.add_argument(
+                "--" + prefix_str + "client.dlp_root_contract_address",
+                default=os.getenv("DLP_ROOT_CONTRACT_ADDRESS") or None,
+                type=str,
+                help="""The address for the DLP Root Contract.""")
+        except argparse.ArgumentError:
+            # re-parsing arguments.
+            pass
+
     def __init__(self, config: vana.Config):
         if config is None:
             config = self.config()
@@ -33,7 +55,7 @@ def __init__(self, config: vana.Config):
         )
         with open(data_registry_contract_path) as f:
             self.data_registry_contract = self.chain_manager.web3.eth.contract(
-                address=contracts[self.network]["DataRegistry"],
+                address=config.client.data_registry_contract_address or contracts[self.network]["DataRegistry"],
                 abi=json.load(f)
             )
         tee_pool_contract_path = os.path.join(
@@ -42,7 +64,7 @@ def __init__(self, config: vana.Config):
         )
         with open(tee_pool_contract_path) as f:
             self.tee_pool_contract = self.chain_manager.web3.eth.contract(
-                address=contracts[self.network]["TeePool"],
+                address=config.client.tee_pool_contract_address or contracts[self.network]["TeePool"],
                 abi=json.load(f)
             )
 

vana/contracts/__init__.py

@@ -2,12 +2,12 @@
     "vana": {
         "TeePool": "0x3c92fD91639b41f13338CE62f19131e7d19eaa0D",
         "DataRegistry": "0x8C8788f98385F6ba1adD4234e551ABba0f82Cb7C",
-        "RootNetworkContract": "0xf408A064d640b620219F510963646Ed2bD5606BB",
+        "RootNetworkContract": "0xff14346dF2B8Fd0c95BF34f1c92e49417b508AD5",
     },
     "islander": {
         "TeePool": "0x3c92fD91639b41f13338CE62f19131e7d19eaa0D",
         "DataRegistry": "0x8C8788f98385F6ba1adD4234e551ABba0f82Cb7C",
-        "RootNetworkContract": "0xf408A064d640b620219F510963646Ed2bD5606BB",
+        "RootNetworkContract": "0xff14346dF2B8Fd0c95BF34f1c92e49417b508AD5",
     },
     "maya": {
         "TeePool": "0xF084Ca24B4E29Aa843898e0B12c465fAFD089965",