try fixing tests

Netflix · Jun 6, 2024 · 63fd05c · 63fd05c
1 parent 0fced00
commit 63fd05c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/metaflow/runner/click_api.py b/metaflow/runner/click_api.py
@@ -183,7 +183,10 @@ def visit_ClassDef(self, node):
         ):
             if node.decorator_list:
                 for each_decorator in node.decorator_list:
-                    decorator_name = each_decorator.func.id
+                    if isinstance(each_decorator, ast.Call):
+                        decorator_name = each_decorator.func.id
+                    elif isinstance(each_decorator, ast.Name):
+                        decorator_name = each_decorator.id
                     if decorator_name in FLOW_DECORATORS_DESC_DICT:
                         class_name = FLOW_DECORATORS_DESC_DICT[decorator_name]
                         options = get_options_from_classname(class_name)