ygg 0.1.30__py3-none-any.whl → 0.1.32__py3-none-any.whl

yggdrasil/databricks/compute/execution_context.py

@@ -1,3 +1,5 @@
+"""Remote execution helpers for Databricks command contexts."""
+
 import base64
 import dataclasses as dc
 import datetime as dt
@@ -33,6 +35,7 @@ logger = logging.getLogger(__name__)
 
 @dc.dataclass
 class RemoteMetadata:
+    """Metadata describing the remote cluster execution environment."""
     site_packages_path: Optional[str] = dc.field(default=None)
     os_env: Dict[str, str] = dc.field(default_factory=dict)
     requirements: Optional[str] = dc.field(default=None)
@@ -42,6 +45,7 @@ class RemoteMetadata:
         self,
         current: Optional[Dict] = None
     ):
+        """Return environment variables present locally but missing remotely."""
         if current is None:
             current = os.environ
 
@@ -81,6 +85,7 @@ class ExecutionContext:
 
     # --- Pickle / cloudpickle support (don’t serialize locks or cached remote metadata) ---
     def __getstate__(self):
+        """Serialize context state, excluding locks and remote metadata."""
         state = self.__dict__.copy()
 
         # name-mangled field for _lock in instance dict:
@@ -89,25 +94,30 @@ class ExecutionContext:
         return state
 
     def __setstate__(self, state):
+        """Restore context state, rehydrating locks if needed."""
         state["_lock"] = state.get("_lock", threading.RLock())
 
         self.__dict__.update(state)
 
     def __enter__(self) -> "ExecutionContext":
+        """Enter a context manager, opening a remote execution context."""
         self.cluster.__enter__()
         self._was_connected = self.context_id is not None
         return self.connect()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the context manager and close the remote context if created."""
         if not self._was_connected:
             self.close()
         self.cluster.__exit__(exc_type, exc_val=exc_val, exc_tb=exc_tb)
 
     def __del__(self):
+        """Best-effort cleanup for the remote execution context."""
         self.close()
 
     @property
     def remote_metadata(self) -> RemoteMetadata:
+        """Fetch and cache remote environment metadata for the cluster."""
         # fast path (no lock)
         rm = self._remote_metadata
         if rm is not None:
@@ -134,7 +144,7 @@ os_env = meta["os_env"] = {}
 for k, v in os.environ.items():
     os_env[k] = v
     
-meta["requirements"] = current_env.export_requirements_matrix()
+meta["requirements"] = current_env.requirements()
 meta["version_info"] = current_env.version_info
 
 print(json.dumps(meta))"""
@@ -151,17 +161,24 @@ print(json.dumps(meta))"""
 
     # ------------ internal helpers ------------
     def _workspace_client(self):
+        """Return the Databricks SDK client for command execution.
+
+        Returns:
+            The underlying WorkspaceClient instance.
+        """
         return self.cluster.workspace.sdk()
 
     def _create_command(
         self,
         language: "Language",
     ) -> any:
-        """
-        Wrap `client.command_execution.create` in a 10s timeout.
-        On timeout:
-          - ensure the cluster is running
-          - retry once with the same timeout
+        """Create a command execution context, retrying if needed.
+
+        Args:
+            language: The Databricks command language to use.
+
+        Returns:
+            The created command execution context response.
         """
         self.cluster.ensure_running()
 
@@ -182,7 +199,14 @@ print(json.dumps(meta))"""
         self,
         language: Optional["Language"] = None
     ) -> "ExecutionContext":
-        """Create a remote command execution context if not already open."""
+        """Create a remote command execution context if not already open.
+
+        Args:
+            language: Optional language override for the context.
+
+        Returns:
+            The connected ExecutionContext instance.
+        """
         if self.context_id is not None:
             logger.debug(
                 "Execution context already open for %s",
@@ -209,7 +233,11 @@ print(json.dumps(meta))"""
         return self
 
     def close(self) -> None:
-        """Destroy the remote command execution context if it exists."""
+        """Destroy the remote command execution context if it exists.
+
+        Returns:
+            None.
+        """
         if self.context_id is None:
             return
 
@@ -241,6 +269,21 @@ print(json.dumps(meta))"""
         result_tag: Optional[str] = None,
         **options
     ):
+        """Execute a string command or a callable in the remote context.
+
+        Args:
+            obj: Command string or callable to execute.
+            args: Optional positional arguments for callables.
+            kwargs: Optional keyword arguments for callables.
+            env_keys: Environment variable names to forward.
+            env_variables: Environment variables to inject remotely.
+            timeout: Optional timeout for execution.
+            result_tag: Optional result tag for parsing output.
+            **options: Additional execution options.
+
+        Returns:
+            The decoded execution result.
+        """
         if isinstance(obj, str):
             return self.execute_command(
                 command=obj,
@@ -261,6 +304,7 @@ print(json.dumps(meta))"""
         raise ValueError(f"Cannot execute {type(obj)}")
 
     def is_in_databricks_environment(self):
