diff --git a/python_modules/dagster/dagster/core/execution/context/compute.py b/python_modules/dagster/dagster/core/execution/context/compute.py --- a/python_modules/dagster/dagster/core/execution/context/compute.py +++ b/python_modules/dagster/dagster/core/execution/context/compute.py @@ -1,6 +1,8 @@ from abc import ABC, abstractmethod, abstractproperty +from typing import Any from dagster import check +from dagster.core.storage.pipeline_run import PipelineRun from dagster.utils.forked_pdb import ForkedPdb from .step import StepExecutionContext @@ -8,7 +10,7 @@ class AbstractComputeExecutionContext(ABC): # pylint: disable=no-init - """Base class for solid context implemented by SolidExecutionContext and DagstermillInNotebookExecutionContext""" + """Base class for solid context implemented by SolidExecutionContext and DagstermillExecutionContext""" @abstractmethod def has_tag(self, key): @@ -35,13 +37,21 @@ """The pipeline being executed.""" @abstractproperty - def resources(self): + def pipeline_run(self) -> PipelineRun: + """The PipelineRun object corresponding to the execution.""" + + @abstractproperty + def resources(self) -> Any: """Resources available in the execution context.""" @abstractproperty def log(self): """The log manager available in the execution context.""" + @abstractproperty + def solid_config(self): + """The parsed config specific to this solid.""" + class SolidExecutionContext(StepExecutionContext, AbstractComputeExecutionContext): """The ``context`` object available to solid compute logic.""" diff --git a/python_modules/dagster/dagster/core/execution/context/step.py b/python_modules/dagster/dagster/core/execution/context/step.py --- a/python_modules/dagster/dagster/core/execution/context/step.py +++ b/python_modules/dagster/dagster/core/execution/context/step.py @@ -69,7 +69,7 @@ def get_tag(self, key): return self._system_step_execution_context.get_tag(key) - def get_system_context(self): + def get_system_context(self) -> SystemStepExecutionContext: """ This allows advanced users (e.g. framework authors) to punch through to the underlying system context. diff --git a/python_modules/dagster/dagster/core/execution/context/system.py b/python_modules/dagster/dagster/core/execution/context/system.py --- a/python_modules/dagster/dagster/core/execution/context/system.py +++ b/python_modules/dagster/dagster/core/execution/context/system.py @@ -128,7 +128,7 @@ return self._execution_context_data.environment_config @property - def pipeline(self): + def pipeline(self) -> IPipeline: return self._execution_context_data.pipeline @property diff --git a/python_modules/dagster/dagster/utils/__init__.py b/python_modules/dagster/dagster/utils/__init__.py --- a/python_modules/dagster/dagster/utils/__init__.py +++ b/python_modules/dagster/dagster/utils/__init__.py @@ -13,6 +13,7 @@ import threading from collections import namedtuple from enum import Enum +from typing import Iterator from warnings import warn import _thread as thread @@ -253,7 +254,7 @@ raise cpe -def safe_tempfile_path_unmanaged(): +def safe_tempfile_path_unmanaged() -> str: # This gets a valid temporary file path in the safest possible way, although there is still no # guarantee that another process will not create a file at this path. The NamedTemporaryFile is # deleted when the context manager exits and the file object is closed. @@ -269,7 +270,7 @@ @contextlib.contextmanager -def safe_tempfile_path(): +def safe_tempfile_path() -> Iterator[str]: try: path = safe_tempfile_path_unmanaged() yield path