Source code for deephaven_enterprise.controller_import

from __future__ import annotations

from importlib.machinery import ModuleSpec
from deephaven_enterprise import dh_globals
from deephaven import jcompat

import sys


def exec_script(path: str, globals: Dict[str, Any] = None) -> None:
    """
    Attempt to exec a controller script

    :param path: full path to the script, beginning with "/"
    :param globals: the globals to use for execution of the notebook, if you want to have the variables from your
    script session available use "globals()"
    """
    src = dh_globals.j_controller_client_factory.getUnsubscribed().getScriptBody(path, True,
                                                                                 dh_globals.script_loader_state)
    if src is None:
        raise FileNotFoundError("Controller script not found: " + path)
    exec(compile(src, path, 'exec'), globals)


def meta_import(root: str = "controller") -> None:
    """
    Set up a meta importer to enable using a Python import statement that references paths from the Persistent Query Controller.

    For example, after calling 'meta_import("controller")'; you can then call 'import controller.file' to make
    the "/file.py" from the controller available as a Python module.

    :param root: the module name that prefixes all controller sourced modules, defaults to "controller"
    """
    sys.meta_path.append(__ControllerImporter(root))


class _ImportedControllerScript:
    def __init__(self, path, source=None):
        self.path = path
        self.source = source

    def create_module(self, spec):
        return None

    def exec_module(self, module):
        if self.source is not None:
            exec(compile(self.source, '<string>', 'exec'), module.__dict__)


def _j_collection_to_list(jcollection) -> List[Any]:
    """Converts a java list to a python list."""
    if not jcollection:
        return []

    res: List[Any] = []
    it = jcollection.iterator()
    while it.hasNext():
        res.append(jcompat.wrap_j_object(it.next()))

    return res


class __ControllerImporter:
    root: str

    def __init__(self, root: str):
        if root is None or root == "":
            raise Exception("root for controller imports must not be empty!")
        if "." in root:
            raise Exception(
                "root for controller imports must not contain a period!")
        self.root = root
        self.j_controller_client = dh_globals.j_controller_client_factory.getUnsubscribed()
        self.loader_state = dh_globals.script_loader_state

    def find_spec(self, name, path, target=None):
        if name != self.root and not name.startswith(self.root + "."):
            # this is not a controller import; return None so that other finder/loaders can try
            return None

        relevant_path = name[len(self.root) + 1:]
        paths = _j_collection_to_list(
            self.j_controller_client.getScriptPaths(self.loader_state, True))

        to_directory = relevant_path.replace(".", "/")

        if to_directory + ".py" in paths:
            # this is the right thing
            src_name = to_directory + ".py"
        elif to_directory + "/__init__.py" in paths:
            src_name = to_directory + "/__init__.py"
        elif relevant_path == "":
            src_name = None
        else:
            # we should check that some path here would be a child of our path
            prefix_exists: bool = False
            for x in paths:
                if x.startswith(to_directory + "/"):
                    prefix_exists = True
                    break
            if not prefix_exists:
                return None
            src_name = None

        if src_name is None:
            src = None
        else:
            src = self.j_controller_client.getScriptBody(
                src_name, True, self.loader_state)

        module_spec = ModuleSpec(
            name, _ImportedControllerScript(path, src), is_package=True)
        module_spec.submodule_search_locations = name.split(".")
        return module_spec