+        """Return True when running on a Databricks runtime."""
         return self.cluster.is_in_databricks_environment()
 
     def execute_callable(
@@ -274,12 +318,26 @@ print(json.dumps(meta))"""
         timeout: Optional[dt.timedelta] = None,
         command: Optional[str] = None,
     ) -> Any:
+        """Execute a Python callable remotely and return the decoded result.
+
+        Args:
+            func: Callable or serialized callable to run remotely.
+            args: Positional arguments for the callable.
+            kwargs: Keyword arguments for the callable.
+            env_keys: Environment variable names to forward.
+            env_variables: Environment variables to inject remotely.
+            print_stdout: Whether to print stdout from the command output.
+            timeout: Optional timeout for execution.
+            command: Optional prebuilt command string override.
+
+        Returns:
+            The decoded return value from the remote execution.
+        """
         if self.is_in_databricks_environment():
             args = args or []
             kwargs = kwargs or {}
             return func(*args, **kwargs)
 
-        """Execute a command in this context and return decoded output."""
         self.connect(language=Language.PYTHON)
 
         logger.debug(
@@ -354,7 +412,17 @@ print(json.dumps(meta))"""
         result_tag: Optional[str] = None,
         print_stdout: Optional[bool] = True,
     ) -> str:
-        """Execute a command in this context and return decoded output."""
+        """Execute a command in this context and return decoded output.
+
+        Args:
+            command: The command string to execute.
+            timeout: Optional timeout for execution.
+            result_tag: Optional tag to extract a specific result segment.
+            print_stdout: Whether to print stdout for tagged output.
+
+        Returns:
+            The decoded command output string.
+        """
         self.connect()
 
         client = self._workspace_client()
@@ -402,6 +470,12 @@ print(json.dumps(meta))"""
         - If local_path is a directory:
               remote_path is the *directory root* on remote; the directory
               contents are mirrored under it.
+        Args:
+            local_path: Local file or directory to upload.
+            remote_path: Target path on the remote cluster.
+
+        Returns:
+            None.
         """
         local_path = os.path.abspath(local_path)
         if not os.path.exists(local_path):
@@ -490,6 +564,12 @@ with zipfile.ZipFile(buf, "r") as zf:
         - path to a file    (e.g. "./ygg/__init__.py")
         - module name       (e.g. "ygg")
         - module object     (e.g. import ygg; workspace.upload_local_lib(ygg))
