Source code for deephaven_enterprise.notebook

#
# Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
#
"""This module provides a way for users to modularize their Deephaven Web console notebooks and import and execute them
as Python modules.
"""

from __future__ import annotations

import sys
from importlib.machinery import ModuleSpec
from typing import Any, Optional

import jpy

from deephaven_enterprise.database import Database




def exec_notebook(
    db: Database, path: str, globals: Optional[dict[str, Any]] = None
) -> None:
    """Executes a Deephaven Web console notebook.

    Args:
        db (Database): Core+ Python Database object
        path (str): full path to Notebook, beginning with "/"
        globals (dict[str, Any]): 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_database, path)
    exec(compile(src, path, "exec"), globals)





def meta_import(db: Database, root: str = "notebook") -> None:
    """Sets 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.

    Note: implicitly populated global objects like 'db', 'performance_overview', etc., in the main module are not
    automatically available in the imported module. You must explicitly import them or pass them as arguments to
    functions in the module. To import them, you can use the following:

        from deephaven_enterprise.database import db

        from deephaven_enterprise.performance_overview import performance_overview

    Args:
        db (Database): Core+ Python Database object
        root (str): 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: 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_database, notebook_path, False, ".py"
            )
            if data.isFile():
                src = data.getContents()
            elif data.isFolder():
                try:
                    init_data = self.notebook_reader.getFileOrFolder(
                        self.db.j_database,
                        notebook_path.replace(".", "/") + "/__init__.py",
                        True,
                        "",
                    )
                    src = init_data.getContents()
                except Exception:  # noqa: E722
                    pass

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