dagstermill 0.16.0__tar.gz → 0.27.9__tar.gz

{dagstermill-0.16.0 → dagstermill-0.27.9}/LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright {yyyy} {name of copyright owner}
+   Copyright 2025 Dagster Labs, Inc.
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

dagstermill-0.27.9/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: dagstermill
+Version: 0.27.9
+Summary: run notebooks using the Dagster tools
+Author: Dagster Labs
+Author-email: hello@dagsterlabs.com
+License: Apache-2.0
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9,<3.14
+License-File: LICENSE
+Requires-Dist: dagster==1.11.9
+Requires-Dist: ipykernel!=5.4.0,!=5.4.1,>=4.9.0
+Requires-Dist: ipython_genutils>=0.2.0
+Requires-Dist: packaging>=20.9
+Requires-Dist: papermill>=1.0.0
+Requires-Dist: scrapbook>=0.5.0
+Requires-Dist: nbconvert
+Requires-Dist: jupyter-client<8
+Provides-Extra: test
+Requires-Dist: matplotlib; extra == "test"
+Requires-Dist: scikit-learn>=0.19.0; extra == "test"
+Requires-Dist: tqdm<=4.48; extra == "test"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

dagstermill-0.27.9/README.md

@@ -0,0 +1,4 @@
+# dagstermill
+
+The docs for `dagstermill` can be found
+[here](https://docs.dagster.io/api/python-api/libraries/dagstermill).

dagstermill-0.27.9/dagstermill/__init__.py

@@ -0,0 +1,27 @@
+from dagster_shared.libraries import DagsterLibraryRegistry
+
+import dagstermill.factory as factory  # noqa: F401
+from dagstermill.asset_factory import define_dagstermill_asset as define_dagstermill_asset
+from dagstermill.context import DagstermillExecutionContext as DagstermillExecutionContext
+from dagstermill.errors import DagstermillError as DagstermillError
+from dagstermill.factory import define_dagstermill_op as define_dagstermill_op
+from dagstermill.io_managers import (
+    ConfigurableLocalOutputNotebookIOManager as ConfigurableLocalOutputNotebookIOManager,
+    local_output_notebook_io_manager as local_output_notebook_io_manager,
+)
+from dagstermill.manager import MANAGER_FOR_NOTEBOOK_INSTANCE as _MANAGER_FOR_NOTEBOOK_INSTANCE
+from dagstermill.version import __version__ as __version__
+
+DagsterLibraryRegistry.register("dagstermill", __version__)
+
+get_context = _MANAGER_FOR_NOTEBOOK_INSTANCE.get_context
+
+yield_result = _MANAGER_FOR_NOTEBOOK_INSTANCE.yield_result
+
+yield_event = _MANAGER_FOR_NOTEBOOK_INSTANCE.yield_event
+
+_reconstitute_job_context = _MANAGER_FOR_NOTEBOOK_INSTANCE.reconstitute_job_context
+
+_teardown = _MANAGER_FOR_NOTEBOOK_INSTANCE.teardown_resources
+
+_load_input_parameter = _MANAGER_FOR_NOTEBOOK_INSTANCE.load_input_parameter

dagstermill-0.27.9/dagstermill/__main__.py

@@ -0,0 +1,3 @@
+from dagstermill.cli import main
+
+main()

dagstermill-0.27.9/dagstermill/asset_factory.py

@@ -0,0 +1,227 @@
+import pickle
+import tempfile
+from collections.abc import Iterable, Mapping
+from typing import Any, Callable, Optional, Union, cast
+
+import dagster._check as check
+from dagster import (
+    AssetIn,
+    AssetKey,
+    AssetsDefinition,
+    Failure,
+    Output,
+    PartitionsDefinition,
+    ResourceDefinition,
+    RetryPolicy,
+    RetryRequested,
+    SourceAsset,
+    asset,
+)
+from dagster._annotations import beta, beta_param
+from dagster._config.pythonic_config import Config, infer_schema_from_config_class
+from dagster._config.pythonic_config.type_check_utils import safe_is_subclass
+from dagster._core.definitions.events import CoercibleToAssetKey, CoercibleToAssetKeyPrefix
+from dagster._core.execution.context.compute import OpExecutionContext
+from dagster._core.storage.tags import COMPUTE_KIND_TAG
+from dagster._utils.tags import normalize_tags
+
+from dagstermill.factory import _clean_path_for_windows, execute_notebook
+
+
+def _make_dagstermill_asset_compute_fn(
+    name: str,
+    notebook_path: str,
+    save_notebook_on_failure: bool,
+) -> Callable:
+    def _t_fn(context: OpExecutionContext, **inputs) -> Iterable:
+        check.param_invariant(
+            isinstance(context.run_config, dict),
+            "context",
+            "StepExecutionContext must have valid run_config",
+        )
+
+        with tempfile.TemporaryDirectory() as output_notebook_dir:
+            executed_notebook_path = execute_notebook(
+                context.get_step_execution_context(),
+                name=name,
+                inputs=inputs,
+                save_notebook_on_failure=save_notebook_on_failure,
+                notebook_path=notebook_path,
+                output_notebook_dir=output_notebook_dir,
+            )
+
+            with open(executed_notebook_path, "rb") as fd:
+                yield Output(fd.read())
+
+            # deferred import for perf
+            import scrapbook
+
+            output_nb = scrapbook.read_notebook(executed_notebook_path)
+
+            for key, value in output_nb.scraps.items():
+                if key.startswith("event-"):
+                    with open(value.data, "rb") as fd:
+                        event = pickle.loads(fd.read())
+                        if isinstance(event, (Failure, RetryRequested)):
+                            raise event
+                        else:
+                            yield event
+
+    return _t_fn
+
+
+@beta
+@beta_param(param="resource_defs")
+def define_dagstermill_asset(
+    name: str,
+    notebook_path: str,
+    key_prefix: Optional[CoercibleToAssetKeyPrefix] = None,
+    ins: Optional[Mapping[str, AssetIn]] = None,
+    deps: Optional[Iterable[Union[CoercibleToAssetKey, AssetsDefinition, SourceAsset]]] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    config_schema: Optional[Union[Any, Mapping[str, Any]]] = None,
+    required_resource_keys: Optional[set[str]] = None,
+    resource_defs: Optional[Mapping[str, ResourceDefinition]] = None,
+    description: Optional[str] = None,
+    partitions_def: Optional[PartitionsDefinition] = None,
+    op_tags: Optional[Mapping[str, Any]] = None,
+    group_name: Optional[str] = None,
+    io_manager_key: Optional[str] = None,
+    retry_policy: Optional[RetryPolicy] = None,
+    save_notebook_on_failure: bool = False,
+    non_argument_deps: Optional[Union[set[AssetKey], set[str]]] = None,
+    asset_tags: Optional[Mapping[str, Any]] = None,
+) -> AssetsDefinition:
+    """Creates a Dagster asset for a Jupyter notebook.
+
+    Arguments:
+        name (str): The name for the asset
+        notebook_path (str): Path to the backing notebook
+        key_prefix (Optional[Union[str, Sequence[str]]]): If provided, the asset's key is the
+            concatenation of the key_prefix and the asset's name, which defaults to the name of
+            the decorated function. Each item in key_prefix must be a valid name in dagster (ie only
+            contains letters, numbers, and _) and may not contain python reserved keywords.
+        ins (Optional[Mapping[str, AssetIn]]): A dictionary that maps input names to information
+            about the input.
+        deps (Optional[Sequence[Union[AssetsDefinition, SourceAsset, AssetKey, str]]]): The assets
+            that are upstream dependencies, but do not pass an input value to the notebook.
+        config_schema (Optional[ConfigSchema): The configuration schema for the asset's underlying
+            op. If set, Dagster will check that config provided for the op matches this schema and fail
+            if it does not. If not set, Dagster will accept any config provided for the op.
+        metadata (Optional[Dict[str, Any]]): A dict of metadata entries for the asset.
+        required_resource_keys (Optional[Set[str]]): Set of resource handles required by the notebook.
+        description (Optional[str]): Description of the asset to display in the Dagster UI.
+        partitions_def (Optional[PartitionsDefinition]): Defines the set of partition keys that
+            compose the asset.
+        op_tags (Optional[Dict[str, Any]]): A dictionary of tags for the op that computes the asset.
+            Frameworks may expect and require certain metadata to be attached to a op. Values that
+            are not strings will be json encoded and must meet the criteria that
+            `json.loads(json.dumps(value)) == value`.
+        group_name (Optional[str]): A string name used to organize multiple assets into groups. If not provided,
+            the name "default" is used.
+        resource_defs (Optional[Mapping[str, ResourceDefinition]]):
+            (Beta) A mapping of resource keys to resource definitions. These resources
+            will be initialized during execution, and can be accessed from the
+            context within the notebook.
+        io_manager_key (Optional[str]): A string key for the IO manager used to store the output notebook.
+            If not provided, the default key output_notebook_io_manager will be used.
+        retry_policy (Optional[RetryPolicy]): The retry policy for the op that computes the asset.
+        save_notebook_on_failure (bool): If True and the notebook fails during execution, the failed notebook will be
+            written to the Dagster storage directory. The location of the file will be printed in the Dagster logs.
+            Defaults to False.
+        asset_tags (Optional[Dict[str, Any]]): A dictionary of tags to apply to the asset.
+        non_argument_deps (Optional[Union[Set[AssetKey], Set[str]]]): Deprecated, use deps instead. Set of asset keys that are
+            upstream dependencies, but do not pass an input to the asset.
+
+    Examples:
+        .. code-block:: python
+
+            from dagstermill import define_dagstermill_asset
+            from dagster import asset, AssetIn, AssetKey
+            from sklearn import datasets
+            import pandas as pd
+            import numpy as np
+
+            @asset
+            def iris_dataset():
+                sk_iris = datasets.load_iris()
+                return pd.DataFrame(
+                    data=np.c_[sk_iris["data"], sk_iris["target"]],
+                    columns=sk_iris["feature_names"] + ["target"],
+                )
+
+            iris_kmeans_notebook = define_dagstermill_asset(
+                name="iris_kmeans_notebook",
+                notebook_path="/path/to/iris_kmeans.ipynb",
+                ins={
+                    "iris": AssetIn(key=AssetKey("iris_dataset"))
+                }
+            )
+    """
+    check.str_param(name, "name")
+    check.str_param(notebook_path, "notebook_path")
+    check.bool_param(save_notebook_on_failure, "save_notebook_on_failure")
+
+    required_resource_keys = set(
+        check.opt_set_param(required_resource_keys, "required_resource_keys", of_type=str)
+    )
+    ins = check.opt_mapping_param(ins, "ins", key_type=str, value_type=AssetIn)
+
+    if isinstance(key_prefix, str):
+        key_prefix = [key_prefix]
+
+    key_prefix = check.opt_list_param(key_prefix, "key_prefix", of_type=str)
+
+    default_description = f"This asset is backed by the notebook at {notebook_path}"
+    description = check.opt_str_param(description, "description", default=default_description)
+
+    io_mgr_key = check.opt_str_param(
+        io_manager_key, "io_manager_key", default="output_notebook_io_manager"
+    )
+
+    user_tags = normalize_tags(op_tags)
+    if op_tags is not None:
+        check.invariant(
+            "notebook_path" not in op_tags,
+            "user-defined op tags contains the `notebook_path` key, but the `notebook_path` key"
+            " is reserved for use by Dagster",
+        )
+        check.invariant(
+            COMPUTE_KIND_TAG not in op_tags,
+            f"user-defined op tags contains the `{COMPUTE_KIND_TAG}` key, but the `{COMPUTE_KIND_TAG}` key is reserved for"
+            " use by Dagster",
+        )
+
+    default_tags = {
+        "notebook_path": _clean_path_for_windows(notebook_path),
+        COMPUTE_KIND_TAG: "ipynb",
+    }
+
+    if safe_is_subclass(config_schema, Config):
+        config_schema = infer_schema_from_config_class(cast("type[Config]", config_schema))
+
+    return asset(
+        name=name,
+        key_prefix=key_prefix,
+        ins=ins,
+        deps=deps,
+        metadata=metadata,
+        description=description,
+        config_schema=config_schema,
+        required_resource_keys=required_resource_keys,
+        resource_defs=resource_defs,
+        partitions_def=partitions_def,
+        op_tags={**user_tags, **default_tags},
+        group_name=group_name,
+        output_required=False,
+        io_manager_key=io_mgr_key,
+        retry_policy=retry_policy,
+        non_argument_deps=non_argument_deps,
+        tags=asset_tags,
+    )(
+        _make_dagstermill_asset_compute_fn(
+            name=name,
+            notebook_path=notebook_path,
+            save_notebook_on_failure=save_notebook_on_failure,
+        )
+    )

{dagstermill-0.16.0 → dagstermill-0.27.9}/dagstermill/cli.py

@@ -1,15 +1,15 @@
 import copy
 import os
 import subprocess
-from typing import Dict, Optional
+from collections.abc import Mapping
+from typing import Optional
 
 import click
-import nbformat
-from papermill.iorw import load_notebook_node, write_ipynb
-
 import dagster._check as check
-from dagster._seven.json import loads
+import nbformat
 from dagster._utils import mkdir_p, safe_isfile
+from dagster_shared.seven.json import loads
+from papermill.iorw import load_notebook_node, write_ipynb
 
 
 def get_import_cell():
@@ -35,15 +35,13 @@ def get_kernelspec(kernel: Optional[str] = None):
             )
         ) + list(kernelspecs["kernelspecs"].keys())
         kernel = preferred_kernels[0]
