feat: allow changing the registry address (#45)

Author: antazoey

README.md

@@ -165,3 +165,25 @@ Configuring entries is useful for:
 1. Testing in the `local` network.
 2. Attaining faster performance (no Ethereum call).
 3. Avoiding connecting to Ethereum mainnet.
+
+### Change Registry
+
+Change the registry contract address by configuring it in your `pyproject.toml`:
+
+```toml
+[tool.ape.ens]
+registry_address = "0x123..."
+```
+
+or using `ape-config.yaml`:
+
+```yaml
+ens:
+  registry_address: "0x123..."
+```
+
+You can also switch the registry adhoc during CLI commands:
+
+```shell
+ape ens resolve vitalik.eth --registry-address 0x123...311
+```

ape_ens/_cli.py

@@ -2,6 +2,8 @@
 
 import click
 from ape.cli import ApeCliContextObject, ape_cli_context, network_option
+from ape.exceptions import ConversionError
+from ape.types.address import AddressType
 from eth_utils import to_hex
 
 from ape_ens.ens import ENS
@@ -25,15 +27,40 @@ def cli():
     """
 
 
+def registry_address_option(**kwargs):
+    if "callback" not in kwargs:
+
+        def validate_address(ctx, param, value):
+            if value is None:
+                return value
+
+            from ape import convert
+
+            try:
+                return convert(value, AddressType)
+            except ConversionError:
+                raise click.BadOptionUsage(
+                    "--registry-address", f"Invalid `--registry-address` {value}."
+                )
+
+        kwargs["callback"] = validate_address
+
+    if "help" not in kwargs:
+        kwargs["help"] = "ENS registry address"
+
+    return click.option("--registry-address", **kwargs)
+
+
 @cli.command()
 @ape_cli_context(obj_type=ENSContext)
 @click.argument("name")
 @network_option(default=None)
-def resolve(cli_ctx, name):
+@registry_address_option()
+def resolve(cli_ctx, name, registry_address):
     """
     Resolve an ENS address.
     """
-    if address := cli_ctx.ens.resolve(name):
+    if address := cli_ctx.ens.resolve(name, registry_address=registry_address):
         click.echo(address)
     else:
         click.echo(f"Could not resolve ENS '{name}'.", err=True)
@@ -43,11 +70,12 @@ def resolve(cli_ctx, name):
 @ape_cli_context(obj_type=ENSContext)
 @click.argument("address")
 @network_option(default=None)
-def name_cmd(cli_ctx, address):
+@registry_address_option()
+def name_cmd(cli_ctx, address, registry_address):
     """
     Get the ENS of an address.
     """
-    if name := cli_ctx.ens.name(address):
+    if name := cli_ctx.ens.name(address, registry_address=registry_address):
         click.echo(name)
     else:
         click.echo(f"No ENS name found for '{address}'.", err=True)
@@ -57,11 +85,12 @@ def name_cmd(cli_ctx, address):
 @ape_cli_context(obj_type=ENSContext)
 @click.argument("name")
 @network_option(default=None)
-def owner(cli_ctx, name):
+@registry_address_option()
+def owner(cli_ctx, name, registry_address):
     """
     Get the owner of an ENS domain.
     """
-    if owner_address := cli_ctx.ens.owner(name):
+    if owner_address := cli_ctx.ens.owner(name, registry_address=registry_address):
         click.echo(owner_address)
     else:
         click.echo(f"No owner found for '{name}'.", err=True)

ape_ens/config.py

@@ -1,4 +1,7 @@
+from typing import Optional
+
 from ape.api import PluginConfig
+from ape.types import AddressType
 
 
 class ENSConfig(PluginConfig):
@@ -19,3 +22,9 @@ class ENSConfig(PluginConfig):
     performance and in testing environments without access
     to Ethereum mainnet.
     """
+
+    registry_address: Optional[AddressType] = None
+    """
+    Configure the registry address if it different than the default
+    Ethereum mainnet address.
+    """

ape_ens/converter.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Any
 
 from ape.api import ConverterAPI
+from ape.exceptions import ConversionError
 
 from ape_ens.ens import ENS
 
@@ -40,7 +41,13 @@ def is_convertible(self, value: Any) -> bool:
         if not isinstance(value, str):
             return False
 
-        return self.ens.can_resolve(value)
+        try:
+            return self.ens.can_resolve(value)
+        except Exception:
+            return False
 
     def convert(self, value: str) -> "AddressType":
-        return self.ens.resolve(value)
+        try:
+            return self.ens.resolve(value)
+        except Exception as err:
+            raise ConversionError(str(err)) from err

ape_ens/ens.py

@@ -4,9 +4,10 @@
 from ape.exceptions import ProviderError
 from ape.logging import logger
 from ape.utils.basemodel import ManagerAccessMixin
-from web3.exceptions import CannotHandleRequest
+from web3.exceptions import BadFunctionCallOutput, CannotHandleRequest, Web3RPCError
 from web3.main import ENS as Web3ENS
 
+from ape_ens.exceptions import MissingRegistryError
 from ape_ens.utils.namehash import namehash
 
 if TYPE_CHECKING:
@@ -37,7 +38,7 @@ class ENS(ManagerAccessMixin):
     """
 
     def __init__(self, backend: Optional["Web3ENS"] = None) -> None:
-        self._ens = backend
+        self.__initialized_ens = backend
         self.local_registry: dict[str, AddressType] = {}
 
     @classmethod
@@ -112,16 +113,28 @@ def _get_mainnet_provider(self) -> "Web3Provider":
 
     @cached_property
     def _web3_ens(self) -> "Web3ENS":
-        if ens := self._ens:
+        if ens := self.__initialized_ens:
             # Initialized with ENS (testing?)
             return ens
 
-        return self._mainnet_provider.web3.ens
+        return self._create_web3_ens()
 
     @property
     def config(self) -> "ENSConfig":
         return self.config_manager.ens
 
+    def _create_web3_ens(self, registry_address: Optional["AddressType"] = None) -> "Web3ENS":
+        if registry_address:
+            return Web3ENS.from_web3(self._mainnet_provider.web3, registry_address)
+
+        else:
+            # Check config.
+            if address := self.config.registry_address:
+                return Web3ENS.from_web3(self._mainnet_provider.web3, address)
+
+        # Use default (most common).
+        return self._mainnet_provider.web3.ens
+
     def can_resolve(self, name: str) -> bool:
         """
         Returns ``True`` when ENS can resolve the name.
@@ -144,17 +157,29 @@ def can_resolve(self, name: str) -> bool:
 
         return address is not None
 
-    def resolve(self, name: str, use_cache: Optional[bool] = None) -> Optional["AddressType"]:
+    def resolve(
+        self,
+        name: str,
+        use_cache: Optional[bool] = None,
+        registry_address: Optional["AddressType"] = None,
+    ) -> Optional["AddressType"]:
         """
         Resolve an ENS name.
 
         Args:
             name (str): The name to resolve.
             use_cache (bool): Set to ``False`` to not use the in-memory cache.
+            registry_address (Optional[AddressType]): Optionally, change the registry
+              address.
 
         Returns:
             AddressType | None
         """
+        ens = (
+            self._create_web3_ens(registry_address=registry_address)
+            if registry_address
+            else self._web3_ens
+        )
         if use_cache is None:
             # Use default from config.
             use_cache = self.config.use_cache
@@ -168,37 +193,55 @@ def resolve(self, name: str, use_cache: Optional[bool] = None) -> Optional["Addr
                 self.local_registry[name] = address
                 return address
 
-        if address := self._web3_ens.address(name):
-            if use_cache:
-                self.local_registry[name] = address
+        try:
+            address = ens.address(name)
+        except (Web3RPCError, BadFunctionCallOutput) as err:
+            raise MissingRegistryError(str(err))
 
-            return address
+        if use_cache and address is not None:
+            self.local_registry[name] = address
 
-        return None
+        return address
 
-    def name(self, address: "AddressType") -> Optional[str]:
+    def name(
+        self, address: "AddressType", registry_address: Optional["AddressType"] = None
+    ) -> Optional[str]:
         """
         Reverse look-up an address to get the ENS name.
 
         Args:
             address (AddressType): The address to resolve.
+            registry_address (Optional[AddressType]): Optionally, change the registry.
 
         Returns:
             str | None: The ENS name.
         """
-        return self._web3_ens.name(address)
-
-    def owner(self, name: str) -> Optional["AddressType"]:
+        ens = (
+            self._create_web3_ens(registry_address=registry_address)
+            if registry_address
+            else self._web3_ens
+        )
+        return ens.name(address)
+
+    def owner(
+        self, name: str, registry_address: Optional["AddressType"] = None
+    ) -> Optional["AddressType"]:
         """
         Get the owner of an ENS domain.
 
         Args:
             name (str): The ENS name to check.
+            registry_address (Optional[AddressType]): Optionally, change the registry.
 
         Returns:
             AddressType | None
         """
-        return self._web3_ens.owner(name)
+        ens = (
+            self._create_web3_ens(registry_address=registry_address)
+            if registry_address
+            else self._web3_ens
+        )
+        return ens.owner(name)
 
     def namehash(self, name: str) -> "HexBytes":
         """

ape_ens/exceptions.py

@@ -0,0 +1,13 @@
+from ape.exceptions import ApeException
+
+
+class ApeENSException(ApeException):
+    """
+    Base class for exceptions in the `ape-ens` package.
+    """
+
+
+class MissingRegistryError(ApeENSException):
+    """
+    Raised when ape-ens detect the registry is missing.
+    """

tests/test_cli.py

@@ -15,11 +15,16 @@ def runner():
     return CliRunner()
 
 
-def test_resolve(networks, runner):
+def test_resolve(runner):
     result = runner.invoke(cli, ["resolve", "vitalik.eth"])
     assert "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045" in result.output, result.output
 
 
+def test_resolve_invalid_registry_address(runner):
+    result = runner.invoke(cli, ["resolve", "vitalik.eth", "--registry-address", "asdf"])
+    assert result.exit_code != 0
+
+
 def test_name(runner):
     result = runner.invoke(cli, ["name", "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045"])
     assert "vitalik.eth" in result.output, result.output

tests/test_converter.py

@@ -1,4 +1,5 @@
 import pytest
+from ape.exceptions import ConversionError
 from web3.exceptions import CannotHandleRequest
 
 from tests.conftest import negative_tests
@@ -63,6 +64,17 @@ def test_convert_after_adding_to_local_registry(converter, vitalik, accounts):
     assert actual == dev_account.address
 
 
+def test_convert_change_registry_address(project, converter, vitalik, accounts):
+    fake_registry = accounts[0].address
+    ens = converter.ens
+    converter._ens = None  # Pretend this is the start of the session.
+    with project.temp_config(ens={"registry_address": fake_registry}):
+        with pytest.raises(ConversionError):
+            _ = converter.convert("apepython.eth")
+
+    converter._ens = ens
+
+
 def test_address_cache(converter, address):
     converter.address_cache["test.eth"] = address
     assert converter.convert("test.eth") == address