@@ -70,22 +70,27 @@ def build_param(ast, name, default):
7070 # MAKE_FUNCTION_... or MAKE_CLOSURE_...
7171 assert node [- 1 ].type .startswith ('MAKE_' )
7272
73- annotate_return = None
7473 annotate_arg = node [annotate_last ]
74+ annotate_args = {}
7575
7676 if (annotate_arg == 'annotate_arg'
77- and annotate_arg [0 ] == 'LOAD_CONST'
77+ and annotate_arg [0 ] in ( 'LOAD_CONST' , 'LOAD_NAME' )
7878 and isinstance (annotate_arg [0 ].attr , tuple )):
7979 annotate_tup = annotate_arg [0 ].attr
80- if annotate_tup [- 1 ] == 'return' :
81- annotate_return = node [annotate_last - 1 ][0 ].attr
82- pass
80+ i = - 1
81+ j = annotate_last - 1
82+ l = - len (node )
83+ while j >= l and node [j ] == 'annotate_arg' :
84+ annotate_args [annotate_tup [i ]] = (node [j ][0 ].attr ,
85+ node [j ][0 ] == 'LOAD_CONST' )
86+ i -= 1
87+ j -= 1
8388
8489 args_node = node [- 1 ]
8590 if isinstance (args_node .attr , tuple ):
8691 # positional args are before kwargs
8792 defparams = node [:args_node .attr [0 ]]
88- pos_args , kw_args , annotate_args = args_node .attr
93+ pos_args , kw_args , annotate_argc = args_node .attr
8994 else :
9095 defparams = node [:args_node .attr ]
9196 kw_args = 0
@@ -139,12 +144,26 @@ def build_param(ast, name, default):
139144 argc += 1
140145
141146 i = len (paramnames ) - len (defparams )
142- self .write (", " .join (paramnames [:i ]))
147+ suffix = ''
148+ for param in paramnames [:i ]:
149+ self .write (suffix , param )
150+ if param in annotate_args :
151+ value , string = annotate_args [param ]
152+ if string :
153+ self .write (': "%s"' % value )
154+ else :
155+ self .write (': %s' % value )
156+ suffix = ', '
157+
143158 suffix = ', ' if i > 0 else ''
144159 for n in node :
145160 if n == 'pos_arg' :
146161 self .write (suffix )
147- self .write (paramnames [i ] + '=' )
162+ param = paramnames [i ]
163+ self .write (param )
164+ if param in annotate_args :
165+ self .write (':"%s' % annotate_args [param ])
166+ self .write ('=' )
148167 i += 1
149168 self .preorder (n )
150169 if (line_number != self .line_number ):
@@ -187,8 +206,13 @@ def build_param(ast, name, default):
187206 self .write (": " )
188207 else :
189208 self .write (')' )
190- if annotate_return :
191- self .write (' -> "%s"' % annotate_return )
209+ if 'return' in annotate_args :
210+ value , string = annotate_args ['return' ]
211+ if string :
212+ self .write (' -> "%s"' % value )
213+ else :
214+ self .write (' -> %s' % value )
215+
192216 self .println (":" )
193217
194218 if (len (code .co_consts ) > 0 and
@@ -247,7 +271,7 @@ def build_param(ast, name, default):
247271 if isinstance (args_node .attr , tuple ):
248272 # positional args are after kwargs
249273 defparams = node [1 :args_node .attr [0 ]+ 1 ]
250- pos_args , kw_args , annotate_args = args_node .attr
274+ pos_args , kw_args , annotate_argc = args_node .attr
251275 else :
252276 defparams = node [:args_node .attr ]
253277 kw_args = 0
@@ -377,7 +401,7 @@ def build_param(ast, name, default):
377401 else :
378402 # positional args are before kwargs
379403 defparams = node [:args_node .attr [0 ]]
380- pos_args , kw_args , annotate_args = args_node .attr
404+ pos_args , kw_args , annotate_argc = args_node .attr
381405 else :
382406 defparams = node [:args_node .attr ]
383407 kw_args = 0
0 commit comments