-        print(  # pylint: disable=print-call
-            "No kernel specified, defaulting to '{kernel}'".format(kernel=kernel)
-        )
+        print(f"No kernel specified, defaulting to '{kernel}'")  # noqa: T201
 
     check.invariant(
         kernel in kernelspecs["kernelspecs"],
         "Could not find kernel '{kernel}': available kernels are [{kernels}]".format(
             kernel=kernel,
-            kernels=", ".join(["'{k}'".format(k=k) for k in kernelspecs["kernelspecs"]]),
+            kernels=", ".join([f"'{k}'" for k in kernelspecs["kernelspecs"]]),
         ),
     )
 
@@ -54,8 +52,8 @@ def get_kernelspec(kernel: Optional[str] = None):
     }
 
 
-def get_notebook_scaffolding(kernelspec: Dict[str, str]):
-    check.dict_param(kernelspec, "kernelspec", key_type=str, value_type=str)
+def get_notebook_scaffolding(kernelspec: Mapping[str, str]):
+    check.mapping_param(kernelspec, "kernelspec", key_type=str, value_type=str)
 
     notebook = nbformat.v4.new_notebook()
 
@@ -69,7 +67,7 @@ def get_notebook_scaffolding(kernelspec: Dict[str, str]):
 
 
 @click.command(
-    name="register-notebook", help=("Scaffolds existing notebook for dagstermill compatibility")
+    name="register-notebook", help="Scaffolds existing notebook for dagstermill compatibility"
 )
 @click.option("--notebook", "-note", type=click.Path(exists=True), help="Path to existing notebook")
 def retroactively_scaffold_notebook(notebook: str):
