fix __annotations__ key not found (#2391)

Signed-off-by: Austin Liu <austin362667@gmail.com>
flyteorg · May 7, 2024 · 4dd4d22 · 4dd4d22
1 parent 5290171
commit 4dd4d22
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
@@ -161,8 +161,7 @@ def extract_cols_and_format(
             if hasattr(aa, "__annotations__"):
                 # handle dataclass argument
                 d = collections.OrderedDict()
-                dm = vars(aa)
-                d.update(dm["__annotations__"])
+                d.update(aa.__annotations__)
                 ordered_dict_cols = d
             elif isinstance(aa, dict):
                 d = collections.OrderedDict()