fix typecheck for In.dagster_type (#7242)

dagster-io · Mar 31, 2022 · 99d40dd · 99d40dd
1 parent 14204d5
commit 99d40dd
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 8 deletions.
diff --git a/python_modules/dagster/dagster/core/asset_defs/decorators.py b/python_modules/dagster/dagster/core/asset_defs/decorators.py
@@ -387,11 +387,9 @@ def build_asset_ins(
             asset_key = asset_ins[input_name].asset_key
             metadata = asset_ins[input_name].metadata or {}
             namespace = asset_ins[input_name].namespace
-            dagster_type = None
         else:
             metadata = {}
             namespace = None
-            dagster_type = None
 
         asset_key = asset_key or AssetKey(
             list(filter(None, [*(namespace or asset_namespace or []), input_name]))
@@ -401,7 +399,6 @@ def build_asset_ins(
             metadata=metadata,
             root_manager_key="root_manager",
             asset_key=asset_key,
-            dagster_type=dagster_type,
         )
 
     for asset_key in non_argument_deps:

diff --git a/python_modules/dagster/dagster/core/definitions/input.py b/python_modules/dagster/dagster/core/definitions/input.py
@@ -350,7 +350,7 @@ class In(
     NamedTuple(
         "_In",
         [
-            ("dagster_type", Optional[Union[type, DagsterType]]),
+            ("dagster_type", Union[DagsterType, Type[NoValueSentinel]]),
             ("description", Optional[str]),
             ("default_value", Any),
             ("root_manager_key", Optional[str]),
@@ -386,7 +386,7 @@ class In(
 
     def __new__(
         cls,
-        dagster_type: Optional[Union[Type, DagsterType]] = NoValueSentinel,
+        dagster_type: Union[Type, DagsterType] = NoValueSentinel,
         description: Optional[str] = None,
         default_value: Any = NoValueSentinel,
         root_manager_key: Optional[str] = None,
@@ -396,7 +396,9 @@ def __new__(
     ):
         return super(In, cls).__new__(
             cls,
-            dagster_type=check.opt_inst_param(dagster_type, "dagster_type", (type, DagsterType)),
+            dagster_type=NoValueSentinel
+            if dagster_type is NoValueSentinel
+            else resolve_dagster_type(dagster_type),
             description=check.opt_str_param(description, "description"),
             default_value=default_value,
             root_manager_key=check.opt_str_param(root_manager_key, "root_manager_key"),

diff --git a/python_modules/dagster/dagster/core/definitions/output.py b/python_modules/dagster/dagster/core/definitions/output.py
@@ -388,7 +388,7 @@ class Out(
 
     def __new__(
         cls,
-        dagster_type: Union[Type, DagsterType, Type[NoValueSentinel]] = NoValueSentinel,
+        dagster_type: Union[Type, DagsterType] = NoValueSentinel,
         description: Optional[str] = None,
         is_required: bool = True,
         io_manager_key: Optional[str] = None,

diff --git a/...on_modules/dagster/dagster_tests/core_tests/definitions_tests/decorators_tests/test_op.py b/...on_modules/dagster/dagster_tests/core_tests/definitions_tests/decorators_tests/test_op.py
@@ -1,5 +1,5 @@
 import time
-from typing import Dict, Generator, Tuple
+from typing import Dict, Generator, List, Tuple
 
 import pytest
 
@@ -89,6 +89,14 @@ def my_graph():
     assert my_op(1, "2") == 3
 
 
+def test_ins_dagster_types():
+    assert In(dagster_type=None)
+    assert In(dagster_type=int)
+    assert In(dagster_type=List)
+    assert In(dagster_type=List[int])  # typing type
+    assert In(dagster_type=Int)  # dagster type
+
+
 def test_out():
     @op(out=Out(metadata={"x": 1}))
     def my_op() -> int:
@@ -104,6 +112,14 @@ def my_op() -> int:
     assert my_op() == 1
 
 
+def test_out_dagster_types():
+    assert Out(dagster_type=None)
+    assert Out(dagster_type=int)
+    assert Out(dagster_type=List)
+    assert Out(dagster_type=List[int])  # typing type
+    assert Out(dagster_type=Int)  # dagster type
+
+
 def test_multi_out():
     @op(out={"a": Out(metadata={"x": 1}), "b": Out(metadata={"y": 2})})
     def my_op() -> Tuple[int, str]: