From deb9de6f22caa2ee7cbdf5b8c149ef150c4e64a1 Mon Sep 17 00:00:00 2001
From: Stefan Nelson-Lindall <stef@stripe.com>
Date: Wed, 1 Dec 2021 12:13:41 -0800
Subject: [PATCH 1/6] support for delayed annotations

Signed-off-by: Stefan Nelson-Lindall <stef@stripe.com>
---
 flytekit/core/interface.py                 | 16 +++---
 flytekit/core/launch_plan.py               |  5 +-
 flytekit/core/promise.py                   |  2 +-
 flytekit/core/python_function_task.py      |  7 +--
 flytekit/core/task.py                      |  5 +-
 flytekit/core/workflow.py                  |  7 ++-
 tests/flytekit/unit/core/test_interface.py | 57 ++++++++++++----------
 7 files changed, 51 insertions(+), 48 deletions(-)

diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py
index 4ef55b809c..216e337282 100644
--- a/flytekit/core/interface.py
+++ b/flytekit/core/interface.py
@@ -267,19 +267,24 @@ def _change_unrecognized_type_to_pickle(t: Type[T]) -> Type[T]:
     return t
 
 
-def transform_signature_to_interface(signature: inspect.Signature, docstring: Optional[Docstring] = None) -> Interface:
+def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] = None) -> Interface:
     """
     From the annotations on a task function that the user should have provided, and the output names they want to use
     for each output parameter, construct the TypedInterface object
 
     For now the fancy object, maybe in the future a dumb object.
+
     """
-    outputs = extract_return_annotation(signature.return_annotation)
+    type_hints = typing.get_type_hints(fn)
+    signature = inspect.signature(fn)
+    return_annotation = type_hints.get("return", None)
+
+    outputs = extract_return_annotation(return_annotation)
     for k, v in outputs.items():
         outputs[k] = _change_unrecognized_type_to_pickle(v)
     inputs = OrderedDict()
     for k, v in signature.parameters.items():
-        annotation = v.annotation
+        annotation = type_hints.get(k, None)
         default = v.default if v.default is not inspect.Parameter.empty else None
         # Inputs with default values are currently ignored, we may want to look into that in the future
         inputs[k] = (_change_unrecognized_type_to_pickle(annotation), default)
