jaclang 0.7.14__py3-none-any.whl → 0.7.16__py3-none-any.whl

jaclang/plugin/feature.py

@@ -2,11 +2,15 @@
 
 from __future__ import annotations
 
+import ast as ast3
 import types
-from typing import Any, Callable, Optional, Type, TypeAlias, Union
+from typing import Any, Callable, Mapping, Optional, Sequence, Type, TypeAlias, Union
 
-from jaclang.compiler.absyntree import Module
-from jaclang.core.constructs import (
+import jaclang.compiler.absyntree as ast
+from jaclang.compiler.passes.main.pyast_gen_pass import PyastGenPass
+from jaclang.plugin.default import ExecutionContext
+from jaclang.plugin.spec import JacBuiltin, JacCmdSpec, JacFeatureSpec, P, T
+from jaclang.runtimelib.constructs import (
     Architype,
     EdgeArchitype,
     Memory,
@@ -14,9 +18,6 @@ from jaclang.core.constructs import (
     Root,
     WalkerArchitype,
 )
-from jaclang.plugin.default import ExecutionContext
-from jaclang.plugin.spec import JacBuiltin, JacCmdSpec, JacFeatureSpec, T
-
 
 import pluggy
 
@@ -29,10 +30,11 @@ pm.add_hookspecs(JacBuiltin)
 class JacFeature:
     """Jac Feature."""
 
-    import abc
-    from jaclang.compiler.constant import EdgeDir
-    from jaclang.core.constructs import DSFunc
+    from jaclang.compiler.constant import EdgeDir as EdgeDirType
+    from jaclang.runtimelib.constructs import DSFunc as DSFuncType
 
+    EdgeDir: TypeAlias = EdgeDirType
+    DSFunc: TypeAlias = DSFuncType
     RootType: TypeAlias = Root
     Obj: TypeAlias = Architype
     Node: TypeAlias = NodeArchitype
@@ -94,6 +96,13 @@ class JacFeature:
         """Create a walker architype."""
         return pm.hook.make_walker(on_entry=on_entry, on_exit=on_exit)
 
+    @staticmethod
+    def impl_patch_filename(
+        file_loc: str,
+    ) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        """Update impl file location."""
+        return pm.hook.impl_patch_filename(file_loc=file_loc)
+
     @staticmethod
     def jac_import(
         target: str,
@@ -102,9 +111,9 @@ class JacFeature:
         cachable: bool = True,
         mdl_alias: Optional[str] = None,
         override_name: Optional[str] = None,
-        mod_bundle: Optional[Module | str] = None,
         lng: Optional[str] = "jac",
         items: Optional[dict[str, Union[str, Optional[str]]]] = None,
+        reload_module: Optional[bool] = False,
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         return pm.hook.jac_import(
@@ -114,9 +123,9 @@ class JacFeature:
             cachable=cachable,
             mdl_alias=mdl_alias,
             override_name=override_name,
-            mod_bundle=mod_bundle,
             lng=lng,
             items=items,
+            reload_module=reload_module,
         )
 
     @staticmethod
@@ -301,6 +310,8 @@ class JacFeature:
         inputs: list[tuple[str, str, str, Any]],
         outputs: tuple,
         action: str,
+        _globals: dict,
+        _locals: Mapping,
     ) -> Any:  # noqa: ANN401
         """Jac's with_llm feature."""
         return pm.hook.with_llm(
@@ -313,8 +324,45 @@ class JacFeature:
             inputs=inputs,
             outputs=outputs,
             action=action,
+            _globals=_globals,
+            _locals=_locals,
         )
 
+    @staticmethod
+    def gen_llm_body(_pass: PyastGenPass, node: ast.Ability) -> list[ast3.AST]:
+        """Generate the by LLM body."""
+        return pm.hook.gen_llm_body(_pass=_pass, node=node)
+
+    @staticmethod
+    def by_llm_call(
+        _pass: PyastGenPass,
+        model: ast3.AST,
+        model_params: dict[str, ast.Expr],
+        scope: ast3.AST,
+        inputs: Sequence[Optional[ast3.AST]],
+        outputs: Sequence[Optional[ast3.AST]] | ast3.Call,
+        action: Optional[ast3.AST],
+        include_info: list[tuple[str, ast3.AST]],
+        exclude_info: list[tuple[str, ast3.AST]],
+    ) -> ast3.Call:
+        """Return the LLM Call, e.g. _Jac.with_llm()."""
+        return pm.hook.by_llm_call(
+            _pass=_pass,
+            model=model,
+            model_params=model_params,
+            scope=scope,
+            inputs=inputs,
+            outputs=outputs,
+            action=action,
+            include_info=include_info,
+            exclude_info=exclude_info,
+        )
+
+    @staticmethod
+    def get_by_llm_call_args(_pass: PyastGenPass, node: ast.FuncCall) -> dict:
+        """Get the by LLM call args."""
+        return pm.hook.get_by_llm_call_args(_pass=_pass, node=node)
+
 
 class JacCmd:
     """Jac CLI command."""

jaclang/plugin/spec.py

@@ -2,13 +2,26 @@
 
 from __future__ import annotations
 
+import ast as ast3
 import types
-from typing import Any, Callable, Optional, TYPE_CHECKING, Type, TypeVar, Union
-
-from jaclang.compiler.absyntree import Module
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    Optional,
+    ParamSpec,
+    Sequence,
+    TYPE_CHECKING,
+    Type,
+    TypeVar,
+    Union,
+)
+
+import jaclang.compiler.absyntree as ast
+from jaclang.compiler.passes.main.pyast_gen_pass import PyastGenPass
 
 if TYPE_CHECKING:
-    from jaclang.core.constructs import EdgeArchitype, NodeArchitype
+    from jaclang.runtimelib.constructs import EdgeArchitype, NodeArchitype
     from jaclang.plugin.default import (
         Architype,
         EdgeDir,
@@ -17,13 +30,14 @@ if TYPE_CHECKING:
         Root,
         DSFunc,
     )
-    from jaclang.core.memory import Memory
+    from jaclang.runtimelib.memory import Memory
 
 import pluggy
 
 hookspec = pluggy.HookspecMarker("jac")
 
 T = TypeVar("T")
+P = ParamSpec("P")
 
 
 class JacFeatureSpec:
@@ -90,6 +104,14 @@ class JacFeatureSpec:
         """Create a walker architype."""
         raise NotImplementedError
 
+    @staticmethod
+    @hookspec(firstresult=True)
+    def impl_patch_filename(
+        file_loc: str,
+    ) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        """Update impl file location."""
+        raise NotImplementedError
+
     @staticmethod
     @hookspec(firstresult=True)
     def jac_import(
@@ -99,9 +121,9 @@ class JacFeatureSpec:
         cachable: bool,
         mdl_alias: Optional[str],
         override_name: Optional[str],
-        mod_bundle: Optional[Module | str],
         lng: Optional[str],
         items: Optional[dict[str, Union[str, Optional[str]]]],
+        reload_module: Optional[bool],
     ) -> tuple[types.ModuleType, ...]:
         """Core Import Process."""
         raise NotImplementedError
@@ -283,10 +305,40 @@ class JacFeatureSpec:
         inputs: list[tuple[str, str, str, Any]],
         outputs: tuple,
         action: str,
+        _globals: dict,
+        _locals: Mapping,
     ) -> Any:  # noqa: ANN401
         """Jac's with_llm stmt feature."""
         raise NotImplementedError
 
