Merge pull request #30 from jmanuel1/runtime-import-resolution

concat/__main__.py

@@ -56,8 +56,9 @@ def get_line_at(file: TextIO, location: concat.astutils.Location) -> str:
     help='tokenize input from the given file and print the tokens as a JSON array',
 )
 
-# We should pass any unknown args onto the program we're about to run. There
-# might be a better way to go about this, but I think this is fine for now.
+# We should pass any unknown args onto the program we're about to run.
+# FIXME: There might be a better way to go about this, but I think this is fine
+# for now.
 args, rest = arg_parser.parse_known_args()
 sys.argv = [sys.argv[0], *rest]
 
@@ -89,7 +90,11 @@ def get_line_at(file: TextIO, location: concat.astutils.Location) -> str:
         raise
     else:
         concat.execute.execute(
-            filename, python_ast, {}, should_log_stacks=args.debug
+            filename,
+            python_ast,
+            {},
+            should_log_stacks=args.debug,
+            import_resolution_start_directory=os.path.dirname(filename),
         )
     finally:
         args.file.close()

concat/execute.py

@@ -1,7 +1,5 @@
 """This module takes the transpiler/compiler output and executes it."""
 import ast
-import types
-from typing import Dict, Optional, List, Callable
 import concat.stdlib.compositional
 import concat.stdlib.execution
 import concat.stdlib.importlib
@@ -13,6 +11,11 @@
 import concat.stdlib.pyinterop.math
 import concat.stdlib.pyinterop.method
 import concat.stdlib.pyinterop.user_defined_function as udf
+import contextlib
+import sys
+import types
+from typing import Callable, Dict, Iterator, List, Optional
+import os
 
 
 class ConcatRuntimeError(RuntimeError):
@@ -56,20 +59,37 @@ def _compile(filename: str, ast_: ast.Module) -> types.CodeType:
 
 def _run(
     prog: types.CodeType,
+    import_resolution_start_directory: os.PathLike,
     globals: Optional[Dict[str, object]] = None,
     locals: Optional[Dict[str, object]] = None,
 ) -> None:
     globals = {} if globals is None else globals
     try:
-        # FIXME: Imports should be resolved from the location of the source
-        # file.
-        exec(prog, globals, globals if locals is None else locals)
+        with _override_import_resolution_start_directory(
+            import_resolution_start_directory
+        ):
+            exec(prog, globals, globals if locals is None else locals)
     except Exception as e:
-        # throw away all of the traceback outside the code
-        # traceback = e.__traceback__.tb_next
+        # TODO: throw away all of the traceback outside the code, but have an
+        # option to keep the traceback.
         raise ConcatRuntimeError(globals['stack'], globals['stash']) from e
 
 
+@contextlib.contextmanager
+def _override_import_resolution_start_directory(
+    import_resolution_start_directory: os.PathLike,
+) -> Iterator[None]:
+    # This is similar to how Python sets sys.path[0] when running just a file
+    # (like `python blah.py`, not `python -m blah`). This should probably be
+    # extended when module support is implemented.
+    first_path, sys.path[0] = (
+        sys.path[0],
+        str(import_resolution_start_directory),
+    )
+    yield
+    sys.path[0] = first_path
+
+
 def _do_preamble(globals: Dict[str, object], should_log_stacks=False) -> None:
     """Add key-value pairs expected by Concat code to the passed-in mapping.
 
@@ -167,7 +187,14 @@ def execute(
     globals: Dict[str, object],
     locals: Optional[Dict[str, object]] = None,
     should_log_stacks=False,
+    # By default, sys.path[0] is '' when executing a package.
+    import_resolution_start_directory: os.PathLike = '',
 ) -> None:
     _do_preamble(globals, should_log_stacks)
 
-    _run(_compile(filename, ast), globals, locals)
+    _run(
+        _compile(filename, ast),
+        import_resolution_start_directory,
+        globals,
+        locals,
+    )

concat/stdlib/repl.py

@@ -173,6 +173,7 @@ def show_var(stack: List[object], stash: List[object]):
         while True:
             print(prompt, end='', flush=True)
             try:
+                # FIXME: `stack` might not exist yet if there was no init file.
                 eval(
                     'concat.stdlib.repl.read_form(stack, [])', globals, locals,
                 )

concat/tests/test_execute.py

@@ -1,5 +1,6 @@
 import concat.execute
 import ast
+import pathlib
 from typing import Dict
 import unittest
 
@@ -73,3 +74,18 @@ def test_preamble(self) -> None:
             with self.subTest(msg='presence of "{}"'.format(name), name=name):
                 message = 'preamble did not add "{}"'.format(name)
                 self.assertIn(name, globals, msg=message)
+
+    def test_import_resolution_location(self) -> None:
+        """Test that imports are resolved from the given directory."""
+        test_module_path = (pathlib.Path(__file__) / '../fixtures/').resolve()
+        globals_dict: dict = {}
+        concat.execute.execute(
+            '<test>',
+            ast.parse('import imported_module'),
+            globals_dict,
+            import_resolution_start_directory=test_module_path,
+        )
+        self.assertEqual(
+            pathlib.Path(globals_dict['imported_module'].__file__),
+            test_module_path / 'imported_module.py',
+        )