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)