@@ -83,7 +81,7 @@ def execute_retroactive_scaffold(notebook_path: str):
     write_ipynb(new_nb, notebook_path)
 
 
-@click.command(name="create-notebook", help=("Creates new dagstermill notebook."))
+@click.command(name="create-notebook", help="Creates new dagstermill notebook.")
 @click.option("--notebook", "-note", type=click.Path(), help="Name of notebook")
 @click.option(
     "--force-overwrite",
@@ -113,17 +111,15 @@ def execute_create_notebook(notebook: str, force_overwrite: bool, kernel: str):
 
     if not force_overwrite and safe_isfile(notebook_path):
         click.confirm(
-            (
-                "Warning, {notebook_path} already exists and continuing "
-                "will overwrite the existing notebook. "
-                "Are you sure you want to continue?"
-            ).format(notebook_path=notebook_path),
+            f"Warning, {notebook_path} already exists and continuing "
+            "will overwrite the existing notebook. "
+            "Are you sure you want to continue?",
             abort=True,
         )
 
     with open(notebook_path, "w", encoding="utf8") as f:
         f.write(get_notebook_scaffolding(get_kernelspec(kernel)))
-        click.echo("Created new dagstermill notebook at {path}".format(path=notebook_path))
+        click.echo(f"Created new dagstermill notebook at {notebook_path}")
 
 
 def create_dagstermill_cli():

{dagstermill-0.16.0 → dagstermill-0.27.9}/dagstermill/compat.py

@@ -1,17 +1,17 @@
 import papermill
-from packaging.version import LegacyVersion, parse
+from packaging.version import Version, parse
 from papermill.exceptions import PapermillExecutionError
 
 
 def is_papermill_2():
     version = parse(papermill.__version__)
-    if isinstance(version, LegacyVersion):
-        return False
+    # satisfies typechecker that might think version is a LegacyVersion
+    assert isinstance(version, Version)
     return version.major == 2
 
 
 if is_papermill_2():
-    from nbclient.exceptions import CellExecutionError  # pylint: disable=import-error
+    from nbclient.exceptions import CellExecutionError
 
     ExecutionError = (PapermillExecutionError, CellExecutionError)
 

dagstermill-0.27.9/dagstermill/context.py

@@ -0,0 +1,191 @@
+from collections.abc import Mapping
+from typing import AbstractSet, Any, Optional, cast  # noqa: UP035
+
+from dagster import (
+    DagsterRun,
+    JobDefinition,
+    OpDefinition,
+    _check as check,
+)
+from dagster._annotations import beta, public
+from dagster._core.definitions.dependency import Node, NodeHandle
+from dagster._core.definitions.repository_definition.repository_definition import (
+    RepositoryDefinition,
+)
+from dagster._core.execution.context.op_execution_context import AbstractComputeExecutionContext
+from dagster._core.execution.context.system import PlanExecutionContext, StepExecutionContext
+from dagster._core.log_manager import DagsterLogManager
+from dagster._core.system_config.objects import ResolvedRunConfig
+
+
+@beta
+class DagstermillExecutionContext(AbstractComputeExecutionContext):
+    """Dagstermill-specific execution context.
+
+    Do not initialize directly: use :func:`dagstermill.get_context`.
+    """
+
+    def __init__(
+        self,
+        job_context: PlanExecutionContext,
+        job_def: JobDefinition,
+        resource_keys_to_init: AbstractSet[str],
+        op_name: str,
+        node_handle: NodeHandle,
+        op_config: Any = None,
+    ):
+        self._job_context = check.inst_param(job_context, "job_context", PlanExecutionContext)
+        self._job_def = check.inst_param(job_def, "job_def", JobDefinition)
+        self._resource_keys_to_init = check.set_param(
+            resource_keys_to_init, "resource_keys_to_init", of_type=str
+        )
+        self.op_name = check.str_param(op_name, "op_name")
+        self.node_handle = check.inst_param(node_handle, "node_handle", NodeHandle)
+        self._op_config = op_config
+
+    def has_tag(self, key: str) -> bool:
+        """Check if a logging tag is defined on the context.
+
+        Args:
+            key (str): The key to check.
+
+        Returns:
+            bool
+        """
+        check.str_param(key, "key")
+        return self._job_context.has_tag(key)
+
+    def get_tag(self, key: str) -> Optional[str]:
+        """Get a logging tag defined on the context.
+
+        Args:
+            key (str): The key to get.
+
+        Returns:
+            str
+        """
+        check.str_param(key, "key")
+        return self._job_context.get_tag(key)
+
+    @public
+    @property
+    def run_id(self) -> str:
+        """str: The run_id for the context."""
+        return self._job_context.run_id
+
+    @public
+    @property
+    def run_config(self) -> Mapping[str, Any]:
+        """dict: The run_config for the context."""
+        return self._job_context.run_config
+
+    @property
+    def resolved_run_config(self) -> ResolvedRunConfig:
+        """:class:`dagster.ResolvedRunConfig`: The resolved_run_config for the context."""
+        return self._job_context.resolved_run_config
+
+    @public
+    @property
+    def logging_tags(self) -> Mapping[str, str]:
+        """dict: The logging tags for the context."""
+        return self._job_context.logging_tags
+
+    @public
+    @property
+    def job_name(self) -> str:
+        """str: The name of the executing job."""
+        return self._job_context.job_name
+
+    @public
+    @property
+    def job_def(self) -> JobDefinition:
+        """:class:`dagster.JobDefinition`: The job definition for the context.
+
+        This will be a dagstermill-specific shim.
+        """
+        return self._job_def
+
+    @property
+    def repository_def(self) -> RepositoryDefinition:
+        """:class:`dagster.RepositoryDefinition`: The repository definition for the context."""
+        raise NotImplementedError
+
+    @property
+    def resources(self) -> Any:
+        """collections.namedtuple: A dynamically-created type whose properties allow access to
+        resources.
+        """
+        return self._job_context.scoped_resources_builder.build(
+            required_resource_keys=self._resource_keys_to_init,
+        )
+
+    @public
+    @property
+    def run(self) -> DagsterRun:
+        """:class:`dagster.DagsterRun`: The job run for the context."""
+        return cast("DagsterRun", self._job_context.dagster_run)
+
+    @property
+    def log(self) -> DagsterLogManager:
+        """:class:`dagster.DagsterLogManager`: The log manager for the context.
+
+        Call, e.g., ``log.info()`` to log messages through the Dagster machinery.
+        """
+        return self._job_context.log
+
+    @public
+    @property
+    def op_def(self) -> OpDefinition:
+        """:class:`dagster.OpDefinition`: The op definition for the context.
+
+        In interactive contexts, this may be a dagstermill-specific shim, depending whether an
+        op definition was passed to ``dagstermill.get_context``.
+        """
+        return cast("OpDefinition", self._job_def.node_def_named(self.op_name))
+
+    @property
+    def node(self) -> Node:
+        """:class:`dagster.Node`: The node for the context.
+
+        In interactive contexts, this may be a dagstermill-specific shim, depending whether an
+        op definition was passed to ``dagstermill.get_context``.
+        """
+        return self.job_def.get_node(self.node_handle)
+
+    @public
+    @property
+    def op_config(self) -> Any:
+        """collections.namedtuple: A dynamically-created type whose properties allow access to
+        op-specific config.
+        """
+        if self._op_config:
+            return self._op_config
+
+        op_config = self.resolved_run_config.ops.get(self.op_name)
+        return op_config.config if op_config else None
+
+
+class DagstermillRuntimeExecutionContext(DagstermillExecutionContext):
+    def __init__(
+        self,
+        job_context: PlanExecutionContext,
+        job_def: JobDefinition,
+        resource_keys_to_init: AbstractSet[str],
+        op_name: str,
+        step_context: StepExecutionContext,
+        node_handle: NodeHandle,
+        op_config: Any = None,
+    ):
+        self._step_context = check.inst_param(step_context, "step_context", StepExecutionContext)
+        super().__init__(
+            job_context,
+            job_def,
+            resource_keys_to_init,
+            op_name,
+            node_handle,
+            op_config,
+        )
+
+    @property
+    def step_context(self) -> StepExecutionContext:
+        return self._step_context