@@ -287,7 +292,6 @@ def transform_signature_to_interface(signature: inspect.Signature, docstring: Op
     # This is just for typing.NamedTuples - in those cases, the user can select a name to call the NamedTuple. We
     # would like to preserve that name in our custom collections.namedtuple.
     custom_name = None
-    return_annotation = signature.return_annotation
     if hasattr(return_annotation, "__bases__"):
         bases = return_annotation.__bases__
         if len(bases) == 1 and bases[0] == tuple and hasattr(return_annotation, "_fields"):
@@ -334,7 +338,7 @@ def output_name_generator(length: int) -> Generator[str, None, None]:
         yield default_output_name(x)
 
 
-def extract_return_annotation(return_annotation: Union[Type, Tuple]) -> Dict[str, Type]:
+def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Dict[str, Type]:
     """
     The purpose of this function is to sort out whether a function is returning one thing, or multiple things, and to
     name the outputs accordingly, either by using our default name function, or from a typing.NamedTuple.
@@ -368,7 +372,7 @@ def t(a: int, b: str) -> Dict[str, int]: ...
 
     # Handle Option 6
     # We can think about whether we should add a default output name with type None in the future.
-    if return_annotation is None or return_annotation is inspect.Signature.empty:
+    if return_annotation in (None, type(None), inspect.Signature.empty):
         return {}
 
     # This statement results in true for typing.Namedtuple, single and void return types, so this
diff --git a/flytekit/core/launch_plan.py b/flytekit/core/launch_plan.py
index 481871977e..217db39652 100644
--- a/flytekit/core/launch_plan.py
+++ b/flytekit/core/launch_plan.py
@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-import inspect
 from typing import Any, Callable, Dict, List, Optional, Type
 
 from flytekit.core import workflow as _annotated_workflow
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager, FlyteEntities
-from flytekit.core.interface import Interface, transform_inputs_to_parameters, transform_signature_to_interface
+from flytekit.core.interface import Interface, transform_function_to_interface, transform_inputs_to_parameters
 from flytekit.core.promise import create_and_link_node, translate_inputs_to_literals
 from flytekit.core.reference_entity import LaunchPlanReference, ReferenceEntity
 from flytekit.models import common as _common_models
@@ -399,7 +398,7 @@ def reference_launch_plan(
     """
 
     def wrapper(fn) -> ReferenceLaunchPlan:
-        interface = transform_signature_to_interface(inspect.signature(fn))
+        interface = transform_function_to_interface(fn)
         return ReferenceLaunchPlan(project, domain, name, version, interface.inputs, interface.outputs)
 
     return wrapper
diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py
index 01a19cca1e..7c9b780e19 100644
--- a/flytekit/core/promise.py
+++ b/flytekit/core/promise.py
@@ -500,7 +500,7 @@ def create_task_output(
     if len(promises) == 1:
         if not entity_interface:
             return promises[0]
-        # See transform_signature_to_interface for more information, we're using the existence of a name as a proxy
+        # See transform_function_to_interface for more information, we're using the existence of a name as a proxy
         # for the user having specified a one-element typing.NamedTuple, which means we should _not_ extract it. We
         # should still return a tuple but it should be one of ours.
         if not entity_interface.output_tuple_name:
diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py
index addea248f4..25a363b070 100644
--- a/flytekit/core/python_function_task.py
+++ b/flytekit/core/python_function_task.py
@@ -14,7 +14,6 @@
 """
 
 
-import inspect
 from abc import ABC
 from collections import OrderedDict
 from enum import Enum
@@ -24,7 +23,7 @@
 from flytekit.core.base_task import Task, TaskResolverMixin
 from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, FlyteContext, FlyteContextManager
 from flytekit.core.docstring import Docstring
-from flytekit.core.interface import transform_signature_to_interface
+from flytekit.core.interface import transform_function_to_interface
 from flytekit.core.python_auto_container import PythonAutoContainerTask, default_task_resolver
 from flytekit.core.tracker import is_functools_wrapped_module_level, isnested, istestfunction
 from flytekit.core.workflow import (
@@ -114,9 +113,7 @@ def __init__(
         """
         if task_function is None:
             raise ValueError("TaskFunction is a required parameter for PythonFunctionTask")
-        self._native_interface = transform_signature_to_interface(
-            inspect.signature(task_function), Docstring(callable_=task_function)
-        )
+        self._native_interface = transform_function_to_interface(task_function, Docstring(callable_=task_function))
         mutated_interface = self._native_interface.remove_inputs(ignore_input_vars)
         super().__init__(
             task_type=task_type,
diff --git a/flytekit/core/task.py b/flytekit/core/task.py
index 24fe283399..4bc00baf67 100644
--- a/flytekit/core/task.py
+++ b/flytekit/core/task.py
@@ -1,9 +1,8 @@
 import datetime as _datetime
-import inspect
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 from flytekit.core.base_task import TaskMetadata
-from flytekit.core.interface import transform_signature_to_interface
+from flytekit.core.interface import transform_function_to_interface
 from flytekit.core.python_function_task import PythonFunctionTask
 from flytekit.core.reference_entity import ReferenceEntity, TaskReference
 from flytekit.core.resources import Resources
@@ -237,7 +236,7 @@ def reference_task(
     """
 
     def wrapper(fn) -> ReferenceTask:
-        interface = transform_signature_to_interface(inspect.signature(fn))
+        interface = transform_function_to_interface(fn)
         return ReferenceTask(project, domain, name, version, interface.inputs, interface.outputs)
 
     return wrapper
diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py
index a14be3d0a5..744ecfbb11 100644
--- a/flytekit/core/workflow.py
+++ b/flytekit/core/workflow.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import inspect
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
@@ -15,9 +14,9 @@
 from flytekit.core.docstring import Docstring
 from flytekit.core.interface import (
     Interface,
+    transform_function_to_interface,
     transform_inputs_to_parameters,
     transform_interface_to_typed_interface,
-    transform_signature_to_interface,
 )
 from flytekit.core.launch_plan import LaunchPlan
 from flytekit.core.node import Node
@@ -574,7 +573,7 @@ def __init__(
     ):
         name = f"{workflow_function.__module__}.{workflow_function.__name__}"
         self._workflow_function = workflow_function
-        native_interface = transform_signature_to_interface(inspect.signature(workflow_function), docstring=docstring)
+        native_interface = transform_function_to_interface(workflow_function, docstring=docstring)
 
         # TODO do we need this - can this not be in launchplan only?
         #    This can be in launch plan only, but is here only so that we don't have to re-evaluate. Or
@@ -770,7 +769,7 @@ def reference_workflow(
     """
 
     def wrapper(fn) -> ReferenceWorkflow:
-        interface = transform_signature_to_interface(inspect.signature(fn))
+        interface = transform_function_to_interface(fn)
         return ReferenceWorkflow(project, domain, name, version, interface.inputs, interface.outputs)
 
     return wrapper
diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py
index 81a01517de..8e55ee1bb4 100644
--- a/tests/flytekit/unit/core/test_interface.py
+++ b/tests/flytekit/unit/core/test_interface.py
@@ -1,4 +1,3 @@
-import inspect
 import os
 import typing
 from typing import Dict, List
@@ -7,9 +6,9 @@
 from flytekit.core.docstring import Docstring
 from flytekit.core.interface import (
     extract_return_annotation,
+    transform_function_to_interface,
     transform_inputs_to_parameters,
     transform_interface_to_typed_interface,
-    transform_signature_to_interface,
     transform_variable_map,
 )
 from flytekit.models.core import types as _core_types
@@ -21,7 +20,7 @@ def test_extract_only():
     def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int):
         ...
 
-    return_types = extract_return_annotation(inspect.signature(x).return_annotation)
+    return_types = extract_return_annotation(typing.get_type_hints(x).get("return", None))
     assert len(return_types) == 2
     assert return_types["x_str"] == str
     assert return_types["y_int"] == int
@@ -29,7 +28,7 @@ def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int):
     def t() -> List[int]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 1
     assert return_type["o0"]._name == "List"
     assert return_type["o0"].__origin__ == list
@@ -37,7 +36,7 @@ def t() -> List[int]:
     def t() -> Dict[str, int]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 1
     assert return_type["o0"]._name == "Dict"
     assert return_type["o0"].__origin__ == dict
@@ -45,7 +44,7 @@ def t() -> Dict[str, int]:
     def t(a: int, b: str) -> typing.Tuple[int, str]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 2
     assert return_type["o0"] == int
     assert return_type["o1"] == str
@@ -53,7 +52,7 @@ def t(a: int, b: str) -> typing.Tuple[int, str]:
     def t(a: int, b: str) -> (int, str):
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 2
     assert return_type["o0"] == int
     assert return_type["o1"] == str
@@ -61,27 +60,33 @@ def t(a: int, b: str) -> (int, str):
     def t(a: int, b: str) -> str:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 1
     assert return_type["o0"] == str
 
+    def t(a: int, b: str):
+        ...
+
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
+    assert len(return_type) == 0
+
     def t(a: int, b: str) -> None:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 0
 
     def t(a: int, b: str) -> List[int]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 1
     assert return_type["o0"] == List[int]
 
     def t(a: int, b: str) -> Dict[str, int]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t).get("return", None))
     assert len(return_type) == 1
     assert return_type["o0"] == Dict[str, int]
 
