migrate from click to argparse

common-workflow-lab · Jun 21, 2020 · 4aaed47 · 4aaed47
1 parent df64c68
commit 4aaed47
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 18 deletions.
diff --git a/ipython2cwl/cwltool.py b/ipython2cwl/cwltool.py
@@ -34,7 +34,7 @@ def visit_AnnAssign(self, node):
         type_mapper = {
             CWLFilePathInput.__name__: (
                 'File',
-                'click.Path(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True)',
+                'pathlib.Path',
                 input_flag
             ),
             CWLBooleanInput.__name__: (
@@ -61,14 +61,21 @@ def visit_AnnAssign(self, node):
                     (node, mapper[0], mapper[1], True, mapper[2] == input_flag, mapper[2] == output_flag)
                 )
                 return None
-            elif isinstance(node.annotation, ast.Subscript) \
-                    and node.annotation.value.id == "Optional" \
-                    and node.annotation.slice.value.id in type_mapper:
-                mapper = type_mapper[node.annotation.slice.value.id]
-                self.extracted_nodes.append(
-                    (node, mapper[0] + '?', mapper[1], False, mapper[2] == input_flag, mapper[2] == output_flag)
-                )
-                return None
+            elif isinstance(node.annotation, ast.Subscript):
+                if node.annotation.value.id == "Optional" \
+                        and node.annotation.slice.value.id in type_mapper:
+                    mapper = type_mapper[node.annotation.slice.value.id]
+                    self.extracted_nodes.append(
+                        (node, mapper[0] + '?', mapper[1], False, mapper[2] == input_flag, mapper[2] == output_flag)
+                    )
+                    return None
+                elif node.annotation.value.id == "List" \
+                        and node.annotation.slice.value.id in type_mapper:
+                    mapper = type_mapper[node.annotation.slice.value.id]
+                    self.extracted_nodes.append(
+                        (node, mapper[0] + '[]', mapper[1], True, mapper[2] == input_flag, mapper[2] == output_flag)
+                    )
+                    return None
         except AttributeError:
             pass
         return node
@@ -84,7 +91,7 @@ class AnnotatedIPython2CWLToolConverter:
 
     _VariableNameTypePair = namedtuple(
         'VariableNameTypePair',
-        ['name', 'cwl_typeof', 'click_typeof', 'required', 'is_input', 'is_output']
+        ['name', 'cwl_typeof', 'argparse_typeof', 'required', 'is_input', 'is_output']
     )
 
     """The annotated python code to convert."""
@@ -100,17 +107,23 @@ def __init__(self, annotated_ipython_code: str):
 
     @classmethod
     def _wrap_script_to_method(cls, tree, variables) -> str:
-        main_function = ast.parse(os.linesep.join([
-            'import click',
-            '@click.command()',
-            *[f'@click.option("--{variable.name}", type={variable.click_typeof}, required={variable.required})'
-              for variable in variables],
+        main_template_code = os.linesep.join([
             f"def main({','.join([variable.name for variable in variables])}):",
             "\tpass",
             "if __name__ == '__main__':",
-            "\tmain()"
-        ]))
-        main_function.body[1].body = tree.body
+            *['\t' + line for line in [
+                "import argparse",
+                'import pathlib',
+                "parser = argparse.ArgumentParser()",
+                *[f'parser.add_argument("--{variable.name}", type={variable.argparse_typeof}, required={variable.required})'
+                  for variable in variables],
+                "args = parser.parse_args()",
+                f"main({','.join([f'{v.name}=args.{v.name}' for v in variables])})"
+            ]],
+        ])
+        main_function = ast.parse(main_template_code)
+        [node for node in main_function.body if isinstance(node, ast.FunctionDef) and node.name == 'main'][0]\
+            .body = tree.body
         return astor.to_source(main_function)
 
     def cwl_command_line_tool(self, docker_image_id: str = 'jn2cwl:latest') -> Dict:

diff --git a/tests/test_cwltool.py b/tests/test_cwltool.py
@@ -1,5 +1,7 @@
 from pathlib import Path
 from unittest import TestCase
+
+from iotypes import CWLStringInput
 from ipython2cwl.cwltool import AnnotatedIPython2CWLToolConverter
 import os
 import tempfile
@@ -118,3 +120,52 @@ def test_AnnotatedIPython2CWLToolConverter_optional_arguments(self):
             },
             cwl_tool
         )
+
+    def test_AnnotatedIPython2CWLToolConverter_list_arguments(self):
+        annotated_python_script = os.linesep.join([
+            "import csv",
+            "input_filename: List[CWLFilePathInput] = ['data1.csv', 'data2.csv']",
+            "for fn in input_filename:",
+            "\twith open(input_filename) as f:",
+            "\t\tcsv_reader = csv.reader(f)",
+            "\t\tdata = [line for line in csv_reader]",
+            "\tprint(data)"
+        ])
+        cwl_tool = AnnotatedIPython2CWLToolConverter(annotated_python_script).cwl_command_line_tool()
+        self.assertDictEqual(
+            {
+                'cwlVersion': "v1.1",
+                'class': 'CommandLineTool',
+                'baseCommand': 'notebookTool',
+                'hints': {
+                    'DockerRequirement': {'dockerImageId': 'jn2cwl:latest'}
+                },
+                'inputs': {
+                    'input_filename': {
+                        'type': 'File[]',
+                        'inputBinding': {
+                            'prefix': '--input_filename'
+                        }
+                    }
+                },
+                'outputs': [],
+            },
+            cwl_tool
+        )
+
+    def test_test_AnnotatedIPython2CWLToolConverter_wrap_script_to_method(self):
+        printed_message = ''
+        annotated_python_script = os.linesep.join([
+            'global printed_message',
+            f"msg: {CWLStringInput.__name__} = 'original'",
+            "print('message:', msg)",
+            "printed_message = msg"
+        ])
+        exec(annotated_python_script)
+        self.assertEqual('original', globals()['printed_message'])
+        converter = AnnotatedIPython2CWLToolConverter(annotated_python_script)
+        new_script = converter._wrap_script_to_method(converter._tree, converter._variables)
+        print('\n' + new_script, '\n')
+        exec(new_script)
+        locals()['main']('new message')
+        self.assertEqual('new message', globals()['printed_message'])