dagstermill 0.23.3__py3-none-any.whl → 0.27.6__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of dagstermill might be problematic. Click here for more details.

dagstermill/__init__.py CHANGED
@@ -1,15 +1,16 @@
1
- from dagster._core.libraries import DagsterLibraryRegistry
2
-
3
- from .asset_factory import define_dagstermill_asset as define_dagstermill_asset
4
- from .context import DagstermillExecutionContext as DagstermillExecutionContext
5
- from .errors import DagstermillError as DagstermillError
6
- from .factory import define_dagstermill_op as define_dagstermill_op
7
- from .io_managers import (
1
+ from dagster_shared.libraries import DagsterLibraryRegistry
2
+
3
+ import dagstermill.factory as factory # noqa: F401
4
+ from dagstermill.asset_factory import define_dagstermill_asset as define_dagstermill_asset
5
+ from dagstermill.context import DagstermillExecutionContext as DagstermillExecutionContext
6
+ from dagstermill.errors import DagstermillError as DagstermillError
7
+ from dagstermill.factory import define_dagstermill_op as define_dagstermill_op
8
+ from dagstermill.io_managers import (
8
9
  ConfigurableLocalOutputNotebookIOManager as ConfigurableLocalOutputNotebookIOManager,
9
10
  local_output_notebook_io_manager as local_output_notebook_io_manager,
10
11
  )
11
- from .manager import MANAGER_FOR_NOTEBOOK_INSTANCE as _MANAGER_FOR_NOTEBOOK_INSTANCE
12
- from .version import __version__ as __version__
12
+ from dagstermill.manager import MANAGER_FOR_NOTEBOOK_INSTANCE as _MANAGER_FOR_NOTEBOOK_INSTANCE
13
+ from dagstermill.version import __version__ as __version__
13
14
 
14
15
  DagsterLibraryRegistry.register("dagstermill", __version__)
15
16
 
dagstermill/__main__.py CHANGED
@@ -1,3 +1,3 @@
1
- from .cli import main
1
+ from dagstermill.cli import main
2
2
 
3
3
  main()
@@ -1,6 +1,7 @@
1
1
  import pickle
2
2
  import tempfile
3
- from typing import Any, Callable, Iterable, Mapping, Optional, Set, Type, Union, cast
3
+ from collections.abc import Iterable, Mapping
4
+ from typing import Any, Callable, Optional, Union, cast
4
5
 
5
6
  import dagster._check as check
6
7
  from dagster import (
@@ -16,11 +17,13 @@ from dagster import (
16
17
  SourceAsset,
17
18
  asset,
18
19
  )
20
+ from dagster._annotations import beta, beta_param
19
21
  from dagster._config.pythonic_config import Config, infer_schema_from_config_class
20
22
  from dagster._config.pythonic_config.type_check_utils import safe_is_subclass
21
23
  from dagster._core.definitions.events import CoercibleToAssetKey, CoercibleToAssetKeyPrefix
22
- from dagster._core.definitions.utils import normalize_tags
23
24
  from dagster._core.execution.context.compute import OpExecutionContext
25
+ from dagster._core.storage.tags import COMPUTE_KIND_TAG
26
+ from dagster._utils.tags import normalize_tags
24
27
 
25
28
  from dagstermill.factory import _clean_path_for_windows, execute_notebook
26
29
 
@@ -67,6 +70,8 @@ def _make_dagstermill_asset_compute_fn(
67
70
  return _t_fn
68
71
 
69
72
 
73
+ @beta
74
+ @beta_param(param="resource_defs")
70
75
  def define_dagstermill_asset(
71
76
  name: str,
72
77
  notebook_path: str,
@@ -75,7 +80,7 @@ def define_dagstermill_asset(
75
80
  deps: Optional[Iterable[Union[CoercibleToAssetKey, AssetsDefinition, SourceAsset]]] = None,
76
81
  metadata: Optional[Mapping[str, Any]] = None,
77
82
  config_schema: Optional[Union[Any, Mapping[str, Any]]] = None,
78
- required_resource_keys: Optional[Set[str]] = None,
83
+ required_resource_keys: Optional[set[str]] = None,
79
84
  resource_defs: Optional[Mapping[str, ResourceDefinition]] = None,
80
85
  description: Optional[str] = None,
81
86
  partitions_def: Optional[PartitionsDefinition] = None,
@@ -84,7 +89,8 @@ def define_dagstermill_asset(
84
89
  io_manager_key: Optional[str] = None,
85
90
  retry_policy: Optional[RetryPolicy] = None,
86
91
  save_notebook_on_failure: bool = False,
87
- non_argument_deps: Optional[Union[Set[AssetKey], Set[str]]] = None,
92
+ non_argument_deps: Optional[Union[set[AssetKey], set[str]]] = None,
93
+ asset_tags: Optional[Mapping[str, Any]] = None,
88
94
  ) -> AssetsDefinition:
89
95
  """Creates a Dagster asset for a Jupyter notebook.
90
96
 
@@ -114,7 +120,7 @@ def define_dagstermill_asset(
114
120
  group_name (Optional[str]): A string name used to organize multiple assets into groups. If not provided,
115
121
  the name "default" is used.
116
122
  resource_defs (Optional[Mapping[str, ResourceDefinition]]):
117
- (Experimental) A mapping of resource keys to resource definitions. These resources
123
+ (Beta) A mapping of resource keys to resource definitions. These resources
118
124
  will be initialized during execution, and can be accessed from the
119
125
  context within the notebook.
120
126
  io_manager_key (Optional[str]): A string key for the IO manager used to store the output notebook.
@@ -123,6 +129,7 @@ def define_dagstermill_asset(
123
129
  save_notebook_on_failure (bool): If True and the notebook fails during execution, the failed notebook will be
124
130
  written to the Dagster storage directory. The location of the file will be printed in the Dagster logs.
125
131
  Defaults to False.
132
+ asset_tags (Optional[Dict[str, Any]]): A dictionary of tags to apply to the asset.
126
133
  non_argument_deps (Optional[Union[Set[AssetKey], Set[str]]]): Deprecated, use deps instead. Set of asset keys that are
127
134
  upstream dependencies, but do not pass an input to the asset.
128
135
 
@@ -172,7 +179,7 @@ def define_dagstermill_asset(
172
179
  io_manager_key, "io_manager_key", default="output_notebook_io_manager"
173
180
  )
174
181
 
175
- user_tags = normalize_tags(op_tags).tags
182
+ user_tags = normalize_tags(op_tags)
176
183
  if op_tags is not None:
177
184
  check.invariant(
178
185
  "notebook_path" not in op_tags,
@@ -180,15 +187,18 @@ def define_dagstermill_asset(
180
187
  " is reserved for use by Dagster",
181
188
  )
182
189
  check.invariant(
183
- "kind" not in op_tags,
184
- "user-defined op tags contains the `kind` key, but the `kind` key is reserved for"
190
+ COMPUTE_KIND_TAG not in op_tags,
191
+ f"user-defined op tags contains the `{COMPUTE_KIND_TAG}` key, but the `{COMPUTE_KIND_TAG}` key is reserved for"
185
192
  " use by Dagster",
186
193
  )
187
194
 
188
- default_tags = {"notebook_path": _clean_path_for_windows(notebook_path), "kind": "ipynb"}
195
+ default_tags = {
196
+ "notebook_path": _clean_path_for_windows(notebook_path),
197
+ COMPUTE_KIND_TAG: "ipynb",
198
+ }
189
199
 
190
200
  if safe_is_subclass(config_schema, Config):
191
- config_schema = infer_schema_from_config_class(cast(Type[Config], config_schema))
201
+ config_schema = infer_schema_from_config_class(cast("type[Config]", config_schema))
192
202
 
193
203
  return asset(
194
204
  name=name,
@@ -207,6 +217,7 @@ def define_dagstermill_asset(
207
217
  io_manager_key=io_mgr_key,
208
218
  retry_policy=retry_policy,
209
219
  non_argument_deps=non_argument_deps,
220
+ tags=asset_tags,
210
221
  )(
211
222
  _make_dagstermill_asset_compute_fn(
212
223
  name=name,
dagstermill/cli.py CHANGED
@@ -1,13 +1,14 @@
1
1
  import copy
2
2
  import os
3
3
  import subprocess
4
- from typing import Mapping, Optional
4
+ from collections.abc import Mapping
5
+ from typing import Optional
5
6
 
6
7
  import click
7
8
  import dagster._check as check
8
9
  import nbformat
9
- from dagster._seven.json import loads
10
10
  from dagster._utils import mkdir_p, safe_isfile
11
+ from dagster_shared.seven.json import loads
11
12
  from papermill.iorw import load_notebook_node, write_ipynb
12
13
 
13
14
 
dagstermill/context.py CHANGED
@@ -1,4 +1,5 @@
1
- from typing import AbstractSet, Any, Mapping, Optional, cast
1
+ from collections.abc import Mapping
2
+ from typing import AbstractSet, Any, Optional, cast # noqa: UP035
2
3
 
3
4
  from dagster import (
4
5
  DagsterRun,
@@ -6,17 +7,18 @@ from dagster import (
6
7
  OpDefinition,
7
8
  _check as check,
8
9
  )
9
- from dagster._annotations import public
10
+ from dagster._annotations import beta, public
10
11
  from dagster._core.definitions.dependency import Node, NodeHandle
11
12
  from dagster._core.definitions.repository_definition.repository_definition import (
12
13
  RepositoryDefinition,
13
14
  )
14
- from dagster._core.execution.context.compute import AbstractComputeExecutionContext
15
+ from dagster._core.execution.context.op_execution_context import AbstractComputeExecutionContext
15
16
  from dagster._core.execution.context.system import PlanExecutionContext, StepExecutionContext
16
17
  from dagster._core.log_manager import DagsterLogManager
17
18
  from dagster._core.system_config.objects import ResolvedRunConfig
18
19
 
19
20
 
21
+ @beta
20
22
  class DagstermillExecutionContext(AbstractComputeExecutionContext):
21
23
  """Dagstermill-specific execution context.
22
24
 
@@ -121,7 +123,7 @@ class DagstermillExecutionContext(AbstractComputeExecutionContext):
121
123
  @property
122
124
  def run(self) -> DagsterRun:
123
125
  """:class:`dagster.DagsterRun`: The job run for the context."""
124
- return cast(DagsterRun, self._job_context.dagster_run)
126
+ return cast("DagsterRun", self._job_context.dagster_run)
125
127
 
126
128
  @property
127
129
  def log(self) -> DagsterLogManager:
@@ -139,7 +141,7 @@ class DagstermillExecutionContext(AbstractComputeExecutionContext):
139
141
  In interactive contexts, this may be a dagstermill-specific shim, depending whether an
140
142
  op definition was passed to ``dagstermill.get_context``.
141
143
  """
142
- return cast(OpDefinition, self._job_def.node_def_named(self.op_name))
144
+ return cast("OpDefinition", self._job_def.node_def_named(self.op_name))
143
145
 
144
146
  @property
145
147
  def node(self) -> Node:
dagstermill/engine.py CHANGED
@@ -1,7 +1,7 @@
1
1
  import nbformat
2
2
  from papermill.log import logger
3
3
 
4
- from .compat import ExecutionError, is_papermill_2
4
+ from dagstermill.compat import ExecutionError, is_papermill_2
5
5
 
6
6
  if is_papermill_2():
7
7
  from papermill.clientwrap import PapermillNotebookClient
@@ -86,9 +86,7 @@ else:
86
86
  # the kernel down. Note that atexit doesn't seem to work at all in ipython, and hooking into
87
87
  # the ipython post_execute event doesn't work in papermill.
88
88
  def papermill_process(self, nb_man, resources):
89
- _, resources = super(DagstermillExecutePreprocessor, self).papermill_process(
90
- nb_man, resources
91
- )
89
+ _, resources = super().papermill_process(nb_man, resources)
92
90
 
93
91
  new_cell = nbformat.v4.new_code_cell(
94
92
  source="import dagstermill as __dm_dagstermill\n__dm_dagstermill._teardown()\n"
dagstermill/errors.py CHANGED
@@ -1,4 +1,4 @@
1
- from dagster._core.errors import DagsterError
1
+ from dagster_shared.error import DagsterError
2
2
 
3
3
 
4
4
  class DagstermillError(DagsterError):
@@ -1 +1 @@
1
- from .repository import notebook_repo as notebook_repo
1
+ from dagstermill.examples.repository import notebook_repo as notebook_repo
@@ -23,7 +23,7 @@ from dagster import (
23
23
  with_resources,
24
24
  )
25
25
  from dagster._config.pythonic_config import Config
26
- from dagster._core.definitions.asset_graph import AssetGraph
26
+ from dagster._core.definitions.assets.graph.asset_graph import AssetGraph
27
27
  from dagster._core.definitions.utils import DEFAULT_OUTPUT
28
28
  from dagster._utils import PICKLE_PROTOCOL, file_relative_path
29
29
 
@@ -303,12 +303,12 @@ def error_job():
303
303
 
304
304
  if DAGSTER_PANDAS_PRESENT and SKLEARN_PRESENT and MATPLOTLIB_PRESENT:
305
305
  # We need type-ignores here because type checkers don't understand the `*_PRESENT` kwargs.
306
- clean_data = test_nb_op("clean_data", outs={DEFAULT_OUTPUT: Out(DataFrame)})
306
+ clean_data = test_nb_op("clean_data", outs={DEFAULT_OUTPUT: Out(DataFrame)}) # pyright: ignore[reportPossiblyUnboundVariable]
307
307
 
308
308
  # FIXME add an output to this
309
- tutorial_LR = test_nb_op("tutorial_LR", ins={"df": In(DataFrame)})
309
+ tutorial_LR = test_nb_op("tutorial_LR", ins={"df": In(DataFrame)}) # pyright: ignore[reportPossiblyUnboundVariable]
310
310
 
311
- tutorial_RF = test_nb_op("tutorial_RF", ins={"df": In(DataFrame)})
311
+ tutorial_RF = test_nb_op("tutorial_RF", ins={"df": In(DataFrame)}) # pyright: ignore[reportPossiblyUnboundVariable]
312
312
 
313
313
  @job(resource_defs=common_resource_defs)
314
314
  def tutorial_job():
@@ -570,6 +570,7 @@ custom_io_mgr_key_asset = dagstermill.define_dagstermill_asset(
570
570
  name="custom_io_mgr_key",
571
571
  notebook_path=nb_test_path("hello_world"),
572
572
  io_manager_key="my_custom_io_manager",
573
+ asset_tags={"foo": "bar"},
573
574
  )
574
575
 
575
576
  yield_event_asset = dagstermill.define_dagstermill_asset(
dagstermill/factory.py CHANGED
@@ -4,7 +4,8 @@ import pickle
4
4
  import sys
5
5
  import tempfile
6
6
  import uuid
7
- from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, Set, Type, Union, cast
7
+ from collections.abc import Iterable, Mapping, Sequence
8
+ from typing import Any, Callable, Optional, Union, cast
8
9
 
9
10
  import nbformat
10
11
  import papermill
@@ -14,29 +15,31 @@ from dagster import (
14
15
  Out,
15
16
  Output,
16
17
  _check as check,
17
- _seven,
18
18
  )
19
+ from dagster._annotations import beta
19
20
  from dagster._config.pythonic_config import Config, infer_schema_from_config_class
20
21
  from dagster._config.pythonic_config.type_check_utils import safe_is_subclass
21
22
  from dagster._core.definitions.events import AssetMaterialization, Failure, RetryRequested
22
23
  from dagster._core.definitions.metadata import MetadataValue
23
24
  from dagster._core.definitions.reconstruct import ReconstructableJob
24
- from dagster._core.definitions.utils import normalize_tags
25
25
  from dagster._core.execution.context.compute import OpExecutionContext
26
26
  from dagster._core.execution.context.input import build_input_context
27
27
  from dagster._core.execution.context.system import StepExecutionContext
28
28
  from dagster._core.execution.plan.outputs import StepOutputHandle
29
+ from dagster._core.storage.tags import COMPUTE_KIND_TAG
29
30
  from dagster._serdes import pack_value
30
- from dagster._seven import get_system_temp_directory
31
31
  from dagster._utils import mkdir_p, safe_tempfile_path
32
32
  from dagster._utils.error import serializable_error_info_from_exc_info
33
+ from dagster._utils.tags import normalize_tags
34
+ from dagster_shared import seven
35
+ from dagster_shared.seven import get_system_temp_directory
33
36
  from papermill.engines import papermill_engines
34
37
  from papermill.iorw import load_notebook_node, write_ipynb
35
38
 
36
- from .compat import ExecutionError
37
- from .engine import DagstermillEngine
38
- from .errors import DagstermillError
39
- from .translator import DagsterTranslator
39
+ from dagstermill.compat import ExecutionError
40
+ from dagstermill.engine import DagstermillEngine
41
+ from dagstermill.errors import DagstermillError
42
+ from dagstermill.translator import DagsterTranslator
40
43
 
41
44
 
42
45
  def _clean_path_for_windows(notebook_path: str) -> str:
@@ -60,6 +63,7 @@ def _find_first_tagged_cell_index(nb, tag):
60
63
  # This is based on papermill.parameterize.parameterize_notebook
61
64
  # Typically, papermill injects the injected-parameters cell *below* the parameters cell
62
65
  # but we want to *replace* the parameters cell, which is what this function does.
66
+ @beta
63
67
  def replace_parameters(context, nb, parameters):
64
68
  """Assigned parameters into the appropriate place in the input notebook.
65
69
 
@@ -101,11 +105,12 @@ def replace_parameters(context, nb, parameters):
101
105
  after = nb.cells
102
106
 
103
107
  nb.cells = before + [newcell] + after
104
- nb.metadata.papermill["parameters"] = _seven.json.dumps(parameters)
108
+ nb.metadata.papermill["parameters"] = seven.json.dumps(parameters)
105
109
 
106
110
  return nb
107
111
 
108
112
 
113
+ @beta
109
114
  def get_papermill_parameters(
110
115
  step_context: StepExecutionContext,
111
116
  inputs: Mapping[str, object],
@@ -159,6 +164,7 @@ def get_papermill_parameters(
159
164
  return parameters
160
165
 
161
166
 
167
+ @beta
162
168
  def execute_notebook(
163
169
  step_context: StepExecutionContext,
164
170
  name: str,
@@ -340,13 +346,14 @@ def _make_dagstermill_compute_fn(
340
346
  return _t_fn
341
347
 
342
348
 
349
+ @beta
343
350
  def define_dagstermill_op(
344
351
  name: str,
345
352
  notebook_path: str,
346
353
  ins: Optional[Mapping[str, In]] = None,
347
354
  outs: Optional[Mapping[str, Out]] = None,
348
355
  config_schema: Optional[Union[Any, Mapping[str, Any]]] = None,
349
- required_resource_keys: Optional[Set[str]] = None,
356
+ required_resource_keys: Optional[set[str]] = None,
350
357
  output_notebook_name: Optional[str] = None,
351
358
  asset_key_prefix: Optional[Union[Sequence[str], str]] = None,
352
359
  description: Optional[str] = None,
@@ -401,7 +408,7 @@ def define_dagstermill_op(
401
408
  required_resource_keys.add(io_mgr_key)
402
409
  outs = {
403
410
  **outs,
404
- cast(str, output_notebook_name): Out(io_manager_key=io_mgr_key),
411
+ cast("str", output_notebook_name): Out(io_manager_key=io_mgr_key),
405
412
  }
406
413
 
407
414
  if isinstance(asset_key_prefix, str):
@@ -412,7 +419,7 @@ def define_dagstermill_op(
412
419
  default_description = f"This op is backed by the notebook at {notebook_path}"
413
420
  description = check.opt_str_param(description, "description", default=default_description)
414
421
 
415
- user_tags = normalize_tags(tags).tags
422
+ user_tags = normalize_tags(tags)
416
423
  if tags is not None:
417
424
  check.invariant(
418
425
  "notebook_path" not in tags,
@@ -420,14 +427,17 @@ def define_dagstermill_op(
420
427
  " is reserved for use by Dagster",
421
428
  )
422
429
  check.invariant(
423
- "kind" not in tags,
430
+ COMPUTE_KIND_TAG not in tags,
424
431
  "user-defined op tags contains the `kind` key, but the `kind` key is reserved for"
425
432
  " use by Dagster",
426
433
  )
427
- default_tags = {"notebook_path": _clean_path_for_windows(notebook_path), "kind": "ipynb"}
434
+ default_tags = {
435
+ "notebook_path": _clean_path_for_windows(notebook_path),
436
+ COMPUTE_KIND_TAG: "ipynb",
437
+ }
428
438
 
429
439
  if safe_is_subclass(config_schema, Config):
430
- config_schema = infer_schema_from_config_class(cast(Type[Config], config_schema))
440
+ config_schema = infer_schema_from_config_class(cast("type[Config]", config_schema))
431
441
 
432
442
  return OpDefinition(
433
443
  name=name,
@@ -1,6 +1,7 @@
1
1
  import os
2
+ from collections.abc import Sequence
2
3
  from pathlib import Path
3
- from typing import Any, List, Optional, Sequence
4
+ from typing import Any, Optional
4
5
 
5
6
  import dagster._check as check
6
7
  from dagster import (
@@ -10,6 +11,7 @@ from dagster import (
10
11
  InitResourceContext,
11
12
  IOManager,
12
13
  )
14
+ from dagster._annotations import beta
13
15
  from dagster._core.definitions.metadata import MetadataValue
14
16
  from dagster._core.execution.context.input import InputContext
15
17
  from dagster._core.execution.context.output import OutputContext
@@ -33,7 +35,7 @@ class OutputNotebookIOManager(IOManager):
33
35
 
34
36
  class LocalOutputNotebookIOManager(OutputNotebookIOManager):
35
37
  def __init__(self, base_dir: str, asset_key_prefix: Optional[Sequence[str]] = None):
36
- super(LocalOutputNotebookIOManager, self).__init__(asset_key_prefix=asset_key_prefix)
38
+ super().__init__(asset_key_prefix=asset_key_prefix)
37
39
  self.base_dir = base_dir
38
40
  self.write_mode = "wb"
39
41
  self.read_mode = "rb"
@@ -82,6 +84,7 @@ class LocalOutputNotebookIOManager(OutputNotebookIOManager):
82
84
  return file_obj.read()
83
85
 
84
86
 
87
+ @beta
85
88
  class ConfigurableLocalOutputNotebookIOManager(ConfigurableIOManagerFactory):
86
89
  """Built-in IO Manager for handling output notebook."""
87
90
 
@@ -92,7 +95,7 @@ class ConfigurableLocalOutputNotebookIOManager(ConfigurableIOManagerFactory):
92
95
  " directory if not provided."
93
96
  ),
94
97
  )
95
- asset_key_prefix: List[str] = Field(
98
+ asset_key_prefix: list[str] = Field(
96
99
  default=[],
97
100
  description=(
98
101
  "Asset key prefix to apply to assets materialized for output notebooks. Defaults to no"
@@ -111,6 +114,7 @@ class ConfigurableLocalOutputNotebookIOManager(ConfigurableIOManagerFactory):
111
114
  )
112
115
 
113
116
 
117
+ @beta
114
118
  @dagster_maintained_io_manager
115
119
  @io_manager(config_schema=ConfigurableLocalOutputNotebookIOManager.to_config_schema())
116
120
  def local_output_notebook_io_manager(init_context) -> LocalOutputNotebookIOManager:
dagstermill/manager.py CHANGED
@@ -1,7 +1,9 @@
1
1
  import os
2
2
  import pickle
3
3
  import uuid
4
- from typing import TYPE_CHECKING, AbstractSet, Any, Mapping, Optional, cast
4
+ from asyncio import AbstractEventLoop
5
+ from collections.abc import Mapping
6
+ from typing import TYPE_CHECKING, AbstractSet, Any, Optional, cast # noqa: UP035
5
7
 
6
8
  from dagster import (
7
9
  AssetMaterialization,
@@ -14,6 +16,7 @@ from dagster import (
14
16
  TypeCheck,
15
17
  _check as check,
16
18
  )
19
+ from dagster._annotations import beta
17
20
  from dagster._core.definitions.dependency import NodeHandle
18
21
  from dagster._core.definitions.events import RetryRequested
19
22
  from dagster._core.definitions.graph_definition import GraphDefinition
@@ -27,7 +30,6 @@ from dagster._core.execution.api import create_execution_plan, scoped_job_contex
27
30
  from dagster._core.execution.plan.outputs import StepOutputHandle
28
31
  from dagster._core.execution.plan.plan import ExecutionPlan
29
32
  from dagster._core.execution.plan.state import KnownExecutionState
30
- from dagster._core.execution.plan.step import ExecutionStep
31
33
  from dagster._core.execution.resources_init import (
32
34
  get_required_resource_keys_to_init,
33
35
  resource_initialization_event_generator,
@@ -42,12 +44,13 @@ from dagster._loggers import colored_console_logger
42
44
  from dagster._serdes import unpack_value
43
45
  from dagster._utils import EventGenerationManager
44
46
 
45
- from .context import DagstermillExecutionContext, DagstermillRuntimeExecutionContext
46
- from .errors import DagstermillError
47
- from .serialize import PICKLE_PROTOCOL
47
+ from dagstermill.context import DagstermillExecutionContext, DagstermillRuntimeExecutionContext
48
+ from dagstermill.errors import DagstermillError
49
+ from dagstermill.serialize import PICKLE_PROTOCOL
48
50
 
49
51
  if TYPE_CHECKING:
50
52
  from dagster._core.definitions.node_definition import NodeDefinition
53
+ from dagster._core.execution.plan.step import ExecutionStep
51
54
 
52
55
 
53
56
  class DagstermillResourceEventGenerationManager(EventGenerationManager):
@@ -60,14 +63,10 @@ class DagstermillResourceEventGenerationManager(EventGenerationManager):
60
63
  return iter(())
61
64
 
62
65
  def teardown(self):
63
- return [
64
- teardown_event
65
- for teardown_event in super(
66
- DagstermillResourceEventGenerationManager, self
67
- ).generate_teardown_events()
68
- ]
66
+ return [teardown_event for teardown_event in super().generate_teardown_events()]
69
67
 
70
68
 
69
+ @beta
71
70
  class Manager:
72
71
  def __init__(self):
73
72
  self.job = None
@@ -87,6 +86,7 @@ class Manager:
87
86
  resource_keys_to_init: Optional[AbstractSet[str]],
88
87
  instance: Optional[DagsterInstance],
89
88
  emit_persistent_events: Optional[bool],
89
+ event_loop: Optional[AbstractEventLoop],
90
90
  ):
91
91
  """Drop-in replacement for
92
92
  `dagster._core.execution.resources_init.resource_initialization_manager`. It uses a
@@ -101,6 +101,7 @@ class Manager:
101
101
  resource_keys_to_init=resource_keys_to_init,
102
102
  instance=instance,
103
103
  emit_persistent_events=emit_persistent_events,
104
+ event_loop=event_loop,
104
105
  )
105
106
  self.resource_manager = DagstermillResourceEventGenerationManager(
106
107
  generator, ScopedResourcesBuilder
@@ -195,9 +196,9 @@ class Manager:
195
196
  op_name=op.name,
196
197
  node_handle=node_handle,
197
198
  step_context=cast(
198
- StepExecutionContext,
199
+ "StepExecutionContext",
199
200
  job_context.for_step(
200
- cast(ExecutionStep, execution_plan.get_step_by_key(step_key)),
201
+ cast("ExecutionStep", execution_plan.get_step_by_key(step_key)),
201
202
  known_state=known_state,
202
203
  ),
203
204
  ),
@@ -311,18 +312,18 @@ class Manager:
311
312
  # deferred import for perf
312
313
  import scrapbook
313
314
 
314
- if not self.op_def.has_output(output_name):
315
+ if not self.op_def.has_output(output_name): # pyright: ignore[reportOptionalMemberAccess]
315
316
  raise DagstermillError(
316
- f"Op {self.op_def.name} does not have output named {output_name}.Expected one of"
317
- f" {[str(output_def.name) for output_def in self.op_def.output_defs]}"
317
+ f"Op {self.op_def.name} does not have output named {output_name}.Expected one of" # pyright: ignore[reportOptionalMemberAccess]
318
+ f" {[str(output_def.name) for output_def in self.op_def.output_defs]}" # pyright: ignore[reportOptionalMemberAccess]
318
319
  )
319
320
 
320
321
  # pass output value cross process boundary using io manager
321
- step_context = self.context._step_context # noqa: SLF001
322
+ step_context = self.context._step_context # noqa: SLF001 # pyright: ignore[reportAttributeAccessIssue,reportOptionalMemberAccess]
322
323
  # Note: yield_result currently does not support DynamicOutput
323
324
 
324
325
  # dagstermill assets do not support yielding additional results within the notebook:
325
- if len(step_context.job_def.asset_layer.executable_asset_keys) > 0:
326
+ if len(step_context.job_def.asset_layer.executable_asset_keys) > 0: # pyright: ignore[reportArgumentType]
326
327
  raise DagstermillError(
327
328
  "dagstermill assets do not currently support dagstermill.yield_result"
328
329
  )
@@ -369,7 +370,7 @@ class Manager:
369
370
  import scrapbook
370
371
 
371
372
  event_id = f"event-{uuid.uuid4()}"
372
- out_file_path = os.path.join(self.marshal_dir, event_id)
373
+ out_file_path = os.path.join(self.marshal_dir, event_id) # pyright: ignore[reportCallIssue,reportArgumentType]
373
374
  with open(out_file_path, "wb") as fd:
374
375
  fd.write(pickle.dumps(dagster_event, PICKLE_PROTOCOL))
375
376
 
dagstermill/py.typed ADDED
@@ -0,0 +1 @@
1
+ partial
dagstermill/serialize.py CHANGED
@@ -1,11 +1,11 @@
1
- from dagster import _seven
1
+ from dagster_shared import seven
2
2
 
3
3
  PICKLE_PROTOCOL = 2
4
4
 
5
5
 
6
6
  def is_json_serializable(value):
7
7
  try:
8
- _seven.json.dumps(value)
8
+ seven.json.dumps(value)
9
9
  return True
10
10
  except TypeError:
11
11
  return False
@@ -0,0 +1,45 @@
1
+ import os
2
+ from contextlib import contextmanager
3
+
4
+ from dagster import execute_job
5
+ from dagster._core.definitions.metadata import NotebookMetadataValue, PathMetadataValue
6
+ from dagster._core.definitions.reconstruct import ReconstructableJob
7
+ from dagster._core.test_utils import instance_for_test
8
+
9
+
10
+ def get_path(materialization_event):
11
+ for value in materialization_event.event_specific_data.materialization.metadata.values():
12
+ if isinstance(value, (PathMetadataValue, NotebookMetadataValue)):
13
+ return value.value
14
+
15
+
16
+ def cleanup_result_notebook(result):
17
+ if not result:
18
+ return
19
+ materialization_events = [
20
+ x for x in result.all_events if x.event_type_value == "ASSET_MATERIALIZATION"
21
+ ]
22
+ for materialization_event in materialization_events:
23
+ result_path = get_path(materialization_event)
24
+ if os.path.exists(result_path): # pyright: ignore[reportArgumentType]
25
+ os.unlink(result_path) # pyright: ignore[reportArgumentType]
26
+
27
+
28
+ @contextmanager
29
+ def exec_for_test(fn_name, env=None, raise_on_error=True, **kwargs):
30
+ result = None
31
+ recon_job = ReconstructableJob.for_module("dagstermill.examples.repository", fn_name)
32
+
33
+ with instance_for_test() as instance:
34
+ try:
35
+ with execute_job(
36
+ job=recon_job,
37
+ run_config=env,
38
+ instance=instance,
39
+ raise_on_error=raise_on_error,
40
+ **kwargs,
41
+ ) as result:
42
+ yield result
43
+ finally:
44
+ if result:
45
+ cleanup_result_notebook(result)
dagstermill/translator.py CHANGED
@@ -1,6 +1,6 @@
1
1
  import papermill
2
2
  import papermill.translators
3
- from dagster import _seven
3
+ from dagster_shared import seven
4
4
 
5
5
  RESERVED_INPUT_NAMES = [
6
6
  "__dm_context",
@@ -16,7 +16,7 @@ RESERVED_INPUT_NAMES = [
16
16
 
17
17
  INJECTED_BOILERPLATE = """
18
18
  # Injected parameters
19
- from dagster import seven as __dm_seven
19
+ from dagster_shared import seven as __dm_seven
20
20
  import dagstermill as __dm_dagstermill
21
21
  context = __dm_dagstermill._reconstitute_job_context(
22
22
  **{{
@@ -51,7 +51,7 @@ class DagsterTranslator(papermill.translators.PythonTranslator):
51
51
  )
52
52
 
53
53
  for key in job_context_args:
54
- job_context_args[key] = _seven.json.dumps(job_context_args[key])
54
+ job_context_args[key] = seven.json.dumps(job_context_args[key])
55
55
 
56
56
  content = INJECTED_BOILERPLATE.format(job_context_args=job_context_args)
57
57
 
dagstermill/version.py CHANGED
@@ -1 +1 @@
1
- __version__ = "0.23.3"
1
+ __version__ = "0.27.6"
@@ -0,0 +1,35 @@
1
+ Metadata-Version: 2.4
2
+ Name: dagstermill
3
+ Version: 0.27.6
4
+ Summary: run notebooks using the Dagster tools
5
+ Author: Dagster Labs
6
+ Author-email: hello@dagsterlabs.com
7
+ License: Apache-2.0
8
+ Classifier: Programming Language :: Python :: 3.9
9
+ Classifier: Programming Language :: Python :: 3.10
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: License :: OSI Approved :: Apache Software License
12
+ Classifier: Operating System :: OS Independent
13
+ Requires-Python: >=3.9,<3.14
14
+ License-File: LICENSE
15
+ Requires-Dist: dagster==1.11.6
16
+ Requires-Dist: ipykernel!=5.4.0,!=5.4.1,>=4.9.0
17
+ Requires-Dist: ipython_genutils>=0.2.0
18
+ Requires-Dist: packaging>=20.9
19
+ Requires-Dist: papermill>=1.0.0
20
+ Requires-Dist: scrapbook>=0.5.0
21
+ Requires-Dist: nbconvert
22
+ Requires-Dist: jupyter-client<8
23
+ Provides-Extra: test
24
+ Requires-Dist: matplotlib; extra == "test"
25
+ Requires-Dist: scikit-learn>=0.19.0; extra == "test"
26
+ Requires-Dist: tqdm<=4.48; extra == "test"
27
+ Dynamic: author
28
+ Dynamic: author-email
29
+ Dynamic: classifier
30
+ Dynamic: license
31
+ Dynamic: license-file
32
+ Dynamic: provides-extra
33
+ Dynamic: requires-dist
34
+ Dynamic: requires-python
35
+ Dynamic: summary
@@ -0,0 +1,24 @@
1
+ dagstermill/__init__.py,sha256=kTAwcogNwhvjcMeWhYJKRP-uF6MazOZhiSoCiABbm2Y,1263
2
+ dagstermill/__main__.py,sha256=GqkyStw1mh9zw2gMbQBQpIqW0hpr3I8unVWrw_VqdW8,41
3
+ dagstermill/asset_factory.py,sha256=SfnLlMNYxfB3KgIDpHlZ4vGLsDzS40Pe33n6lPwXQvs,9843
4
+ dagstermill/cli.py,sha256=J0vLeYGwcytvmyyL6ar29QDQ0G0Cdcmk_yM5JgnoqNw,4527
5
+ dagstermill/compat.py,sha256=GCfUaGC3eIEhlZP_VFtua8hFKNtwXocOAfYxuKZ5X3I,526
6
+ dagstermill/context.py,sha256=U-Os8kqiD2SBhF7im7bhn8xddmOPFt7-mZYu8fg7C0o,6175
7
+ dagstermill/engine.py,sha256=exxp3VUOKcTWjEOlc0NQZkxeRPmR3JYTFp7BBH5xM7I,5788
8
+ dagstermill/errors.py,sha256=4w6hTVfK_oX64Eusi66CvfdWZNVpCb96VEYMdCvTaD0,141
9
+ dagstermill/factory.py,sha256=AlY-v8fhHjrncRvED3qQZQSYTUpAw0xtN_5iLRT6kjE,18872
10
+ dagstermill/io_managers.py,sha256=zRIMpZ_j61942pTttrXdwwkS8Qv20x0r3V0JWf1xnSg,4440
11
+ dagstermill/manager.py,sha256=_VHYoAZ7NPaYbLaA1-2Ph5_mmb-LSQN-OQRadD6soKg,16235
12
+ dagstermill/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
13
+ dagstermill/serialize.py,sha256=nvNv29bnxrRA5hXYcDJm8HV7Xsj0C4sRB2oUpHXCres,193
14
+ dagstermill/test_utils.py,sha256=_OX3bXMpqCZDRF-Cq7UsBWOMaWuqgDPpIBuddk-ll-8,1585
15
+ dagstermill/translator.py,sha256=6rKVaN3lcx8d5M1Dlir578s-fsATrTNq-ZJDmUDn9wk,2182
16
+ dagstermill/version.py,sha256=BwKhBzWMdVser1JHOUEX0Aa2nBqgua67wsNi17fRle0,23
17
+ dagstermill/examples/__init__.py,sha256=_2u28VKkHmd16vz0Q2N7Fw5QKsJ65HEnIfJL26gNUIw,75
18
+ dagstermill/examples/repository.py,sha256=P_5rkdjsdUZ-fUci_ImZAx9JZMv7iCd1eZGPzQ2LUZc,16021
19
+ dagstermill-0.27.6.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
20
+ dagstermill-0.27.6.dist-info/METADATA,sha256=DYI0SibfUIU1DLhWdNfeYv6vF2fg8-TcD0qnM08yF_w,1109
21
+ dagstermill-0.27.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
22
+ dagstermill-0.27.6.dist-info/entry_points.txt,sha256=885a7vvhABYWEj7W28elkzSmIcKO3REkdd5h4Z4DEJs,53
23
+ dagstermill-0.27.6.dist-info/top_level.txt,sha256=YDelJKdA5YIIrjsObdd8U4E9YhuXJLRe9NKfUzud9Uc,12
24
+ dagstermill-0.27.6.dist-info/RECORD,,
@@ -1,5 +1,5 @@
1
1
  Wheel-Version: 1.0
2
- Generator: bdist_wheel (0.41.2)
2
+ Generator: setuptools (80.9.0)
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any
5
5
 
@@ -186,7 +186,7 @@
186
186
  same "printed page" as the copyright notice for easier
187
187
  identification within third-party archives.
188
188
 
189
- Copyright 2023 Dagster Labs, Inc".
189
+ Copyright 2025 Dagster Labs, Inc.
190
190
 
191
191
  Licensed under the Apache License, Version 2.0 (the "License");
192
192
  you may not use this file except in compliance with the License.
@@ -1,28 +0,0 @@
1
- Metadata-Version: 2.1
2
- Name: dagstermill
3
- Version: 0.23.3
4
- Summary: run notebooks using the Dagster tools
5
- Author: Dagster Labs
6
- Author-email: hello@dagsterlabs.com
7
- License: Apache-2.0
8
- Classifier: Programming Language :: Python :: 3.8
9
- Classifier: Programming Language :: Python :: 3.9
10
- Classifier: Programming Language :: Python :: 3.10
11
- Classifier: Programming Language :: Python :: 3.11
12
- Classifier: License :: OSI Approved :: Apache Software License
13
- Classifier: Operating System :: OS Independent
14
- Requires-Python: >=3.8,<3.13
15
- License-File: LICENSE
16
- Requires-Dist: dagster ==1.7.3
17
- Requires-Dist: ipykernel !=5.4.0,!=5.4.1,>=4.9.0
18
- Requires-Dist: ipython-genutils >=0.2.0
19
- Requires-Dist: packaging >=20.9
20
- Requires-Dist: papermill >=1.0.0
21
- Requires-Dist: scrapbook >=0.5.0
22
- Requires-Dist: nbconvert
23
- Requires-Dist: jupyter-client <8
24
- Provides-Extra: test
25
- Requires-Dist: matplotlib ; extra == 'test'
26
- Requires-Dist: scikit-learn >=0.19.0 ; extra == 'test'
27
- Requires-Dist: tqdm <=4.48 ; extra == 'test'
28
-
@@ -1,22 +0,0 @@
1
- dagstermill/__init__.py,sha256=WZRYSjL1yAM1w5cqmC1t2T7hvuW-RGGY5tV9vZGHZwU,1133
2
- dagstermill/__main__.py,sha256=bYt9eEaoRQWdejEHFD8REx9jxVEdZptECFsV7F49Ink,30
3
- dagstermill/asset_factory.py,sha256=3moIwVr3dzFuDRVAdRZlQlO00nfGFAiA8cvmOgJi1zw,9463
4
- dagstermill/cli.py,sha256=lffeaQgenSR7s8MhNUBNA_BeVehlMUrdg7gnQl1WTdY,4494
5
- dagstermill/compat.py,sha256=GCfUaGC3eIEhlZP_VFtua8hFKNtwXocOAfYxuKZ5X3I,526
6
- dagstermill/context.py,sha256=3_SkRBM9wDB6T64Lc4yJZGp8mpA_PMXkHlF8U1I19A4,6104
7
- dagstermill/engine.py,sha256=4CZW-eni4TnG8C8VhPkZZgj6djItdOUKhHPzvaVhaYo,5843
8
- dagstermill/errors.py,sha256=WOmpAGp-J1XhyGK_LT3ZZKsBwF5dvrWbqSaSldtoh6Y,141
9
- dagstermill/factory.py,sha256=R9IiuAFi35nHneuWsegcXEH1iTp4qUuO0zmtizmBZb0,18638
10
- dagstermill/io_managers.py,sha256=D9UA4KufzrOeCan3TJyPogD0uIJBjK80ExikdMHqweI,4403
11
- dagstermill/manager.py,sha256=HPk1PtVOc7MDFb3F3t-XECvhPYRPRAXNBSapXYAIUBA,15788
12
- dagstermill/serialize.py,sha256=eXW3c26CiILT_uebyFcAKBnsiNxnjyGT_ch3PfGyjek,188
13
- dagstermill/translator.py,sha256=h1VPAOWtdjLKzFjRmyN9hO_R6qJuAETNkJydfdgwWGM,2170
14
- dagstermill/version.py,sha256=HTz96uK4_QHPX5i3NrEcr4eXcxlyepsy_nU-3niBQe8,23
15
- dagstermill/examples/__init__.py,sha256=kzan-9zFjxaJ8o9bqUso44gcGiOmJrlq4JYO-yIBQao,55
16
- dagstermill/examples/repository.py,sha256=hKQ1_GzHFiSf4oZmr1HYQmyRSMG0HDq25m0Y0oz0eXQ,15827
17
- dagstermill-0.23.3.dist-info/LICENSE,sha256=TMatHW4_G9ldRdodEAp-l2Xa2WvsdeOh60E3v1R2jis,11349
18
- dagstermill-0.23.3.dist-info/METADATA,sha256=_vNUVd86D0at1bx1O2eZFPBYN-SCX_UISSsjVqCg2Wg,985
19
- dagstermill-0.23.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
20
- dagstermill-0.23.3.dist-info/entry_points.txt,sha256=885a7vvhABYWEj7W28elkzSmIcKO3REkdd5h4Z4DEJs,53
21
- dagstermill-0.23.3.dist-info/top_level.txt,sha256=YDelJKdA5YIIrjsObdd8U4E9YhuXJLRe9NKfUzud9Uc,12
22
- dagstermill-0.23.3.dist-info/RECORD,,