Improved type annotation assignment logic, fixed test

groundlight · Sep 5, 2023 · 1ee0eae · 1ee0eae
1 parent db3c8b6
commit 1ee0eae
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/groundlight/cli.py b/src/groundlight/cli.py
@@ -22,8 +22,9 @@ def wrapper(*args, **kwargs):
     for name, annotation in method.__annotations__.items():
         if get_origin(annotation) is Union:
             if str in annotation.__args__:
-                new_annotation = str
-        wrapper.__annotations__[name] = new_annotation
+                wrapper.__annotations__[name] = str
+            else:
+                wrapper.__annotations__[name] = annotation
 
     return wrapper