From 5575433fdf323e92dfb4cef36930bbb8b051d950 Mon Sep 17 00:00:00 2001 From: Fernando Diaz Date: Thu, 1 Jul 2021 21:53:06 +0200 Subject: [PATCH] Make transform_signature_to_interface compatible with 3.9 NamedTuple (#534) Signed-off-by: Fernando Diaz --- flytekit/core/interface.py | 9 ++++++--- flytekit/core/type_engine.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index bf7fc7c8f1..992fae3a0d 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -243,9 +243,12 @@ def transform_signature_to_interface(signature: inspect.Signature) -> Interface: # 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 - if hasattr(signature.return_annotation, "_field_types"): - if hasattr(signature.return_annotation, "__name__") and signature.return_annotation.__name__ != "": - custom_name = signature.return_annotation.__name__ + 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"): + if hasattr(return_annotation, "__name__") and return_annotation.__name__ != "": + custom_name = return_annotation.__name__ return Interface(inputs, outputs, output_tuple_name=custom_name) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 482cb96bc6..fc4481c484 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -368,7 +368,7 @@ def to_python_value(cls, ctx: FlyteContext, lv: Literal, expected_python_type: T @classmethod def named_tuple_to_variable_map(cls, t: typing.NamedTuple) -> _interface_models.VariableMap: variables = {} - for idx, (var_name, var_type) in enumerate(t._field_types.items()): + for idx, (var_name, var_type) in enumerate(t.__annotations__.items()): literal_type = cls.to_literal_type(var_type) variables[var_name] = _interface_models.Variable(type=literal_type, description=f"{idx}") return _interface_models.VariableMap(variables=variables)