[JAC-MACHINE]: Isolate to other context vars

jaclang/cli/cli.py

@@ -101,19 +101,23 @@ def run(
             cachable=cache,
             override_name="__main__" if main else None,
         )
+        JacMachine.detach()
         if ret_module is None:
             loaded_mod = None
         else:
             (loaded_mod,) = ret_module
     elif filename.endswith(".jir"):
         with open(filename, "rb") as f:
+            JacMachine(base).attach_program(
+                JacProgram(mod_bundle=pickle.load(f), bytecode=None)
+            )
             ret_module = jac_import(
                 target=mod,
                 base_path=base,
                 cachable=cache,
                 override_name="__main__" if main else None,
-                jac_program=JacProgram(mod_bundle=pickle.load(f), bytecode=None),
             )
+            JacMachine.detach()
             if ret_module is None:
                 loaded_mod = None
             else:
@@ -234,6 +238,7 @@ def enter(
         base = base if base else "./"
         mod_name = mod_name[:-4]
         (mod,) = jac_import(target=mod_name, base_path=base)
+        JacMachine.detach()
         if not mod:
             print("Errors occurred while importing the module.")
         else:
@@ -386,8 +391,9 @@ def dot(
 
     if filename.endswith(".jac"):
         jac_machine = JacMachine(base)
-        jac_import(target=mod, base_path=base, jac_machine=jac_machine)
+        jac_import(target=mod, base_path=base)
         module = jac_machine.loaded_modules.get(mod)
+        JacMachine.detach()
         globals().update(vars(module))
         try:
             node = globals().get(initial, eval(initial)) if initial else None

jaclang/compiler/tests/test_importer.py

@@ -20,7 +20,7 @@ def test_import_basic_python(self) -> None:
     def test_modules_correct(self) -> None:
         """Test basic self loading."""
         jac_machine = JacMachine(__file__)
-        jac_import("fixtures.hello_world", base_path=__file__, jac_machine=jac_machine)
+        jac_import("fixtures.hello_world", base_path=__file__)
         self.assertIn(
             "module 'fixtures.hello_world'",
             str(jac_machine.loaded_modules),

jaclang/plugin/default.py

@@ -238,8 +238,6 @@ def jac_import(
         lng: Optional[str],
         items: Optional[dict[str, Union[str, Optional[str]]]],
         reload_module: Optional[bool],
-        jac_machine: Optional[JacMachine],
-        jac_program: Optional[JacProgram],
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         spec = ImportPathSpec(
@@ -252,10 +250,13 @@ def jac_import(
             lng,
             items,
         )
-        jac_machine = jac_machine or JacMachine(base_path)
-        jac_machine.attach_program(
-            jac_program or JacProgram(mod_bundle=None, bytecode=None)
-        )
+
+        if (jac_machine := JacMachine.get()) is None:
+            jac_machine = JacMachine(base_path)
+
+        if not jac_machine.jac_program:
+            jac_machine.attach_program(JacProgram(mod_bundle=None, bytecode=None))
+
         if lng == "py":
             import_result = PythonImporter(jac_machine).run_import(spec)
         else:

jaclang/plugin/feature.py

@@ -19,7 +19,6 @@
     WalkerArchitype,
 )
 from jaclang.runtimelib.context import ExecutionContext
-from jaclang.runtimelib.machine import JacMachine, JacProgram
 
 import pluggy
 
@@ -111,8 +110,6 @@ def jac_import(
         lng: Optional[str] = "jac",
         items: Optional[dict[str, Union[str, Optional[str]]]] = None,
         reload_module: Optional[bool] = False,
-        jac_machine: Optional[JacMachine] = None,
-        jac_program: Optional[JacProgram] = None,
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         return pm.hook.jac_import(
@@ -125,8 +122,6 @@ def jac_import(
             lng=lng,
             items=items,
             reload_module=reload_module,
-            jac_machine=jac_machine,
-            jac_program=jac_program,
         )
 
     @staticmethod

jaclang/plugin/spec.py

@@ -30,7 +30,6 @@
     )
     from jaclang.runtimelib.constructs import EdgeArchitype, NodeAnchor, NodeArchitype
     from jaclang.runtimelib.context import ExecutionContext
-    from jaclang.runtimelib.machine import JacMachine, JacProgram
     from jaclang.runtimelib.memory import Memory
 
 import pluggy
@@ -119,8 +118,6 @@ def jac_import(
         lng: Optional[str],
         items: Optional[dict[str, Union[str, Optional[str]]]],
         reload_module: Optional[bool],
-        jac_machine: Optional[JacMachine],
-        jac_program: Optional[JacProgram],
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         raise NotImplementedError

jaclang/runtimelib/machine.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import types
+from contextvars import ContextVar
 from typing import Optional
 
 from jaclang.compiler.absyntree import Module
@@ -14,6 +15,9 @@
 logger = logging.getLogger(__name__)
 
 
+JACMACHINE_CONTEXT = ContextVar["JacMachine | None"]("JacMachine")
+
+
 class JacMachine:
     """JacMachine to handle the VM-related functionalities and loaded programs."""
 
@@ -30,6 +34,8 @@ def __init__(self, base_path: str = "") -> None:
         )
         self.jac_program: Optional[JacProgram] = None
 
+        JACMACHINE_CONTEXT.set(self)
+
     def attach_program(self, jac_program: "JacProgram") -> None:
         """Attach a JacProgram to the machine."""
         self.jac_program = jac_program
@@ -59,6 +65,16 @@ def load_module(self, module_name: str, module: types.ModuleType) -> None:
         self.loaded_modules[module_name] = module
         sys.modules[module_name] = module
 
+    @staticmethod
+    def get() -> "JacMachine | None":
+        """Get current jac machine."""
+        return JACMACHINE_CONTEXT.get(None)
+
+    @staticmethod
+    def detach() -> None:
+        """Detach current jac machine."""
+        JACMACHINE_CONTEXT.set(None)
+
 
 class JacProgram:
     """Class to hold the mod_bundle and bytecode for Jac modules."""

jaclang/tests/test_language.py

@@ -23,6 +23,7 @@ class JacLanguageTests(TestCase):
     def setUp(self) -> None:
         """Set up test."""
         SUPER_ROOT_ANCHOR.edges.clear()
+        JacMachine.detach()
         return super().setUp()
 
     def test_sub_abilities(self) -> None:
@@ -227,11 +228,7 @@ def test_deep_imports_mods(self) -> None:
             if i in sys.modules:
                 del sys.modules[i]
         jac_machine = JacMachine(self.fixture_abs_path("./"))
-        jac_import(
-            "deep_import_mods",
-            base_path=self.fixture_abs_path("./"),
-            jac_machine=jac_machine,
-        )
+        jac_import("deep_import_mods", base_path=self.fixture_abs_path("./"))
         mods = jac_machine.loaded_modules.keys()
         for i in targets:
             self.assertIn(i, mods)