Use typing.get_type_hints for input/output and resource/config type inference (#14707)

## Summary & Motivation

Fixes #14571 

Currently, `Parameter.annotation` is used when inferring the dagster
type from an op compute function signature. This causes an error if
either: (a) a string is used directly by the user for the parameter
annotation; (b) `from __future__ import annotations` is used. This is
because both cases cause the value of `Parameter.annotation` to be a
string, and our dagster type resolution logic does not handle strings.

This PR replaces direct access to `Parameter.annotation` with
`typing.get_type_hints`, which handles evaluation of a string annotation
into an actual class which can be handled by our dagster type resolution
logic. Note that this only works if the referents of a string annotation
is module-level scope-- it will fail for a local scope.

Still this is an upgrade and will fix the problem in most cases. It
will:

- Allows users to use string annotations or `from __future__ import
annotations` together with dagster type signature inference
- Allows us to use `from __future__ import annotations` in our own code,
which is necessary to allow Ruff to auto-sort typing-only imports-- see
#14675. Not clear whether this
is a good idea.

## How I Tested These Changes

New unit test using string annotations.

Author: smackesey

python_modules/dagster/dagster/_core/decorator_utils.py

@@ -1,8 +1,26 @@
+import functools
+import re
 import textwrap
 from inspect import Parameter, signature
-from typing import Any, Callable, Optional, Sequence, Set, TypeVar, Union
-
-from typing_extensions import Concatenate, ParamSpec, TypeGuard
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    TypeVar,
+    Union,
+)
+
+from typing_extensions import (
+    Concatenate,
+    ParamSpec,
+    TypeGuard,
+    get_type_hints as typing_get_type_hints,
+)
+
+from dagster._core.errors import DagsterInvalidDefinitionError
 
 R = TypeVar("R")
 T = TypeVar("T")
@@ -36,6 +54,26 @@ def get_function_params(fn: Callable[..., Any]) -> Sequence[Parameter]:
     return list(signature(fn).parameters.values())
 
 
+def get_type_hints(fn: Callable) -> Mapping[str, Any]:
+    target = fn.func if isinstance(fn, functools.partial) else fn
+
+    try:
+        return typing_get_type_hints(target, include_extras=True)
+    except NameError as e:
+        match = re.search(r"'(\w+)'", str(e))
+        assert match
+        annotation = match[1]
+        raise DagsterInvalidDefinitionError(
+            f'Failed to resolve type annotation "{annotation}" in function {target.__name__}. This'
+            " can occur when the parameter has a string annotation that references either: (1) a"
+            " type defined in a local scope (2) a type that is defined or imported in an `if"
+            " TYPE_CHECKING` block. Note that if you are including `from __future__ import"
+            " annotations`, all annotations in that module are stored as strings. Suggested"
+            " solutions include: (1) convert the annotation to a non-string annotation; (2) move"
+            " the type referenced by the annotation out of local scope or a `TYPE_CHECKING` block."
+        )
+
+
 def validate_expected_params(
     params: Sequence[Parameter], expected_params: Sequence[str]
 ) -> Optional[str]:

python_modules/dagster/dagster/_core/definitions/decorators/sensor_decorator.py

@@ -187,7 +187,7 @@ def _wrapped_fn(*args, **kwargs) -> Any:
 
         # Preserve any resource arguments from the underlying function, for when we inspect the
         # wrapped function later on
-        _wrapped_fn.__signature__ = inspect.signature(fn)
+        _wrapped_fn = update_wrapper(_wrapped_fn, wrapped=fn)
 
         return AssetSensorDefinition(
             name=sensor_name,

python_modules/dagster/dagster/_core/definitions/inference.py

@@ -1,6 +1,14 @@
 from inspect import Parameter, Signature, isgeneratorfunction, signature
-from typing import Any, Callable, Mapping, NamedTuple, Optional, Sequence
-
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Sequence,
+)
+
+from dagster._core.decorator_utils import get_type_hints
 from dagster._seven import is_module_available
 
 from .utils import NoValueSentinel
