Make transform_signature_to_interface compatible with 3.9 NamedTuple

Signed-off-by: Fernando Diaz <[email protected]>
flyteorg · Jul 1, 2021 · 3efb42c · 3efb42c
1 parent 2c9fab4
commit 3efb42c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py
@@ -243,9 +243,11 @@ 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
+    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
@@ -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)