+    @staticmethod
+    @hookspec(firstresult=True)
+    def gen_llm_body(_pass: PyastGenPass, node: ast.Ability) -> list[ast3.AST]:
+        """Generate the by LLM body."""
+        raise NotImplementedError
+
+    @staticmethod
+    @hookspec(firstresult=True)
+    def by_llm_call(
+        _pass: PyastGenPass,
+        model: ast3.AST,
+        model_params: dict[str, ast.Expr],
+        scope: ast3.AST,
+        inputs: Sequence[Optional[ast3.AST]],
+        outputs: Sequence[Optional[ast3.AST]] | ast3.Call,
+        action: Optional[ast3.AST],
+        include_info: list[tuple[str, ast3.AST]],
+        exclude_info: list[tuple[str, ast3.AST]],
+    ) -> ast3.Call:
+        """Return the LLM Call, e.g. _Jac.with_llm()."""
+        raise NotImplementedError
+
+    @staticmethod
+    @hookspec(firstresult=True)
+    def get_by_llm_call_args(_pass: PyastGenPass, node: ast.FuncCall) -> dict:
+        """Get the by LLM call args."""
+        raise NotImplementedError
+
 
 class JacBuiltin:
     """Jac Builtins."""

jaclang/{core → runtimelib}/architype.py

@@ -8,7 +8,7 @@ from typing import Any, Callable, Optional
 from uuid import UUID, uuid4
 
 from jaclang.compiler.constant import EdgeDir
