improve typing parsing on list & optionals

Author: giannisdoukas

ipython2cwl/cwltoolextractor.py

@@ -26,23 +26,31 @@
 
 class AnnotatedVariablesExtractor(ast.NodeTransformer):
     input_type_mapper = {
-        CWLFilePathInput.__name__: (
+        (CWLFilePathInput.__name__,): (
             'File',
             'pathlib.Path',
         ),
-        CWLBooleanInput.__name__: (
+        (CWLBooleanInput.__name__,): (
             'boolean',
             'lambda flag: flag.upper() == "TRUE"',
         ),
-        CWLIntInput.__name__: (
+        (CWLIntInput.__name__,): (
             'int',
             'int',
         ),
-        CWLStringInput.__name__: (
+        (CWLStringInput.__name__,): (
             'string',
             'str',
         ),
     }
+    input_type_mapper = {**input_type_mapper, **{
+        ('List', *(t for t in types_names)): (types[0] + "[]", types[1])
+        for types_names, types in input_type_mapper.items()
+    }, **{
+        ('Optional', *(t for t in types_names)): (types[0] + "?", types[1])
+        for types_names, types in input_type_mapper.items()
+    }}
+
     output_type_mapper = {
         CWLFilePathOutput.__name__
     }
@@ -51,33 +59,29 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.extracted_nodes = []
 
+    def __get_annotation__(self, type_annotation):
+        annotation = None
+        if isinstance(type_annotation, ast.Name):
+            annotation = (type_annotation.id,)
+        elif isinstance(type_annotation, ast.Str):
+            annotation = (type_annotation.s,)
+            ann_expr = ast.parse(type_annotation.s.strip()).body[0]
+            if hasattr(ann_expr, 'value') and isinstance(ann_expr.value, ast.Subscript):
+                annotation = self.__get_annotation__(ann_expr.value)
+        elif isinstance(type_annotation, ast.Subscript):
+            annotation = (type_annotation.value.id, *self.__get_annotation__(type_annotation.slice.value))
+        return annotation
+
     def visit_AnnAssign(self, node):
         try:
-            if (isinstance(node.annotation, ast.Name) and node.annotation.id in self.input_type_mapper) or \
-                    (isinstance(node.annotation, ast.Str) and node.annotation.s in self.input_type_mapper):
-                if hasattr(node.annotation, 'id'):
-                    mapper = self.input_type_mapper[node.annotation.id]
-                else:
-                    mapper = self.input_type_mapper[node.annotation.s]
+            annotation = self.__get_annotation__(node.annotation)
+            if annotation in self.input_type_mapper:
+                mapper = self.input_type_mapper[annotation]
                 self.extracted_nodes.append(
                     (node, mapper[0], mapper[1], True, True, False)
                 )
                 return None
-            elif isinstance(node.annotation, ast.Subscript):
-                if node.annotation.value.id == "Optional" \
-                        and node.annotation.slice.value.id in self.input_type_mapper:
-                    mapper = self.input_type_mapper[node.annotation.slice.value.id]
-                    self.extracted_nodes.append(
-                        (node, mapper[0] + '?', mapper[1], False, True, False)
-                    )
-                    return None
-                elif node.annotation.value.id == "List" \
-                        and node.annotation.slice.value.id in self.input_type_mapper:
-                    mapper = self.input_type_mapper[node.annotation.slice.value.id]
-                    self.extracted_nodes.append(
-                        (node, mapper[0] + '[]', mapper[1], True, True, False)
-                    )
-                    return None
+
             elif (isinstance(node.annotation, ast.Name) and node.annotation.id in self.output_type_mapper) or \
                     (isinstance(node.annotation, ast.Str) and node.annotation.s in self.output_type_mapper):
                 self.extracted_nodes.append(
@@ -90,7 +94,7 @@ def visit_AnnAssign(self, node):
                     targets=[node.target],
                     value=node.value
                 )
-        except AttributeError:
+        except Exception:
             pass
         return node
 

tests/test_cwltoolextractor.py

@@ -328,3 +328,52 @@ def test_AnnotatedIPython2CWLToolConverter_exclamation_mark_command(self):
         exec(new_script_without_magics)
         locals()['main']('original\n!ls -l')
         self.assertEqual('original\n!ls -l', globals()['printed_message'])
+
+    def test_AnnotatedIPython2CWLToolConverter_optional_array_input(self):
+        s1 = os.linesep.join([
+            'x1: CWLBooleanInput = True',
+        ])
+        s2 = os.linesep.join([
+            'x1: "CWLBooleanInput" = True',
+        ])
+        # all variables must be the same
+        self.assertEqual(
+            AnnotatedIPython2CWLToolConverter(s1)._variables[0],
+            AnnotatedIPython2CWLToolConverter(s2)._variables[0],
+        )
+
+        s1 = os.linesep.join([
+            'x1: Optional[CWLBooleanInput] = True',
+        ])
+        s2 = os.linesep.join([
+            'x1: "Optional[CWLBooleanInput]" = True',
+        ])
+        s3 = os.linesep.join([
+            'x1: Optional["CWLBooleanInput"] = True',
+        ])
+        # all variables must be the same
+        self.assertEqual(
+            AnnotatedIPython2CWLToolConverter(s1)._variables[0],
+            AnnotatedIPython2CWLToolConverter(s2)._variables[0],
+        )
+        self.assertEqual(
+            AnnotatedIPython2CWLToolConverter(s1)._variables[0],
+            AnnotatedIPython2CWLToolConverter(s3)._variables[0],
+        )
+
+        # test that does not crash
+        self.assertListEqual([], AnnotatedIPython2CWLToolConverter(os.linesep.join([
+            'x1: RandomHint = True'
+        ]))._variables)
+        self.assertListEqual([], AnnotatedIPython2CWLToolConverter(os.linesep.join([
+            'x1: List[RandomHint] = True'
+        ]))._variables)
+        self.assertListEqual([], AnnotatedIPython2CWLToolConverter(os.linesep.join([
+            'x1: List["RandomHint"] = True'
+        ]))._variables)
+        self.assertListEqual([], AnnotatedIPython2CWLToolConverter(os.linesep.join([
+            'x1: "List[List[Union[RandomHint, Foo]]]" = True'
+        ]))._variables)
+        self.assertListEqual([], AnnotatedIPython2CWLToolConverter(os.linesep.join([
+            'x1: "RANDOM CHARACTERS!!!!!!" = True'
+        ]))._variables)