@@ -95,11 +100,11 @@ def x(a: int, b: str) -> typing.NamedTuple("NT1", x_str=str, y_int=int):
     def y(a: int, b: str) -> nt1:
         return nt1("hello world", 5)
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(x).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(x).get("return", None)))
     assert result["x_str"].type.simple == 3
     assert result["y_int"].type.simple == 1
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(y).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(y).get("return", None)))
     assert result["x_str"].type.simple == 3
     assert result["y_int"].type.simple == 1
 
@@ -108,7 +113,7 @@ def test_unnamed_typing_tuple():
     def z(a: int, b: str) -> typing.Tuple[int, str]:
         return 5, "hello world"
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(z).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(z).get("return", None)))
     assert result["o0"].type.simple == 1
     assert result["o1"].type.simple == 3
 
@@ -117,7 +122,7 @@ def test_regular_tuple():
     def q(a: int, b: str) -> (int, str):
         return 5, "hello world"
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None)))
     assert result["o0"].type.simple == 1
     assert result["o1"].type.simple == 3
 
@@ -126,7 +131,7 @@ def test_single_output_new_decorator():
     def q(a: int, b: str) -> int:
         return a + len(b)
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None)))
     assert result["o0"].type.simple == 1
 
 
@@ -134,7 +139,7 @@ def test_sig_files():
     def q() -> os.PathLike:
         ...
 
