Merge pull request #149 from charles-cooper/native-imports

implement "native" imports for boa

boa/interpret.py

@@ -1,4 +1,6 @@
+import importlib
 import json
+import sys
 import textwrap
 from pathlib import Path
 from typing import Any, Union
@@ -24,6 +26,37 @@
 _disk_cache = None
 
 
+class BoaImporter(importlib.abc.MetaPathFinder):
+    def find_module(self, fullname, package_path, target=None):
+        # Return a loader
+        return self
+
+    def load_module(self, fullname):
+        # Return a module
+        if fullname in sys.modules:
+            return sys.modules[fullname]
+
+        path = Path(fullname.replace(".", "/")).with_suffix(".vy")
+        for prefix in sys.path:
+            to_try = Path(prefix) / path
+            try:
+                ret = load_partial(to_try)
+                break
+            except (FileNotFoundError, NotADirectoryError):
+                pass
+        else:
+            raise ImportError(fullname)
+
+        # comply with PEP-302:
+        ret.__name__ = to_try.name
+        ret.__file__ = str(to_try)
+        sys.modules[fullname] = ret
+        return ret
+
+
+sys.meta_path.append(BoaImporter())
+
+
 def set_cache_dir(cache_dir="~/.cache/titanoboa"):
     global _disk_cache
     if cache_dir is None:

tests/unitary/test_import_overloading.py

@@ -0,0 +1,57 @@
+import contextlib
+import os
+import sys
+
+from boa.interpret import VyperDeployer
+
+
+@contextlib.contextmanager
+def mock_sys_path(path):
+    anchor = sys.path
+    try:
+        sys.path = [path]
+        yield
+    finally:
+        sys.path = anchor
+
+
+@contextlib.contextmanager
+def workdir(path):
+    tmp = os.getcwd()
+    try:
+        os.chdir(path)
+        with mock_sys_path("."):
+            yield
+    finally:
+        os.chdir(tmp)
+
+
+def test_imports(tmp_path):
+    code = """
+totalSupply: public(uint256)
+
+@external
+def __init__(initial_supply: uint256):
+    self.totalSupply = initial_supply
+    """
+
+    filepath = tmp_path / "foo" / "bar.vy"
+    filepath.parent.mkdir(parents=True)
+
+    with filepath.open("w") as f:
+        f.write(code)
+
+    # note that `with mock_sys_path(tmp_path)` does not work here!
+    # apparently, there are different semantics for module loading
+    # depending on if we are in the current directory or not.
+    with workdir(tmp_path):
+        from foo import bar
+
+        assert isinstance(bar, VyperDeployer)
+        contract = bar.deploy(100)
+        assert contract.totalSupply() == 100
+
+        from foo import bar as baz
+
+        assert isinstance(baz, VyperDeployer)
+        assert baz is bar