-from jaclang.core.utils import collect_node_connections
+from jaclang.runtimelib.utils import collect_node_connections
 
 
 @dataclass(eq=False)

jaclang/{core → runtimelib}/context.py

@@ -8,6 +8,7 @@ from typing import Callable, Optional
 from uuid import UUID
 
 from .architype import Architype, Root
+from .machine import JacMachine, JacProgram
 from .memory import Memory, ShelveStorage
 
 
@@ -22,13 +23,19 @@ class ExecutionContext:
         super().__init__()
         self.mem = ShelveStorage()
         self.root = None
+        self.jac_machine = JacMachine()
+        jac_program = JacProgram(mod_bundle=None, bytecode=None)
+        self.jac_machine.attach_program(jac_program)
 
-    def init_memory(self, session: str = "") -> None:
+    def init_memory(self, base_path: str = "", session: str = "") -> None:
         """Initialize memory."""
         if session:
             self.mem = ShelveStorage(session)
         else:
             self.mem = Memory()
+        self.jac_machine = JacMachine(base_path)
+        jac_program = JacProgram(mod_bundle=None, bytecode=None)
+        self.jac_machine.attach_program(jac_program)
 
     def get_root(self) -> Root:
         """Get the root object."""

jaclang/runtimelib/importer.py

@@ -0,0 +1,361 @@
+"""Special Imports for Jac Code."""
+
+from __future__ import annotations
+
+import importlib
+import os
+import sys
+import types
+from os import getcwd, path
+from typing import Optional, Union
+
+from jaclang.runtimelib.machine import JacMachine
+from jaclang.runtimelib.utils import sys_path_context
+from jaclang.utils.helpers import dump_traceback
+from jaclang.utils.log import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ImportPathSpec:
+    """Import Specification."""
+
+    def __init__(
+        self,
+        target: str,
+        base_path: str,
+        absorb: bool,
+        cachable: bool,
+        mdl_alias: Optional[str],
+        override_name: Optional[str],
+        lng: Optional[str],
+        items: Optional[dict[str, Union[str, Optional[str]]]],
+    ) -> None:
+        """Initialize the ImportPathSpec object."""
+        self.target = target
+        self.base_path = base_path
+        self.absorb = absorb
+        self.cachable = cachable
+        self.mdl_alias = mdl_alias
+        self.override_name = override_name
+        self.language = lng
+        self.items = items
+        self.dir_path, self.file_name = path.split(path.join(*(target.split("."))))
+        self.module_name = path.splitext(self.file_name)[0]
+        self.package_path = self.dir_path.replace(path.sep, ".")
+        self.caller_dir = self.get_caller_dir()
+        self.full_target = path.abspath(path.join(self.caller_dir, self.file_name))
+
+    def get_caller_dir(self) -> str:
+        """Get the directory of the caller."""
+        caller_dir = (
+            self.base_path
+            if path.isdir(self.base_path)
+            else path.dirname(self.base_path)
+        )
+        caller_dir = caller_dir if caller_dir else getcwd()
+        chomp_target = self.target
+        if chomp_target.startswith("."):
+            chomp_target = chomp_target[1:]
+            while chomp_target.startswith("."):
+                caller_dir = path.dirname(caller_dir)
+                chomp_target = chomp_target[1:]
+        return path.join(caller_dir, self.dir_path)
+
+
+class ImportReturn:
+    """Import Return Object."""
+
+    def __init__(
+        self,
+        ret_mod: types.ModuleType,
+        ret_items: list[types.ModuleType],
+        importer: Importer,
+    ) -> None:
+        """Initialize the ImportReturn object."""
+        self.ret_mod = ret_mod
+        self.ret_items = ret_items
+        self.importer = importer
+
+    def process_items(
+        self,
+        module: types.ModuleType,
+        items: dict[str, Union[str, Optional[str]]],
+        caller_dir: str,
+        lang: Optional[str],
+        cachable: bool = True,
+    ) -> None:
+        """Process items within a module by handling renaming and potentially loading missing attributes."""
+
+        def handle_item_loading(
+            item: types.ModuleType, alias: Union[str, Optional[str]]
+        ) -> None:
+            if item:
+                self.ret_items.append(item)
+                setattr(module, name, item)
+                if alias and alias != name and not isinstance(alias, bool):
+                    setattr(module, alias, item)
+
+        for name, alias in items.items():
+            try:
+                item = getattr(module, name)
+                handle_item_loading(item, alias)
+            except AttributeError:
+                if lang == "jac":
+                    jac_file_path = (
+                        os.path.join(module.__path__[0], f"{name}.jac")
+                        if hasattr(module, "__path__")
+                        else module.__file__
+                    )
+
+                    if jac_file_path and os.path.isfile(jac_file_path):
+                        item = self.load_jac_mod_as_item(
+                            module=module,
+                            name=name,
+                            jac_file_path=jac_file_path,
+                            cachable=cachable,
+                            caller_dir=caller_dir,
+                        )
+                        handle_item_loading(item, alias)
+                else:
+                    if hasattr(module, "__path__"):
+                        full_module_name = f"{module.__name__}.{name}"
+                        item = importlib.import_module(full_module_name)
+                        handle_item_loading(item, alias)
+
+    def load_jac_mod_as_item(
+        self,
+        module: types.ModuleType,
+        name: str,
+        jac_file_path: str,
+        cachable: bool,
+        caller_dir: str,
+    ) -> Optional[types.ModuleType]:
+        """Load a single .jac file into the specified module component."""
+        try:
+            package_name = (
+                f"{module.__name__}.{name}"
+                if hasattr(module, "__path__")
+                else module.__name__
+            )
+            if isinstance(self.importer, JacImporter):
+                new_module = sys.modules.get(
+                    package_name,
+                    self.importer.create_jac_py_module(
+                        self.importer.get_sys_mod_name(jac_file_path),
+                        module.__name__,
+                        jac_file_path,
+                    ),
+                )
+            codeobj = self.importer.jac_machine.get_bytecode(
+                name, jac_file_path, caller_dir=caller_dir, cachable=cachable
+            )
+            if not codeobj:
+                raise ImportError(f"No bytecode found for {jac_file_path}")
+
+            exec(codeobj, new_module.__dict__)
+            return getattr(new_module, name, new_module)
+        except ImportError as e:
+            logger.error(dump_traceback(e))
+            # logger.error(
+            #     f"Failed to load {name} from {jac_file_path} in {module.__name__}: {str(e)}"
+            # )
+            return None
+
+
+class Importer:
+    """Abstract base class for all importers."""
+
+    def __init__(self, jac_machine: JacMachine) -> None:
+        """Initialize the Importer object."""
+        self.jac_machine = jac_machine
+        self.result: Optional[ImportReturn] = None
+
+    def run_import(self, spec: ImportPathSpec) -> ImportReturn:
+        """Run the import process."""
+        raise NotImplementedError
+
+    def update_sys(self, module: types.ModuleType, spec: ImportPathSpec) -> None:
+        """Update sys.modules with the newly imported module."""
+        if spec.module_name not in sys.modules:
+            sys.modules[spec.module_name] = module
+
+
+class PythonImporter(Importer):
+    """Importer for Python modules."""
+
+    def __init__(self, jac_machine: JacMachine) -> None:
+        """Initialize the PythonImporter object."""
+        self.jac_machine = jac_machine
+
+    def run_import(self, spec: ImportPathSpec) -> ImportReturn:
+        """Run the import process for Python modules."""
+        try:
+            loaded_items: list = []
+            if spec.target.startswith("."):
+                spec.target = spec.target.lstrip(".")
+                full_target = path.normpath(path.join(spec.caller_dir, spec.target))
+                imp_spec = importlib.util.spec_from_file_location(
+                    spec.target, full_target + ".py"
+                )
+                if imp_spec and imp_spec.loader:
+                    imported_module = importlib.util.module_from_spec(imp_spec)
+                    sys.modules[imp_spec.name] = imported_module
+                    imp_spec.loader.exec_module(imported_module)
+                else:
+                    raise ImportError(
+                        f"Cannot find module {spec.target} at {full_target}"
+                    )
+            else:
+                imported_module = importlib.import_module(name=spec.target)
+
+            main_module = __import__("__main__")
+            if spec.absorb:
+                for name in dir(imported_module):
+                    if not name.startswith("_"):
+                        setattr(main_module, name, getattr(imported_module, name))
+
+            elif spec.items:
+                for name, alias in spec.items.items():
+                    if isinstance(alias, bool):
+                        alias = name
+                    try:
+                        item = getattr(imported_module, name)
+                        if item not in loaded_items:
+                            setattr(
+                                main_module,
+                                alias if isinstance(alias, str) else name,
+                                item,
+                            )
+                            loaded_items.append(item)
+                    except AttributeError as e:
+                        if hasattr(imported_module, "__path__"):
+                            item = importlib.import_module(f"{spec.target}.{name}")
+                            if item not in loaded_items:
+                                setattr(
+                                    main_module,
+                                    alias if isinstance(alias, str) else name,
+                                    item,
+                                )
+                                loaded_items.append(item)
+                        else:
+                            raise e
+
+            else:
+                setattr(
+                    __import__("__main__"),
+                    spec.mdl_alias if isinstance(spec.mdl_alias, str) else spec.target,
+                    imported_module,
+                )
+            self.result = ImportReturn(imported_module, loaded_items, self)
+            return self.result
+
+        except ImportError as e:
+            raise e
+
+
+class JacImporter(Importer):
+    """Importer for Jac modules."""
+
+    def __init__(self, jac_machine: JacMachine) -> None:
+        """Initialize the JacImporter object."""
+        self.jac_machine = jac_machine
+
+    def get_sys_mod_name(self, full_target: str) -> str:
+        """Generate the system module name based on full target path and package path."""
+        if full_target == self.jac_machine.base_path_dir:
+            return path.basename(self.jac_machine.base_path_dir)
+        relative_path = path.relpath(full_target, start=self.jac_machine.base_path_dir)
+        base_name = path.splitext(relative_path)[0]
+        sys_mod_name = base_name.replace(os.sep, ".").strip(".")
+        return sys_mod_name
+
+    def handle_directory(
+        self, module_name: str, full_mod_path: str
+    ) -> types.ModuleType:
+        """Import from a directory that potentially contains multiple Jac modules."""
+        module_name = self.get_sys_mod_name(full_mod_path)
+        module = types.ModuleType(module_name)
+        module.__name__ = module_name
+        module.__path__ = [full_mod_path]
+        module.__file__ = None
+
+        if module_name not in sys.modules:
+            sys.modules[module_name] = module
+        return module
+
+    def create_jac_py_module(
+        self,
+        module_name: str,
+        package_path: str,
+        full_target: str,
+    ) -> types.ModuleType:
+        """Create a module."""
+        module = types.ModuleType(module_name)
+        module.__file__ = full_target
+        module.__name__ = module_name
+        if package_path:
+            base_path = full_target.split(package_path.replace(".", path.sep))[0]
+            parts = package_path.split(".")
+            for i in range(len(parts)):
+                package_name = ".".join(parts[: i + 1])
+                if package_name not in sys.modules:
+                    full_mod_path = path.join(
+                        base_path, package_name.replace(".", path.sep)
+                    )
+                    self.handle_directory(
+                        module_name=package_name,
+                        full_mod_path=full_mod_path,
+                    )
+        sys.modules[module_name] = module
+        return module
+
+    def run_import(
+        self, spec: ImportPathSpec, reload: Optional[bool] = False
+    ) -> ImportReturn:
+        """Run the import process for Jac modules."""
+        unique_loaded_items: list[types.ModuleType] = []
+        module = None
+        if os.path.isfile(spec.full_target + ".jac"):
+            module_name = self.get_sys_mod_name(spec.full_target + ".jac")
+            module_name = spec.override_name if spec.override_name else module_name
+        else:
+            module_name = self.get_sys_mod_name(spec.full_target)
+
+        module = sys.modules.get(module_name)
+
+        if not module or module.__name__ == "__main__" or reload:
+            if os.path.isdir(spec.full_target):
+                module = self.handle_directory(spec.module_name, spec.full_target)
+            else:
+                spec.full_target += ".jac" if spec.language == "jac" else ".py"
+                module = self.create_jac_py_module(
+                    module_name,
+                    spec.package_path,
+                    spec.full_target,
+                )
+                codeobj = self.jac_machine.get_bytecode(
+                    module_name,
+                    spec.full_target,
+                    caller_dir=spec.caller_dir,
+                    cachable=spec.cachable,
+                )
+                if not codeobj:
+                    raise ImportError(f"No bytecode found for {spec.full_target}")
+                with sys_path_context(spec.caller_dir):
+                    try:
+                        exec(codeobj, module.__dict__)
+                    except Exception as e:
+                        logger.error(dump_traceback(e))
+                        raise e
+        import_return = ImportReturn(module, unique_loaded_items, self)
+        if spec.items:
+            import_return.process_items(
+                module=module,
+                items=spec.items,
+                caller_dir=spec.caller_dir,
+                cachable=spec.cachable,
+                lang=spec.language,
+            )
+        self.result = import_return
+        return self.result