Allow Python imports to be aliased like a Basilisp require (#319)

* Allow Python imports to be aliased like a Basilisp require

* Docstrings

* Formatting

Author: chrisrink10

src/basilisp/core/__init__.lpy

@@ -1159,6 +1159,22 @@
        (finally
          (. (var ~vvar) ~'pop-bindings)))))
 
+(import* [time :as py-time])
+
+(defn ^:private perf-counter
+  []
+  (py-time/perf-counter))
+
+(defmacro time
+  "Time the execution of expr. Return the result of expr and print the
+  time execution took in milliseconds."
+  [expr]
+  `(let [start (perf-counter)]
+     (try
+       ~expr
+       (finally
+         (println (* 1000 (- (perf-counter) start)) "msecs")))))
+
 ;;;;;;;;;;;;;;;;;;;;;;
 ;; Threading Macros ;;
 ;;;;;;;;;;;;;;;;;;;;;;

src/basilisp/lang/compiler.py

@@ -328,8 +328,8 @@ def quoted(self):
         yield
         self._is_quoted.pop()
 
-    def add_import(self, imp: sym.Symbol, mod: types.ModuleType):
-        self.current_ns.add_import(imp, mod)
+    def add_import(self, imp: sym.Symbol, mod: types.ModuleType, *aliases: sym.Symbol):
+        self.current_ns.add_import(imp, mod, *aliases)
 
     @property
     def imports(self) -> lmap.Map:
@@ -1272,35 +1272,52 @@ def _if_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
 def _import_ast(ctx: CompilerContext, form: llist.List) -> ASTStream:
     """Append Import statements into the compiler context nodes."""
     assert form.first == _IMPORT
-    assert all([isinstance(f, sym.Symbol) for f in form.rest])
 
     last = None
-    for s in form.rest:
+    for f in form.rest:
+        if isinstance(f, sym.Symbol):
+            module_name = f
+            module_alias = module_name.name.split(".", maxsplit=1)[0]
+        elif isinstance(f, vec.Vector):
+            module_name = f.entry(0)
+            assert isinstance(
+                module_name, sym.Symbol
+            ), "Python module name must be a symbol"
+            assert kw.keyword("as") == f.entry(1)
+            module_alias = f.entry(2).name
+        else:
+            raise CompilerException("Symbol or vector expected for import*")
+
         try:
-            module = importlib.import_module(s.name)
-            ctx.add_import(s, module)
+            module = importlib.import_module(module_name.name)
+            if module_name.name != module_alias:
+                ctx.add_import(module_name, module, sym.symbol(module_alias))
+            else:
+                ctx.add_import(module_name, module)
         except ModuleNotFoundError:
-            raise ImportError(f"Module '{s.name}' not found")
+            raise ImportError(f"Module '{module_name.name}' not found")
 
         with ctx.quoted():
-            module_name = s.name.split(".", maxsplit=1)[0]
-            yield _dependency(ast.Global(names=[module_name]))
+            module_alias = munge(module_alias)
+            yield _dependency(ast.Global(names=[module_alias]))
             yield _dependency(
                 ast.Assign(
-                    targets=[ast.Name(id=module_name, ctx=ast.Store())],
+                    targets=[ast.Name(id=module_alias, ctx=ast.Store())],
                     value=ast.Call(
                         func=_load_attr("builtins.__import__"),
-                        args=[ast.Str(s.name)],
+                        args=[ast.Str(module_name.name)],
                         keywords=[],
                     ),
                 )
             )
-            last = ast.Name(id=module_name, ctx=ast.Load())
+            last = ast.Name(id=module_alias, ctx=ast.Load())
             yield _dependency(
                 ast.Call(
                     func=_load_attr(f"{_NS_VAR_VALUE}.add_import"),
                     args=list(
-                        chain(_unwrap_nodes(_to_ast(ctx, s)), [last])  # type: ignore
+                        chain(  # type: ignore
+                            _unwrap_nodes(_to_ast(ctx, module_name)), [last]
+                        )
                     ),
                     keywords=[],
                 )
@@ -2046,17 +2063,31 @@ def _resolve_sym(ctx: CompilerContext, form: sym.Symbol) -> Optional[str]:  # no
             if v is not None:
                 return _resolve_sym_var(ctx, v)
         ns_sym = sym.symbol(form.ns)
-        if ns_sym in ctx.current_ns.imports:
+        if ns_sym in ctx.current_ns.imports or ns_sym in ctx.current_ns.import_aliases:
             # We still import Basilisp code, so we'll want to make sure
             # that the symbol isn't referring to a Basilisp Var first
             v = Var.find(form)
             if v is not None:
                 return _resolve_sym_var(ctx, v)
 
+            # Python modules imported using `import*` may be imported
+            # with an alias, which sets the local name of the alias to the
+            # alias value (much like Python's `from module import member`
+            # statement). In this case, we'll want to check for the module
+            # using its non-aliased name, but then generate code to access
+            # the member using the alias.
+            if ns_sym in ctx.current_ns.import_aliases:
+                module_name: sym.Symbol = ctx.current_ns.import_aliases[ns_sym]
+            else:
+                module_name = ns_sym
+
             # Otherwise, try to direct-link it like a Python variable
+            safe_module_name = munge(module_name.name)
+            assert (
+                safe_module_name in sys.modules
+            ), f"Module '{safe_module_name}' is not imported"
+            ns_module = sys.modules[safe_module_name]
             safe_ns = munge(form.ns)
-            assert safe_ns in sys.modules, f"Module '{safe_ns}' is not imported"
-            ns_module = sys.modules[safe_ns]
 
             # Try without allowing builtins first
             safe_name = munge(form.name)

src/basilisp/lang/runtime.py

@@ -305,7 +305,15 @@ class Namespace:
 
     _NAMESPACES = atom.Atom(lmap.Map.empty())
 
-    __slots__ = ("_name", "_module", "_interns", "_refers", "_aliases", "_imports")
+    __slots__ = (
+        "_name",
+        "_module",
+        "_interns",
+        "_refers",
+        "_aliases",
+        "_imports",
+        "_import_aliases",
+    )
 
     def __init__(self, name: sym.Symbol, module: types.ModuleType = None) -> None:
         self._name = name
@@ -319,6 +327,7 @@ def __init__(self, name: sym.Symbol, module: types.ModuleType = None) -> None:
                 .to_dict()
             )
         )
+        self._import_aliases: atom.Atom = atom.Atom(lmap.Map.empty())
         self._interns: atom.Atom = atom.Atom(lmap.Map.empty())
         self._refers: atom.Atom = atom.Atom(lmap.Map.empty())
 
@@ -360,6 +369,11 @@ def imports(self) -> lmap.Map:
         namespace."""
         return self._imports.deref()
 
+    @property
+    def import_aliases(self) -> lmap.Map:
+        """A mapping of a symbolic alias and a Python module name."""
+        return self._import_aliases.deref()
+
     @property
     def interns(self) -> lmap.Map:
         """A mapping between a symbolic name and a Var. The Var may point to
@@ -415,14 +429,32 @@ def find(self, sym: sym.Symbol) -> Optional[Var]:
             return self.refers.entry(sym, None)
         return v
 
-    def add_import(self, sym: sym.Symbol, module: types.ModuleType) -> None:
-        """Add the Symbol as an imported Symbol in this Namespace."""
+    def add_import(
+        self, sym: sym.Symbol, module: types.ModuleType, *aliases: sym.Symbol
+    ) -> None:
+        """Add the Symbol as an imported Symbol in this Namespace. If aliases are given,
+        the aliases will be applied to the """
         self._imports.swap(lambda m: m.assoc(sym, module))
+        if aliases:
+            self._import_aliases.swap(
+                lambda m: m.assoc(
+                    *itertools.chain.from_iterable([(alias, sym) for alias in aliases])
+                )
+            )
 
     def get_import(self, sym: sym.Symbol) -> Optional[types.ModuleType]:
         """Return the module if a moduled named by sym has been imported into
-        this Namespace, None otherwise."""
-        return self.imports.entry(sym, None)
+        this Namespace, None otherwise.
+
+        First try to resolve a module directly with the given name. If no module
+        can be resolved, attempt to resolve the module using import aliases."""
+        mod = self.imports.entry(sym, None)
+        if mod is None:
+            alias = self.import_aliases.get(sym, None)
+            if alias is None:
+                return None
+            return self.imports.entry(alias, None)
+        return mod
 
     def add_refer(self, sym: sym.Symbol, var: Var) -> None:
         """Refer var in this namespace under the name sym."""

tests/basilisp/compiler_test.py

@@ -593,6 +593,24 @@ def test_truthiness(ns: runtime.Namespace):
     assert kw.keyword("a") == lcompile("(if #{true} :a :b)")
 
 
+def test_import(ns: runtime.Namespace):
+    with pytest.raises(compiler.CompilerException):
+        lcompile("(import* :time)")
+
+    with pytest.raises(compiler.CompilerException):
+        lcompile('(import* "time")')
+
+    with pytest.raises(ImportError):
+        lcompile("(import* real.fake.module)")
+
+    import time
+
+    assert time.perf_counter == lcompile("(import* time) time/perf-counter")
+    assert time.perf_counter == lcompile(
+        "(import* [time :as py-time]) py-time/perf-counter"
+    )
+
+
 def test_interop_new(ns: runtime.Namespace):
     assert "hi" == lcompile('(builtins.str. "hi")')
     assert "1" == lcompile("(builtins.str. 1)")