2222 SETUP_TEMPLATE = f .read ()
2323
2424
25- # TODO: does not support recursion if main function exists
25+ # TODO: check if supports recursion if main function exists
2626
2727class AnnotatedVariablesExtractor (ast .NodeTransformer ):
2828 input_type_mapper = {
29- CWLFilePathInput .__name__ : (
29+ ( CWLFilePathInput .__name__ ,) : (
3030 'File' ,
3131 'pathlib.Path' ,
3232 ),
33- CWLBooleanInput .__name__ : (
33+ ( CWLBooleanInput .__name__ ,) : (
3434 'boolean' ,
3535 'lambda flag: flag.upper() == "TRUE"' ,
3636 ),
37- CWLIntInput .__name__ : (
37+ ( CWLIntInput .__name__ ,) : (
3838 'int' ,
3939 'int' ,
4040 ),
41- CWLStringInput .__name__ : (
41+ ( CWLStringInput .__name__ ,) : (
4242 'string' ,
4343 'str' ,
4444 ),
4545 }
46+ input_type_mapper = {** input_type_mapper , ** {
47+ ('List' , * (t for t in types_names )): (types [0 ] + "[]" , types [1 ])
48+ for types_names , types in input_type_mapper .items ()
49+ }, ** {
50+ ('Optional' , * (t for t in types_names )): (types [0 ] + "?" , types [1 ])
51+ for types_names , types in input_type_mapper .items ()
52+ }}
53+
4654 output_type_mapper = {
4755 CWLFilePathOutput .__name__
4856 }
@@ -51,30 +59,29 @@ def __init__(self, *args, **kwargs):
5159 super ().__init__ (* args , ** kwargs )
5260 self .extracted_nodes = []
5361
62+ def __get_annotation__ (self , type_annotation ):
63+ annotation = None
64+ if isinstance (type_annotation , ast .Name ):
65+ annotation = (type_annotation .id ,)
66+ elif isinstance (type_annotation , ast .Str ):
67+ annotation = (type_annotation .s ,)
68+ ann_expr = ast .parse (type_annotation .s .strip ()).body [0 ]
69+ if hasattr (ann_expr , 'value' ) and isinstance (ann_expr .value , ast .Subscript ):
70+ annotation = self .__get_annotation__ (ann_expr .value )
71+ elif isinstance (type_annotation , ast .Subscript ):
72+ annotation = (type_annotation .value .id , * self .__get_annotation__ (type_annotation .slice .value ))
73+ return annotation
74+
5475 def visit_AnnAssign (self , node ):
5576 try :
56- if ( isinstance ( node . annotation , ast . Name ) and node .annotation . id in self . input_type_mapper ) or \
57- ( isinstance ( node . annotation , ast . Str ) and node . annotation . s in self .input_type_mapper ) :
58- mapper = self .input_type_mapper [node . annotation . id ]
77+ annotation = self . __get_annotation__ ( node .annotation )
78+ if annotation in self .input_type_mapper :
79+ mapper = self .input_type_mapper [annotation ]
5980 self .extracted_nodes .append (
60- (node , mapper [0 ], mapper [1 ], True , True , False )
81+ (node , mapper [0 ], mapper [1 ], not mapper [ 0 ]. endswith ( '?' ) , True , False )
6182 )
6283 return None
63- elif isinstance (node .annotation , ast .Subscript ):
64- if node .annotation .value .id == "Optional" \
65- and node .annotation .slice .value .id in self .input_type_mapper :
66- mapper = self .input_type_mapper [node .annotation .slice .value .id ]
67- self .extracted_nodes .append (
68- (node , mapper [0 ] + '?' , mapper [1 ], False , True , False )
69- )
70- return None
71- elif node .annotation .value .id == "List" \
72- and node .annotation .slice .value .id in self .input_type_mapper :
73- mapper = self .input_type_mapper [node .annotation .slice .value .id ]
74- self .extracted_nodes .append (
75- (node , mapper [0 ] + '[]' , mapper [1 ], True , True , False )
76- )
77- return None
84+
7885 elif (isinstance (node .annotation , ast .Name ) and node .annotation .id in self .output_type_mapper ) or \
7986 (isinstance (node .annotation , ast .Str ) and node .annotation .s in self .output_type_mapper ):
8087 self .extracted_nodes .append (
@@ -87,7 +94,7 @@ def visit_AnnAssign(self, node):
8794 targets = [node .target ],
8895 value = node .value
8996 )
90- except AttributeError :
97+ except Exception :
9198 pass
9299 return node
93100
@@ -152,6 +159,7 @@ def from_jupyter_notebook_node(cls, node: NotebookNode) -> 'AnnotatedIPython2CWL
152159
153160 @classmethod
154161 def _wrap_script_to_method (cls , tree , variables ) -> str :
162+ add_args = cls .__get_add_arguments__ ([v for v in variables if v .is_input ])
155163 main_template_code = os .linesep .join ([
156164 f"def main({ ',' .join ([v .name for v in variables if v .is_input ])} ,
157165 "\t pass" ,
@@ -160,19 +168,34 @@ def _wrap_script_to_method(cls, tree, variables) -> str:
160168 "import argparse" ,
161169 'import pathlib' ,
162170 "parser = argparse.ArgumentParser()" ,
163- * [f'parser.add_argument("--{ variable .name }
164- f'type={ variable .argparse_typeof }
165- f'required={ variable .required }
166- for variable in variables ],
171+ * add_args ,
167172 "args = parser.parse_args()" ,
168- f"main({ ',' .join ([f'{ v .name } { v .name } for v in variables if v .is_input ])}
173+ f"main({ ',' .join ([f'{ v .name } { v .name } ' for v in variables if v .is_input ])}
169174 ]],
170175 ])
171176 main_function = ast .parse (main_template_code )
172177 [node for node in main_function .body if isinstance (node , ast .FunctionDef ) and node .name == 'main' ][0 ] \
173178 .body = tree .body
174179 return astor .to_source (main_function )
175180
181+ @classmethod
182+ def __get_add_arguments__ (cls , variables ):
183+ args = []
184+ for variable in variables :
185+ is_array = variable .cwl_typeof .endswith ('[]' )
186+ is_optional = variable .cwl_typeof .endswith ('?' )
187+ arg : str = f'parser.add_argument("--{ variable .name }
188+ arg += f'type={ variable .argparse_typeof }
189+ arg += f'required={ variable .required }
190+ if is_array :
191+ arg += f'nargs="+", '
192+ if is_optional :
193+ arg += f'default=None, '
194+ arg = arg .strip ()
195+ arg += ')'
196+ args .append (arg )
197+ return args
198+
176199 def cwl_command_line_tool (self , docker_image_id : str = 'jn2cwl:latest' ) -> Dict :
177200 """
178201 Creates the description of the CWL Command Line Tool.
0 commit comments