[JAC-MACHINE]: Isolate to other context vars

jaclang/cli/cli.py

@@ -100,19 +100,24 @@ def run(
             base_path=base,
             cachable=cache,
             override_name="__main__" if main else None,
+            is_origin=True,
         )
+
         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),
+                is_origin=True,
             )
             if ret_module is None:
                 loaded_mod = None
@@ -233,7 +238,7 @@ def enter(
         base, mod_name = os.path.split(filename)
         base = base if base else "./"
         mod_name = mod_name[:-4]
-        (mod,) = jac_import(target=mod_name, base_path=base)
+        (mod,) = jac_import(target=mod_name, base_path=base, is_origin=True)
         if not mod:
             print("Errors occurred while importing the module.")
         else:
@@ -386,7 +391,7 @@ 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, is_origin=True)
         module = jac_machine.loaded_modules.get(mod)
         globals().update(vars(module))
         try:

jaclang/compiler/tests/test_importer.py

@@ -14,13 +14,13 @@ class TestLoader(TestCase):
 
     def test_import_basic_python(self) -> None:
         """Test basic self loading."""
-        (h,) = jac_import("fixtures.hello_world", base_path=__file__)
+        (h,) = jac_import("fixtures.hello_world", base_path=__file__, is_origin=True)
         self.assertEqual(h.hello(), "Hello World!")  # type: ignore
 
     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__, is_origin=True)
         self.assertIn(
             "module 'fixtures.hello_world'",
             str(jac_machine.loaded_modules),

jaclang/plugin/default.py

@@ -79,6 +79,18 @@ def get_datasource() -> Memory:
         """Get current execution context."""
         return ExecutionContext.get().datasource
 
+    @staticmethod
+    @hookimpl
+    def get_machine(base_path: str) -> JacMachine:
+        """Get current execution context."""
+        return JacMachine.get(base_path)
+
+    @staticmethod
+    @hookimpl
+    def detach_machine() -> None:
+        """Detach current jac machine."""
+        JacMachine.detach()
+
     @staticmethod
     @hookimpl
     def make_architype(
@@ -238,8 +250,7 @@ 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],
+        is_origin: bool,
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         spec = ImportPathSpec(
@@ -252,14 +263,19 @@ 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)
-        )
+
+        jac_machine = Jac.get_machine(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:
             import_result = JacImporter(jac_machine).run_import(spec, reload_module)
+
+        if is_origin:
+            Jac.detach_machine()
+
         return (
             (import_result.ret_mod,)
             if absorb or not items

jaclang/plugin/feature.py

@@ -19,7 +19,7 @@
     WalkerArchitype,
 )
 from jaclang.runtimelib.context import ExecutionContext
-from jaclang.runtimelib.machine import JacMachine, JacProgram
+from jaclang.runtimelib.machine import JacMachine
 
 import pluggy
 
@@ -53,6 +53,16 @@ def get_datasource() -> Memory:
         """Get current execution context."""
         return pm.hook.get_datasource()
 
+    @staticmethod
+    def get_machine(base_path: str = "") -> JacMachine:
+        """Get current jac machine."""
+        return pm.hook.get_machine(base_path=base_path)
+
+    @staticmethod
+    def detach_machine() -> None:
+        """Detach current jac machine."""
+        return pm.hook.detach_machine()
+
     @staticmethod
     def make_architype(
         cls: type,
@@ -111,8 +121,7 @@ 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,
+        is_origin: bool = False,
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         return pm.hook.jac_import(
@@ -125,8 +134,7 @@ def jac_import(
             lng=lng,
             items=items,
             reload_module=reload_module,
-            jac_machine=jac_machine,
-            jac_program=jac_program,
+            is_origin=is_origin,
         )
 
     @staticmethod

jaclang/plugin/spec.py

@@ -30,7 +30,7 @@
     )
     from jaclang.runtimelib.constructs import EdgeArchitype, NodeAnchor, NodeArchitype
     from jaclang.runtimelib.context import ExecutionContext
-    from jaclang.runtimelib.machine import JacMachine, JacProgram
+    from jaclang.runtimelib.machine import JacMachine
     from jaclang.runtimelib.memory import Memory
 
 import pluggy
@@ -56,6 +56,18 @@ def get_datasource() -> Memory:
         """Get current execution context datasource."""
         raise NotImplementedError
 
+    @staticmethod
+    @hookspec(firstresult=True)
+    def get_machine(base_path: str) -> JacMachine:
+        """Get current jac machine."""
+        raise NotImplementedError
+
+    @staticmethod
+    @hookspec(firstresult=True)
+    def detach_machine() -> None:
+        """Detach current jac machine."""
+        raise NotImplementedError
+
     @staticmethod
     @hookspec(firstresult=True)
     def make_architype(
@@ -119,8 +131,7 @@ 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],
+        is_origin: bool,
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         raise NotImplementedError

jaclang/runtimelib/machine.py

@@ -5,6 +5,7 @@
 import os
 import sys
 import types
+from contextvars import ContextVar
 from typing import Optional
 
 from jaclang.compiler.absyntree import Module
@@ -17,6 +18,9 @@
 logger = logging.getLogger(__name__)
 
 
+JACMACHINE_CONTEXT = ContextVar["JacMachine | None"]("JacMachine")
+
+
 class JacMachine:
     """JacMachine to handle the VM-related functionalities and loaded programs."""
 
@@ -33,6 +37,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
@@ -99,6 +105,18 @@ def list_edges(self, module_name: str) -> list[str]:
             return nodes
         return []
 
+    @staticmethod
+    def get(base_path: str = "") -> "JacMachine":
+        """Get current jac machine."""
+        if (jac_machine := JACMACHINE_CONTEXT.get(None)) is None:
+            jac_machine = JacMachine(base_path)
+        return jac_machine
+
+    @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/fixtures/foo.jac

@@ -5,13 +5,13 @@ import:jac from bar, bar_walk;
 # Test runner to initialize the walker
 can test_run {
     # Print the loaded modules
-    modules = Jac.context().jac_machine.list_modules();
+    modules = Jac.get_machine().list_modules();
     "Loaded Modules:" |> print;
     for mod_name in modules {
         f"Module: {mod_name}" |> print;
     }
     # Print walkers
-    walkers = Jac.context().jac_machine.list_walkers(mod_name);
+    walkers = Jac.get_machine().list_walkers(mod_name);
         if walkers{
             f"Walkers in {mod_name}:" |> print;
             for walker in walkers{
@@ -20,15 +20,15 @@ can test_run {
         }
 
     # Print nodes
-    nodes = Jac.context().jac_machine.list_nodes(mod_name);
+    nodes = Jac.get_machine().list_nodes(mod_name);
     if nodes{
         f"Nodes in {mod_name}:" |> print;
         for node in nodes{
             f"  - Node: {node}" |> print;
         }
     }
     # Print edges
-    edges = Jac.context().jac_machine.list_edges(mod_name);
+    edges = Jac.get_machine().list_edges(mod_name);
     if edges{
         f"Edges in {mod_name}:" |> print;
         for edge in edges{