Source code for deephaven_enterprise.notebook

from __future__ import annotations

from importlib.machinery import ModuleSpec
import sys
import jpy

import deephaven_enterprise.database


def exec_notebook(db: deephaven_enterprise.database.Database, path: str, globals: dict[str, Any] = None) -> None:
    """
    Attempt to exec a Notebook

    :param db: DnD Python Database object
    :param path: full path to Notebook, 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()"
    """
    nr = jpy.get_type("io.deephaven.enterprise.dnd.notebook.NotebookReader")
    src = nr.getNotebookText(db.j_db, path)
    exec(compile(src, path, 'exec'), globals)


def meta_import(db: deephaven_enterprise.database.Database, root: str = "notebook") -> None:
    """
    Set up a meta importer to enable using a Python import statement that references Deephaven Web console notebooks.

    For example, after calling 'meta_import(db, "notebook")'; you can then call 'import notebook.file' to make
    the "/file.py" in your Web File Explorer available as a Python module.

    :param db: DnD Python Database object
    :param root: the module name that prefixes all web notebooks, defaults to "notebook"
    """
    sys.meta_path.append(__NotebookImporter(db, root))


class _ImportedNotebook:
    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__)


class __NotebookImporter:

    db: deephaven_enterprise.database.Database
    root: str

    def __init__(self, db, root: str):
        self.db = db
        if root is None or root == "":
            raise Exception("root for notebook imports must not be empty!")
        if "." in root:
            raise Exception(
                "root for notebook imports must not contain a period!")
        self.root = root
        self.notebook_reader = jpy.get_type(
            "io.deephaven.enterprise.dnd.notebook.NotebookReader")

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

        src = None
        if name.startswith(self.root + "."):
            # we want to keep one dot present
            notebook_path = name[len(self.root):]
            data = self.notebook_reader.getFileOrFolder(
                self.db.j_db, notebook_path, False, ".py")
            if data.isFile():
                src = data.getContents()
            elif data.isFolder():
                try:
                    init_data = self.notebook_reader.getFileOrFolder(self.db.j_db, notebook_path.replace(".", "/") + "/__init__.py", True, "")
                    src = init_data.getContents()
                except:
                    pass

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