-    result = transform_variable_map(extract_return_annotation(inspect.signature(q).return_annotation))
+    result = transform_variable_map(extract_return_annotation(typing.get_type_hints(q).get("return", None)))
     assert isinstance(result["o0"].type.blob, _core_types.BlobType)
 
 
@@ -142,7 +147,7 @@ def test_file_types():
     def t1() -> FlyteFile[typing.TypeVar("svg")]:
         ...
 
-    return_type = extract_return_annotation(inspect.signature(t1).return_annotation)
+    return_type = extract_return_annotation(typing.get_type_hints(t1).get("return", None))
     assert return_type["o0"].extension() == FlyteFile[typing.TypeVar("svg")].extension()
 
 
@@ -152,7 +157,7 @@ def test_parameters_and_defaults():
     def z(a: int, b: str) -> typing.Tuple[int, str]:
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z))
+    our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
     assert params.parameters["a"].required
     assert params.parameters["a"].default is None
@@ -162,7 +167,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
     def z(a: int, b: str = "hello") -> typing.Tuple[int, str]:
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z))
+    our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
     assert params.parameters["a"].required
     assert params.parameters["a"].default is None
@@ -172,7 +177,7 @@ def z(a: int, b: str = "hello") -> typing.Tuple[int, str]:
     def z(a: int = 7, b: str = "eleven") -> typing.Tuple[int, str]:
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z))
+    our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
     assert not params.parameters["a"].required
     assert params.parameters["a"].default.scalar.primitive.integer == 7
@@ -193,7 +198,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
         """
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z))
+    our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     params = transform_inputs_to_parameters(ctx, our_interface)
     assert params.parameters["a"].var.description == "foo"
     assert params.parameters["b"].var.description == "bar"
@@ -211,7 +216,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
         """
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z))
+    our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
     assert typed_interface.inputs.get("a").description == "foo"
     assert typed_interface.inputs.get("b").description == "bar"
@@ -236,7 +241,7 @@ def z(a: int, b: str) -> typing.Tuple[int, str]:
         """
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z))
+    our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
     assert typed_interface.inputs.get("a").description == "foo"
     assert typed_interface.inputs.get("b").description == "bar"
@@ -264,7 +269,7 @@ def z(a: int, b: str) -> typing.NamedTuple("NT", x_str=str, y_int=int):
         """
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z), Docstring(callable_=z))
+    our_interface = transform_function_to_interface(z, Docstring(callable_=z))
     typed_interface = transform_interface_to_typed_interface(our_interface)
     assert typed_interface.inputs.get("a").description == "foo"
     assert typed_interface.inputs.get("b").description == "bar"
