11"Usage: unparse.py <path to source file>"
22import sys
3+ import _ast
4+ import cStringIO
5+ import os
36
47class Unparser :
58 """Methods in this class recursively traverse an AST and
@@ -70,6 +73,18 @@ def _Import(self, t):
7073 if a .asname :
7174 self .write (" as " + a .asname )
7275
76+ def _ImportFrom (self , t ):
77+ self .fill ("from " )
78+ self .write (t .module )
79+ self .write (" import " )
80+ for i , a in enumerate (t .names ):
81+ if i == 0 :
82+ self .write (", " )
83+ self .write (a .name )
84+ if a .asname :
85+ self .write (" as " + a .asname )
86+ # XXX(jpe) what is level for?
87+
7388 def _Assign (self , t ):
7489 self .fill ()
7590 for target in t .targets :
@@ -88,6 +103,36 @@ def _Return(self, t):
88103 if t .value :
89104 self .dispatch (t .value )
90105
106+ def _Pass (self , t ):
107+ self .fill ("pass" )
108+
109+ def _Break (self , t ):
110+ self .fill ("break" )
111+
112+ def _Continue (self , t ):
113+ self .fill ("continue" )
114+
115+ def _Delete (self , t ):
116+ self .fill ("del " )
117+ self .dispatch (t .targets )
118+
119+ def _Assert (self , t ):
120+ self .fill ("assert " )
121+ self .dispatch (t .test )
122+ if t .msg :
123+ self .write (", " )
124+ self .dispatch (t .msg )
125+
126+ def _Exec (self , t ):
127+ self .fill ("exec " )
128+ self .dispatch (t .body )
129+ if t .globals :
130+ self .write (" in " )
131+ self .dispatch (t .globals )
132+ if t .locals :
133+ self .write (", " )
134+ self .dispatch (t .locals )
135+
91136 def _Print (self , t ):
92137 self .fill ("print " )
93138 do_comma = False
@@ -102,6 +147,67 @@ def _Print(self, t):
102147 if not t .nl :
103148 self .write ("," )
104149
150+ def _Global (self , t ):
151+ self .fill ("global" )
152+ for i , n in enumerate (t .names ):
153+ if i != 0 :
154+ self .write ("," )
155+ self .write (" " + n )
156+
157+ def _Yield (self , t ):
158+ self .fill ("yield" )
159+ if t .value :
160+ self .write (" (" )
161+ self .dispatch (t .value )
162+ self .write (")" )
163+
164+ def _Raise (self , t ):
165+ self .fill ('raise ' )
166+ if t .type :
167+ self .dispatch (t .type )
168+ if t .inst :
169+ self .write (", " )
170+ self .dispatch (t .inst )
171+ if t .tback :
172+ self .write (", " )
173+ self .dispatch (t .tback )
174+
175+ def _TryExcept (self , t ):
176+ self .fill ("try" )
177+ self .enter ()
178+ self .dispatch (t .body )
179+ self .leave ()
180+
181+ for ex in t .handlers :
182+ self .dispatch (ex )
183+ if t .orelse :
184+ self .fill ("else" )
185+ self .enter ()
186+ self .dispatch (t .orelse )
187+ self .leave ()
188+
189+ def _TryFinally (self , t ):
190+ self .fill ("try" )
191+ self .enter ()
192+ self .dispatch (t .body )
193+ self .leave ()
194+
195+ self .fill ("finally" )
196+ self .enter ()
197+ self .dispatch (t .finalbody )
198+ self .leave ()
199+
200+ def _excepthandler (self , t ):
201+ self .fill ("except " )
202+ if t .type :
203+ self .dispatch (t .type )
204+ if t .name :
205+ self .write (", " )
206+ self .dispatch (t .name )
207+ self .enter ()
208+ self .dispatch (t .body )
209+ self .leave ()
210+
105211 def _ClassDef (self , t ):
106212 self .write ("\n " )
107213 self .fill ("class " + t .name )
@@ -119,9 +225,24 @@ def _FunctionDef(self, t):
119225 self .write ("\n " )
120226 self .fill ("def " + t .name + "(" )
121227 self .dispatch (t .args )
228+ self .write (")" )
229+ self .enter ()
230+ self .dispatch (t .body )
231+ self .leave ()
232+
233+ def _For (self , t ):
234+ self .fill ("for " )
235+ self .dispatch (t .target )
236+ self .write (" in " )
237+ self .dispatch (t .iter )
122238 self .enter ()
123239 self .dispatch (t .body )
124240 self .leave ()
241+ if t .orelse :
242+ self .fill ("else" )
243+ self .enter ()
244+ self .dispatch (t .orelse )
245+ self .leave
125246
126247 def _If (self , t ):
127248 self .fill ("if " )
@@ -136,11 +257,9 @@ def _If(self, t):
136257 self .dispatch (t .orelse )
137258 self .leave ()
138259
139- def _For (self , t ):
140- self .fill ("for " )
141- self .dispatch (t .target )
142- self .write (" in " )
143- self .dispatch (t .iter )
260+ def _While (self , t ):
261+ self .fill ("while " )
262+ self .dispatch (t .test )
144263 self .enter ()
145264 self .dispatch (t .body )
146265 self .leave ()
@@ -150,13 +269,28 @@ def _For(self, t):
150269 self .dispatch (t .orelse )
151270 self .leave
152271
272+ def _With (self , t ):
273+ self .fill ("with " )
274+ self .dispatch (t .context_expr )
275+ if t .optional_vars :
276+ self .write (" as " )
277+ self .dispatch (t .optional_vars )
278+ self .enter ()
279+ self .dispatch (t .body )
280+ self .leave ()
281+
153282 # expr
154283 def _Str (self , tree ):
155284 self .write (repr (tree .s ))
156285
157286 def _Name (self , t ):
158287 self .write (t .id )
159288
289+ def _Repr (self , t ):
290+ self .write ("`" )
291+ self .dispatch (t .value )
292+ self .write ("`" )
293+
160294 def _Num (self , t ):
161295 self .write (repr (t .n ))
162296
@@ -167,6 +301,37 @@ def _List(self, t):
167301 self .write (", " )
168302 self .write ("]" )
169303
304+ def _ListComp (self , t ):
305+ self .write ("[" )
306+ self .dispatch (t .elt )
307+ for gen in t .generators :
308+ self .dispatch (gen )
309+ self .write ("]" )
310+
311+ def _GeneratorExp (self , t ):
312+ self .write ("(" )
313+ self .dispatch (t .elt )
314+ for gen in t .generators :
315+ self .dispatch (gen )
316+ self .write (")" )
317+
318+ def _comprehension (self , t ):
319+ self .write (" for " )
320+ self .dispatch (t .target )
321+ self .write (" in " )
322+ self .dispatch (t .iter )
323+ for if_clause in t .ifs :
324+ self .write (" if " )
325+ self .dispatch (if_clause )
326+
327+ def _IfExp (self , t ):
328+ self .dispatch (t .body )
329+ self .write (" if " )
330+ self .dispatch (t .test )
331+ if t .orelse :
332+ self .write (" else " )
333+ self .dispatch (t .orelse )
334+
170335 def _Dict (self , t ):
171336 self .write ("{" )
172337 for k ,v in zip (t .keys , t .values ):
@@ -194,8 +359,8 @@ def _UnaryOp(self, t):
194359 self .write (")" )
195360
196361 binop = { "Add" :"+" , "Sub" :"-" , "Mult" :"*" , "Div" :"/" , "Mod" :"%" ,
197- "RShift" :"<<" , "BitOr" :"|" , "BitXor" :"^" , "BitAnd" :"&" ,
198- "FloorDiv" :"//" }
362+ "LShift" : ">>" , " RShift" :"<<" , "BitOr" :"|" , "BitXor" :"^" , "BitAnd" :"&" ,
363+ "FloorDiv" :"//" , "Pow" : "**" }
199364 def _BinOp (self , t ):
200365 self .write ("(" )
201366 self .dispatch (t .left )
@@ -213,6 +378,15 @@ def _Compare(self, t):
213378 self .dispatch (e )
214379 self .write (")" )
215380
381+ boolops = {_ast .And : 'and' , _ast .Or : 'or' }
382+ def _BoolOp (self , t ):
383+ self .write ("(" )
384+ self .dispatch (t .values [0 ])
385+ for v in t .values [1 :]:
386+ self .write (" %s " % self .boolops [t .op .__class__ ])
387+ self .dispatch (v )
388+ self .write (")" )
389+
216390 def _Attribute (self ,t ):
217391 self .dispatch (t .value )
218392 self .write ("." )
@@ -234,12 +408,12 @@ def _Call(self, t):
234408 if comma : self .write (", " )
235409 else : comma = True
236410 self .write ("*" )
237- self .dispatch (t .stararg )
411+ self .dispatch (t .starargs )
238412 if t .kwargs :
239413 if comma : self .write (", " )
240414 else : comma = True
241415 self .write ("**" )
242- self .dispatch (t .stararg )
416+ self .dispatch (t .kwargs )
243417 self .write (")" )
244418
245419 def _Subscript (self , t ):
@@ -249,6 +423,9 @@ def _Subscript(self, t):
249423 self .write ("]" )
250424
251425 # slice
426+ def _Ellipsis (self , t ):
427+ self .write ("..." )
428+
252429 def _Index (self , t ):
253430 self .dispatch (t .value )
254431
@@ -262,6 +439,12 @@ def _Slice(self, t):
262439 self .write (":" )
263440 self .dispatch (t .step )
264441
442+ def _ExtSlice (self , t ):
443+ for i , d in enumerate (t .dims ):
444+ if i != 0 :
445+ self .write (': ' )
446+ self .dispatch (d )
447+
265448 # others
266449 def _arguments (self , t ):
267450 first = True
@@ -283,13 +466,51 @@ def _arguments(self, t):
283466 if t .kwarg :
284467 if first :first = False
285468 else : self .write (", " )
286- self .write ("**" + self .kwarg )
287- self .write (")" )
469+ self .write ("**" + t .kwarg )
470+
471+ def _keyword (self , t ):
472+ self .write (t .arg )
473+ self .write ("=" )
474+ self .dispatch (t .value )
475+
476+ def _Lambda (self , t ):
477+ self .write ("lambda " )
478+ self .dispatch (t .args )
479+ self .write (": " )
480+ self .dispatch (t .body )
288481
289- def roundtrip (filename ):
482+ def roundtrip (filename , output = sys . stdout ):
290483 source = open (filename ).read ()
291484 tree = compile (source , filename , "exec" , 0x400 )
292- Unparser (tree )
485+ Unparser (tree , output )
486+
487+
488+
489+ def testdir (a ):
490+ try :
491+ names = [n for n in os .listdir (a ) if n .endswith ('.py' )]
492+ except OSError :
493+ print >> sys .stderr , "Directory not readable: %s" % a
494+ else :
495+ for n in names :
496+ fullname = os .path .join (a , n )
497+ if os .path .isfile (fullname ):
498+ output = cStringIO .StringIO ()
499+ print 'Testing %s' % fullname
500+ try :
501+ roundtrip (fullname , output )
502+ except Exception , e :
503+ print ' Failed to compile, exception is %s' % repr (e )
504+ elif os .path .isdir (fullname ):
505+ testdir (fullname )
506+
507+ def main (args ):
508+ if args [0 ] == '--testdir' :
509+ for a in args [1 :]:
510+ testdir (a )
511+ else :
512+ for a in args :
513+ roundtrip (a )
293514
294515if __name__ == '__main__' :
295- roundtrip (sys .argv [1 ])
516+ main (sys .argv [1 : ])
0 commit comments