+        Args:
+            libraries: Library path, name, module, or iterable of these.
+            with_dependencies: Whether to include dependencies (unused).
+
+        Returns:
+            The resolved library or list of libraries uploaded.
         """
         if isinstance(libraries, (list, tuple, set)):
             return [
@@ -517,7 +597,16 @@ with zipfile.ZipFile(buf, "r") as zf:
         result_tag: Optional[str],
         print_stdout: Optional[bool] = True
     ) -> str:
-        """Mirror the old Cluster.execute_command result handling."""
+        """Mirror the old Cluster.execute_command result handling.
+
+        Args:
+            result: Raw command execution response.
+            result_tag: Optional tag to extract a segment from output.
+            print_stdout: Whether to print stdout when using tags.
+
+        Returns:
+            The decoded output string.
+        """
         if not getattr(result, "results", None):
             raise RuntimeError("Command execution returned no results")
 

yggdrasil/databricks/compute/remote.py

@@ -1,3 +1,5 @@
+"""Convenience decorator for running functions on Databricks clusters."""
+
 import datetime as dt
 import logging
 from typing import (
@@ -26,6 +28,20 @@ def databricks_remote_compute(
     env_keys: Optional[List[str]] = None,
     **options
 ) -> Callable[[Callable[..., ReturnType]], Callable[..., ReturnType]]:
+    """Return a decorator that executes functions on a remote cluster.
+
+    Args:
+        cluster_id: Optional cluster id to target.
+        cluster_name: Optional cluster name to target.
+        workspace: Workspace instance or host string for lookup.
+        cluster: Pre-configured Cluster instance to reuse.
+        timeout: Optional execution timeout for remote calls.
+        env_keys: Optional environment variable names to forward.
+        **options: Extra options forwarded to the execution decorator.
+
+    Returns:
+        A decorator that runs functions on the resolved Databricks cluster.
+    """
     if isinstance(workspace, str):
         workspace = Workspace(host=workspace)
 

yggdrasil/databricks/jobs/__init__.py

@@ -0,0 +1,5 @@
+"""Helpers for running Databricks jobs and notebooks."""
+
+from .config import NotebookConfig, WidgetType
+
+__all__ = ["NotebookConfig", "WidgetType"]

yggdrasil/databricks/jobs/config.py

@@ -1,12 +1,14 @@
+"""Databricks widget-backed configuration helpers."""
+
 import builtins
 import dataclasses
 import datetime as dt
 import inspect
+import logging
 from dataclasses import dataclass, fields
 from enum import Enum
 from inspect import isclass
-from typing import Any, Dict, List, get_type_hints, Optional, get_origin
-import logging
+from typing import Any, Dict, List, get_type_hints, get_origin
 
 from ...libs.sparklib import SparkSession
 from ...types.cast.registry import convert
@@ -21,6 +23,15 @@ logger = logging.getLogger(__name__)
 
 
 def type_is_iterable(tpe: type, origin=None):
+    """Return True when the type annotation represents a list/set-like container.
+
+    Args:
+        tpe: The type annotation to inspect.
+        origin: Optional origin to reuse when recursing.
+
+    Returns:
+        True when the type is list-like, otherwise False.
+    """
     if (
         tpe is list or tpe is set
     ):
@@ -40,7 +51,7 @@ ALL_VALUES_TAG = "**all**"
 
 
 class WidgetType(Enum):
-    """Enum defining supported widget types in Databricks"""
+    """Enum defining supported Databricks widget types."""
     TEXT = "text"
     DROPDOWN = "dropdown"
     COMBOBOX = "combobox"
@@ -50,8 +61,15 @@ class WidgetType(Enum):
 
 @dataclass
 class NotebookConfig:
+    """Base class for widget-driven notebook configuration dataclasses."""
+
     @classmethod
     def get_dbutils(cls):
+        """Locate a ``dbutils`` instance from known Databricks injection points.
+
+        Returns:
+            The ``dbutils`` instance if found, otherwise None.
+        """
         # 1) explicit builtin injection (Databricks sometimes does this)
         if hasattr(builtins, "dbutils"):
             return builtins.dbutils
@@ -85,8 +103,7 @@ class NotebookConfig:
 
     @classmethod
     def from_environment(cls):
-        """
-        Build data class from environment variables or databricks dbutils if available.
+        """Build a config instance from Databricks widgets or environment variables.
 
         This method looks for values in the following order:
         1. Databricks widgets (if running in Databricks notebook)
@@ -94,7 +111,7 @@ class NotebookConfig:
         3. Environment variables
 
         Returns:
-            An instance of the dataclass populated with values from the environment
+            An instance of the dataclass populated with values from the environment.
         """
         dbutils = cls.get_dbutils()
         key_values: Dict[str, Any] = {}
@@ -229,6 +246,9 @@ class NotebookConfig:
 
         This method creates appropriate widgets for each field in the dataclass,
         with optional default values and customization options.
+
+        Returns:
+            None. Widgets are created in the notebook environment.
         """
         dbutils = cls.get_dbutils()
         if dbutils is None or not hasattr(dbutils, "widgets"):
@@ -296,6 +316,11 @@ class NotebookConfig:
 
     @classmethod
     def init_job(cls):
+        """Initialize widgets, tweak Spark session defaults, and return config.
+
+        Returns:
+            An instance of the dataclass populated from widgets or environment.
+        """
         cls.init_widgets()
 
         if SparkSession is not None:
@@ -308,31 +333,3 @@ class NotebookConfig:
                 spark_session.conf.set("spark.sql.session.timeZone", "UTC")
 
         return cls.from_environment()
-
-
-class ExampleEnum(Enum):
-    """Example enum for widget demonstration"""
-    OPTION1 = "option1"
-    OPTION2 = "option2"
-    OPTION3 = "option3"
-
-
-@dataclass
-class CompleteNotebookConfig(NotebookConfig):
-    """Example JobConfig with various field types to demonstrate widget handling"""
-    # Basic types
-    text_field: str
-    integer_field: int = 42
-    float_field: float = 3.14
-    boolean_field: bool = True
-
-    # Special types
-    date_field: dt.date = dt.date(2023, 1, 1)
-    datetime_field: dt.datetime = dt.datetime(2023, 1, 1, 12, 0, 0)
-    enum_field: ExampleEnum = ExampleEnum.OPTION1
-
-    # Collection types
-    list_of_strings: List[str] = None  # Will be displayed as multiselect
-
-    # Optional fields
-    optional_text: Optional[str] = None

yggdrasil/databricks/sql/__init__.py

@@ -1,3 +1,5 @@
+"""Databricks SQL helpers and engine wrappers."""
+
 from .engine import SQLEngine, StatementResult
 
 # Backwards compatibility