@@ -282,7 +287,7 @@ def __init__(self, name):
     def z(a: Foo) -> Foo:
         ...
 
-    our_interface = transform_signature_to_interface(inspect.signature(z))
+    our_interface = transform_function_to_interface(z)
     params = transform_inputs_to_parameters(ctx, our_interface)
     assert params.parameters["a"].required
     assert params.parameters["a"].default is None

From 8a6a8a4552713f0c7b362123a79f9e499678c2e7 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Tue, 7 Dec 2021 13:27:30 -0800
Subject: [PATCH 2/6] add future annotations

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 tests/flytekit/unit/core/functools/simple_decorator.py | 1 +
 tests/flytekit/unit/core/functools/test_decorators.py  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tests/flytekit/unit/core/functools/simple_decorator.py b/tests/flytekit/unit/core/functools/simple_decorator.py
index cf568c5aa9..a51a283be5 100644
--- a/tests/flytekit/unit/core/functools/simple_decorator.py
+++ b/tests/flytekit/unit/core/functools/simple_decorator.py
@@ -1,4 +1,5 @@
 """Script used for testing local execution of functool.wraps-wrapped tasks"""
+from __future__ import annotations
 
 import os
 from functools import wraps
diff --git a/tests/flytekit/unit/core/functools/test_decorators.py b/tests/flytekit/unit/core/functools/test_decorators.py
index f87e714dc9..4e0bb5818b 100644
--- a/tests/flytekit/unit/core/functools/test_decorators.py
+++ b/tests/flytekit/unit/core/functools/test_decorators.py
@@ -1,4 +1,5 @@
 """Test local execution of files that use functools to decorate tasks and workflows."""
+from __future__ import annotations
 
 import os
 import subprocess

From f321ba6c7203d4d4bd89b0d02c9e924087f5db24 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Tue, 7 Dec 2021 14:37:48 -0800
Subject: [PATCH 3/6] one more file

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 .../unit/core/test_type_hints_postponed.py    | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 tests/flytekit/unit/core/test_type_hints_postponed.py

diff --git a/tests/flytekit/unit/core/test_type_hints_postponed.py b/tests/flytekit/unit/core/test_type_hints_postponed.py
new file mode 100644
index 0000000000..4fd404d044
--- /dev/null
+++ b/tests/flytekit/unit/core/test_type_hints_postponed.py
@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import typing
+from dataclasses import dataclass
+
+from dataclasses_json import dataclass_json
+
+from flytekit.core import context_manager
+from flytekit.core.context_manager import Image, ImageConfig
+from flytekit.core.task import task
+
+serialization_settings = context_manager.SerializationSettings(
+    project="proj",
+    domain="dom",
+    version="123",
+    image_config=ImageConfig(Image(name="name", fqn="asdf/fdsa", tag="123")),
+    env={},
+)
+
+
+@dataclass_json
+@dataclass
+class Foo(object):
+    x: int
+    y: str
+    z: typing.Dict[str, str]
+
+
+@dataclass_json
+@dataclass
+class Bar(object):
+    x: int
+    y: str
+    z: Foo
+
+
+@task
+def t1() -> Foo:
+    return Foo(x=1, y="foo", z={"hello": "world"})
+
+
+def test_guess_dict4():
+    print(t1.interface.outputs["o0"])
+    assert t1.interface.outputs["o0"].type.metadata["definitions"]["FooSchema"] is not None