@@ -55,11 +63,12 @@ def _infer_output_description_from_docstring(fn: Callable) -> Optional[str]:
 
 
 def infer_output_props(fn: Callable) -> InferredOutputProps:
-    sig = signature(fn)
-
-    annotation = Parameter.empty
-    if not isgeneratorfunction(fn):
-        annotation = sig.return_annotation
+    type_hints = get_type_hints(fn)
+    annotation = (
+        type_hints["return"]
+        if not isgeneratorfunction(fn) and "return" in type_hints
+        else Parameter.empty
+    )
 
     return InferredOutputProps(
         annotation=annotation,
@@ -74,6 +83,7 @@ def has_explicit_return_type(fn: Callable) -> bool:
 
 def _infer_inputs_from_params(
     params: Sequence[Parameter],
+    type_hints: Mapping[str, object],
     descriptions: Optional[Mapping[str, Optional[str]]] = None,
 ) -> Sequence[InferredInputProps]:
     _descriptions: Mapping[str, Optional[str]] = descriptions or {}
@@ -82,14 +92,14 @@ def _infer_inputs_from_params(
         if param.default is not Parameter.empty:
             input_def = InferredInputProps(
                 param.name,
-                param.annotation,
+                type_hints.get(param.name, param.annotation),
                 default_value=param.default,
                 description=_descriptions.get(param.name),
             )
         else:
             input_def = InferredInputProps(
                 param.name,
-                param.annotation,
+                type_hints.get(param.name, param.annotation),
                 description=_descriptions.get(param.name),
             )
 
@@ -101,7 +111,8 @@ def _infer_inputs_from_params(
 def infer_input_props(fn: Callable, context_arg_provided: bool) -> Sequence[InferredInputProps]:
     sig = signature(fn)
     params = list(sig.parameters.values())
+    type_hints = get_type_hints(fn)
     descriptions = _infer_input_description_from_docstring(fn)
     params_to_infer = params[1:] if context_arg_provided else params
-    defs = _infer_inputs_from_params(params_to_infer, descriptions=descriptions)
+    defs = _infer_inputs_from_params(params_to_infer, type_hints, descriptions=descriptions)
     return defs

python_modules/dagster/dagster/_core/definitions/resource_annotation.py

@@ -1,26 +1,31 @@
 from inspect import Parameter
-from typing import Sequence, TypeVar
+from typing import Any, Optional, Sequence, Type, TypeVar
 
 from typing_extensions import Annotated
 
-from dagster._core.decorator_utils import get_function_params
+from dagster._core.decorator_utils import get_function_params, get_type_hints
 from dagster._core.definitions.resource_definition import ResourceDefinition
 
 
 def get_resource_args(fn) -> Sequence[Parameter]:
-    return [param for param in get_function_params(fn) if _is_resource_annotated(param)]
+    type_annotations = get_type_hints(fn)
+    return [
+        param
+        for param in get_function_params(fn)
+        if _is_resource_annotation(type_annotations.get(param.name))
+    ]
 
 
 RESOURCE_PARAM_METADATA = "resource_param"
 
 
-def _is_resource_annotated(param: Parameter) -> bool:
+def _is_resource_annotation(annotation: Optional[Type[Any]]) -> bool:
     from dagster._config.pythonic_config import ConfigurableResourceFactory
 
     extends_resource_definition = False
     try:
-        extends_resource_definition = isinstance(param.annotation, type) and issubclass(
-            param.annotation, (ResourceDefinition, ConfigurableResourceFactory)
+        extends_resource_definition = isinstance(annotation, type) and issubclass(
+            annotation, (ResourceDefinition, ConfigurableResourceFactory)
         )
     except TypeError:
         # Using builtin Python types in python 3.9+ will raise a TypeError when using issubclass
@@ -29,8 +34,8 @@ def _is_resource_annotated(param: Parameter) -> bool:
         pass
 
     return (extends_resource_definition) or (
-        hasattr(param.annotation, "__metadata__")
-        and getattr(param.annotation, "__metadata__") == (RESOURCE_PARAM_METADATA,)
+        hasattr(annotation, "__metadata__")
+        and getattr(annotation, "__metadata__") == (RESOURCE_PARAM_METADATA,)
     )
 
 

python_modules/dagster/dagster/_core/definitions/source_asset.py

@@ -1,8 +1,5 @@
-from __future__ import annotations
-
 import warnings
 from typing import (
-    TYPE_CHECKING,
     AbstractSet,
     Any,
     Callable,
@@ -54,12 +51,6 @@
 from dagster._utils.backcompat import ExperimentalWarning, experimental_arg_warning
 from dagster._utils.merger import merge_dicts
 
-if TYPE_CHECKING:
-    from dagster._core.execution.context.compute import (
-        OpExecutionContext,
-    )
-
-
 # Going with this catch-all for the time-being to permit pythonic resources
 SourceAssetObserveFunction: TypeAlias = Callable[..., Any]
 
@@ -193,6 +184,9 @@ def _get_op_def_compute_fn(self, observe_fn: SourceAssetObserveFunction):
             DecoratedOpFunction,
             is_context_provided,
         )
+        from dagster._core.execution.context.compute import (
+            OpExecutionContext,
+        )
 
         observe_fn_has_context = is_context_provided(get_function_params(observe_fn))
 

python_modules/dagster/dagster_tests/core_tests/resource_tests/pythonic_resources/test_binding.py

@@ -585,3 +585,28 @@ def my_asset(context, my_resource: MyResourceWithDefault):
     )
 
     assert executed["yes"]
+
+
+class MyModuleLevelResource(ConfigurableResource):
+    str_field: str
+
+
+# Note that an explicit string annotation has the same effect as defining a resource in a module
+# using `from __future__ import annotations`. This test will only work against a module-scoped
+# resource-- this is a hard limitation of string annotations in Python as of 2023-07-06 and Python
+# 3.11.
+def test_bind_with_string_annotation():
+    @asset
+    def my_asset(context, my_resource: "MyModuleLevelResource"):
+        return my_resource.str_field
+
+    str_field_value = "foo"
+
+    defs = Definitions(
+        [my_asset], resources={"my_resource": MyModuleLevelResource(str_field=str_field_value)}
+    )
+
+    assert (
+        defs.get_implicit_global_asset_job_def().execute_in_process().output_for_node("my_asset")
+        == str_field_value
+    )

python_modules/dagster/dagster_tests/general_tests/py3_tests/test_inference.py

@@ -111,6 +111,37 @@ def add_one(num: int) -> int:
     assert add_one.output_defs[0].dagster_type.unique_name == "Int"
 
 
+def test_string_typed_input_and_output():
+    @op
+    def add_one(_context, num: "Optional[int]") -> "int":
+        return num + 1 if num else 1
+
+    assert add_one
+    assert len(add_one.input_defs) == 1
+    assert add_one.input_defs[0].name == "num"
+    assert add_one.input_defs[0].dagster_type.display_name == "Int?"
+
+    assert len(add_one.output_defs) == 1
+    assert add_one.output_defs[0].dagster_type.unique_name == "Int"
+
+
+def _make_foo():
+    class Foo:
+        pass
+
+    def foo(x: "Foo") -> "Foo":
+        return x
+
+    return foo
+
+
+def test_invalid_string_typed_input():
+    with pytest.raises(
+        DagsterInvalidDefinitionError, match='Failed to resolve type annotation "Foo"'
+    ):
+        op(_make_foo())
+
+
 def test_wrapped_input_and_output_lambda():
     @op
     def add_one(nums: List[int]) -> Optional[List[int]]: