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 62f982a11b..45b7e2dc94 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, TaskResolverMixin -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 @@ -240,7 +239,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/type_engine.py b/flytekit/core/type_engine.py index c2ac3bcb01..d11e795d20 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -244,7 +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.warn("failed to extract schema for object %s, (will run schemaless) error: %s", str(t), e) + # 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" + ) return _primitives.Generic.to_flyte_literal_type(metadata=schema) 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/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 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 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..87daa91f47 --- /dev/null +++ b/tests/flytekit/unit/core/test_type_delayed.py @@ -0,0 +1,27 @@ +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(): + 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