From 06c6f1bb153329279b6b233a758d2e3db9d85d3a Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Tue, 7 Dec 2021 17:00:14 -0800
Subject: [PATCH 4/6] rename the new test file, raise the exception when schema
 has an error instead of ignoring it - will need to ask around, not sure if
 that's the right behavior as it will break
 test_type_engine.py::test_dataclass_transformer

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/core/type_engine.py                  |  8 +++-
 tests/flytekit/unit/core/test_type_delayed.py | 23 ++++++++++
 .../unit/core/test_type_hints_postponed.py    | 44 -------------------
 3 files changed, 30 insertions(+), 45 deletions(-)
 create mode 100644 tests/flytekit/unit/core/test_type_delayed.py
 delete mode 100644 tests/flytekit/unit/core/test_type_hints_postponed.py

diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
index ff9b257b56..a2a4678f32 100644
--- a/flytekit/core/type_engine.py
+++ b/flytekit/core/type_engine.py
@@ -235,7 +235,13 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
                     v.load_by = LoadDumpOptions.name
             schema = JSONSchema().dump(s)
         except Exception as e:
-            logger.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e)
+            logger.error(
+                f"Failed to extract schema for object {t}, (will run schemaless) error: {e}"
+                f"If you have postponed annotations turned on (PEP 563) turn it off please. Postponed"
+                f"evaluation doesn't work with json dataclasses"
+            )
+            # https://github.com/lovasoa/marshmallow_dataclass/issues/13
+            raise e
 
         return _primitives.Generic.to_flyte_literal_type(metadata=schema)
 
diff --git a/tests/flytekit/unit/core/test_type_delayed.py b/tests/flytekit/unit/core/test_type_delayed.py
new file mode 100644
index 0000000000..0dec20cee3
--- /dev/null
+++ b/tests/flytekit/unit/core/test_type_delayed.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import typing
+from dataclasses import dataclass
+
+from dataclasses_json import dataclass_json
+
+from flytekit.core.type_engine import TypeEngine
+
+
+@dataclass_json
+@dataclass
+class Foo(object):
+    x: int
+    y: str
+    z: typing.Dict[str, str]
+
+
+def test_jsondc_schemaize():
+    try:
+        TypeEngine.to_literal_type(Foo)
+    except Exception as e:
+        assert "unsupported field type" in f"{e}"
diff --git a/tests/flytekit/unit/core/test_type_hints_postponed.py b/tests/flytekit/unit/core/test_type_hints_postponed.py
deleted file mode 100644
index 4fd404d044..0000000000
--- a/tests/flytekit/unit/core/test_type_hints_postponed.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from __future__ import annotations
-
-import typing
-from dataclasses import dataclass
-
-from dataclasses_json import dataclass_json
-
-from flytekit.core import context_manager
-from flytekit.core.context_manager import Image, ImageConfig
-from flytekit.core.task import task
-
-serialization_settings = context_manager.SerializationSettings(
-    project="proj",
-    domain="dom",
-    version="123",
-    image_config=ImageConfig(Image(name="name", fqn="asdf/fdsa", tag="123")),
-    env={},
-)
-
-
-@dataclass_json
-@dataclass
-class Foo(object):
-    x: int
-    y: str
-    z: typing.Dict[str, str]
-
-
-@dataclass_json
-@dataclass
-class Bar(object):
-    x: int
-    y: str
-    z: Foo
-
-
-@task
-def t1() -> Foo:
-    return Foo(x=1, y="foo", z={"hello": "world"})
-
-
-def test_guess_dict4():
-    print(t1.interface.outputs["o0"])
-    assert t1.interface.outputs["o0"].type.metadata["definitions"]["FooSchema"] is not None

From 1b8cc8616373dd75ad692e858074577642c607d9 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Wed, 8 Dec 2021 13:42:56 -0800
Subject: [PATCH 5/6] remove test for now

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 tests/flytekit/unit/core/test_type_engine.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
index 826c8df1fb..78500b7db5 100644
--- a/tests/flytekit/unit/core/test_type_engine.py
+++ b/tests/flytekit/unit/core/test_type_engine.py
@@ -421,13 +421,6 @@ def __init__(self):
         self._a = "Hello"
 
 
-@dataclass_json
-@dataclass
-class UnsupportedNestedStruct(object):
-    a: int
-    s: UnsupportedSchemaType
-
-
 def test_dataclass_transformer():
     schema = {
         "$ref": "#/definitions/TeststructSchema",
@@ -472,12 +465,6 @@ def test_dataclass_transformer():
     assert t.metadata is not None
     assert t.metadata == schema
 
-    t = tf.get_literal_type(UnsupportedNestedStruct)
-    assert t is not None
-    assert t.simple is not None
-    assert t.simple == SimpleType.STRUCT
-    assert t.metadata is None
-
 
 def test_dataclass_int_preserving():
     ctx = FlyteContext.current_context()

From 34bf00dc73be1a1ca05107b0c314b7dd1d46ff9e Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Wed, 8 Dec 2021 14:53:15 -0800
Subject: [PATCH 6/6] revert test, eat the error again, change the postponed
 evaluation test for dc json to check for the broken behavior :(

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/core/type_engine.py                  |  5 ++---
 tests/flytekit/unit/core/test_type_delayed.py | 12 ++++++++----
 tests/flytekit/unit/core/test_type_engine.py  | 13 +++++++++++++
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py
index 108df94466..d11e795d20 100644
--- a/flytekit/core/type_engine.py
+++ b/flytekit/core/type_engine.py
@@ -244,13 +244,12 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
                     v.load_by = LoadDumpOptions.name
             schema = JSONSchema().dump(s)
         except Exception as e:
-            logger.error(
+            # https://github.com/lovasoa/marshmallow_dataclass/issues/13
+            logger.warning(
                 f"Failed to extract schema for object {t}, (will run schemaless) error: {e}"
                 f"If you have postponed annotations turned on (PEP 563) turn it off please. Postponed"
                 f"evaluation doesn't work with json dataclasses"
             )
-            # https://github.com/lovasoa/marshmallow_dataclass/issues/13
-            raise e
 
         return _primitives.Generic.to_flyte_literal_type(metadata=schema)
 
diff --git a/tests/flytekit/unit/core/test_type_delayed.py b/tests/flytekit/unit/core/test_type_delayed.py
index 0dec20cee3..87daa91f47 100644
--- a/tests/flytekit/unit/core/test_type_delayed.py
+++ b/tests/flytekit/unit/core/test_type_delayed.py
@@ -17,7 +17,11 @@ class Foo(object):
 
 
 def test_jsondc_schemaize():
-    try:
-        TypeEngine.to_literal_type(Foo)
-    except Exception as e:
-        assert "unsupported field type" in f"{e}"
+    lt = TypeEngine.to_literal_type(Foo)
+    pt = TypeEngine.guess_python_type(lt)
+
+    # When postponed annotations are enabled, dataclass_json will not work and we'll end up with a
+    # schemaless generic.
+    # This test basically tests the broken behavior. Remove this test if
+    # https://github.com/lovasoa/marshmallow_dataclass/issues/13 is ever fixed.
+    assert pt is dict
diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py
index dd02b11784..1cbc58d94d 100644
--- a/tests/flytekit/unit/core/test_type_engine.py
+++ b/tests/flytekit/unit/core/test_type_engine.py
@@ -432,6 +432,13 @@ def __init__(self):
         self._a = "Hello"
 
 
+@dataclass_json
+@dataclass
+class UnsupportedNestedStruct(object):
+    a: int
+    s: UnsupportedSchemaType
+
+
 def test_dataclass_transformer():
     schema = {
         "$ref": "#/definitions/TeststructSchema",
@@ -475,6 +482,12 @@ def test_dataclass_transformer():
     assert t.metadata is not None
     assert t.metadata == schema
 
+    t = tf.get_literal_type(UnsupportedNestedStruct)
+    assert t is not None
+    assert t.simple is not None
+    assert t.simple == SimpleType.STRUCT
+    assert t.metadata is None
+
 
 def test_dataclass_int_preserving():
     ctx = FlyteContext.current_context()