44import cStringIO
55import os
66
7+ def interleave (inter , f , seq ):
8+ """Call f on each item in seq, calling inter() in between.
9+ """
10+ seq = iter (seq )
11+ try :
12+ f (seq .next ())
13+ except StopIteration :
14+ pass
15+ else :
16+ for x in seq :
17+ inter ()
18+ f (x )
19+
720class Unparser :
821 """Methods in this class recursively traverse an AST and
922 output source code for the abstract syntax; original formatting
@@ -63,26 +76,13 @@ def _Expr(self, tree):
6376
6477 def _Import (self , t ):
6578 self .fill ("import " )
66- first = True
67- for a in t .names :
68- if first :
69- first = False
70- else :
71- self .write (", " )
72- self .write (a .name )
73- if a .asname :
74- self .write (" as " + a .asname )
79+ interleave (lambda : self .write (", " ), self .dispatch , t .names )
7580
7681 def _ImportFrom (self , t ):
7782 self .fill ("from " )
7883 self .write (t .module )
7984 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 )
85+ interleave (lambda : self .write (", " ), self .dispatch , t .names )
8686 # XXX(jpe) what is level for?
8787
8888 def _Assign (self , t ):
@@ -99,8 +99,9 @@ def _AugAssign(self, t):
9999 self .dispatch (t .value )
100100
101101 def _Return (self , t ):
102- self .fill ("return " )
102+ self .fill ("return" )
103103 if t .value :
104+ self .write (" " )
104105 self .dispatch (t .value )
105106
106107 def _Pass (self , t ):
@@ -148,18 +149,16 @@ def _Print(self, t):
148149 self .write ("," )
149150
150151 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 )
152+ self .fill ("global " )
153+ interleave (lambda : self .write (", " ), self .write , t .names )
156154
157155 def _Yield (self , t ):
158- self .fill ("yield" )
156+ self .write ("(" )
157+ self .write ("yield" )
159158 if t .value :
160- self .write (" ( " )
159+ self .write (" " )
161160 self .dispatch (t .value )
162- self .write (")" )
161+ self .write (")" )
163162
164163 def _Raise (self , t ):
165164 self .fill ('raise ' )
@@ -198,8 +197,9 @@ def _TryFinally(self, t):
198197 self .leave ()
199198
200199 def _excepthandler (self , t ):
201- self .fill ("except " )
200+ self .fill ("except" )
202201 if t .type :
202+ self .write (" " )
203203 self .dispatch (t .type )
204204 if t .name :
205205 self .write (", " )
@@ -299,9 +299,7 @@ def _Num(self, t):
299299
300300 def _List (self , t ):
301301 self .write ("[" )
302- for e in t .elts :
303- self .dispatch (e )
304- self .write (", " )
302+ interleave (lambda : self .write (", " ), self .dispatch , t .elts )
305303 self .write ("]" )
306304
307305 def _ListComp (self , t ):
@@ -328,30 +326,31 @@ def _comprehension(self, t):
328326 self .dispatch (if_clause )
329327
330328 def _IfExp (self , t ):
329+ self .write ("(" )
331330 self .dispatch (t .body )
332331 self .write (" if " )
333332 self .dispatch (t .test )
334- if t . orelse :
335- self .write ( " else " )
336- self .dispatch ( t . orelse )
333+ self . write ( " else " )
334+ self .dispatch ( t . orelse )
335+ self .write ( ")" )
337336
338337 def _Dict (self , t ):
339338 self .write ("{" )
340- for k , v in zip ( t . keys , t . values ):
339+ def writem (( k , v ) ):
341340 self .dispatch (k )
342- self .write (" : " )
341+ self .write (": " )
343342 self .dispatch (v )
344- self .write (", " )
343+ interleave ( lambda : self .write (", " ), writem , zip ( t . keys , t . values ) )
345344 self .write ("}" )
346345
347346 def _Tuple (self , t ):
348- if not t .elts :
349- self .write ("()" )
350- return
351347 self .write ("(" )
352- for e in t .elts :
353- self .dispatch (e )
354- self .write (", " )
348+ if len (t .elts ) == 1 :
349+ (elt ,) = t .elts
350+ self .dispatch (elt )
351+ self .write ("," )
352+ else :
353+ interleave (lambda : self .write (", " ), self .dispatch , t .elts )
355354 self .write (")" )
356355
357356 unop = {"Invert" :"~" , "Not" : "not" , "UAdd" :"+" , "USub" :"-" }
@@ -367,7 +366,7 @@ def _UnaryOp(self, t):
367366 def _BinOp (self , t ):
368367 self .write ("(" )
369368 self .dispatch (t .left )
370- self .write (") " + self .binop [t .op .__class__ .__name__ ] + "( " )
369+ self .write (" " + self .binop [t .op .__class__ .__name__ ] + " " )
371370 self .dispatch (t .right )
372371 self .write (")" )
373372
@@ -377,17 +376,15 @@ def _Compare(self, t):
377376 self .write ("(" )
378377 self .dispatch (t .left )
379378 for o , e in zip (t .ops , t .comparators ):
380- self .write (") " + self .cmpops [o .__class__ .__name__ ] + " ( " )
379+ self .write (" " + self .cmpops [o .__class__ .__name__ ] + " " )
381380 self .dispatch (e )
382381 self .write (")" )
383382
384383 boolops = {_ast .And : 'and' , _ast .Or : 'or' }
385384 def _BoolOp (self , t ):
386385 self .write ("(" )
387- self .dispatch (t .values [0 ])
388- for v in t .values [1 :]:
389- self .write (" %s " % self .boolops [t .op .__class__ ])
390- self .dispatch (v )
386+ s = " %s " % self .boolops [t .op .__class__ ]
387+ interleave (lambda : self .write (s ), self .dispatch , t .values )
391388 self .write (")" )
392389
393390 def _Attribute (self ,t ):
@@ -443,10 +440,7 @@ def _Slice(self, t):
443440 self .dispatch (t .step )
444441
445442 def _ExtSlice (self , t ):
446- for i , d in enumerate (t .dims ):
447- if i != 0 :
448- self .write (': ' )
449- self .dispatch (d )
443+ interleave (lambda : self .write (', ' ), self .dispatch , t .dims )
450444
451445 # others
452446 def _arguments (self , t ):
@@ -482,9 +476,14 @@ def _Lambda(self, t):
482476 self .write (": " )
483477 self .dispatch (t .body )
484478
479+ def _alias (self , t ):
480+ self .write (t .name )
481+ if t .asname :
482+ self .write (" as " + t .asname )
483+
485484def roundtrip (filename , output = sys .stdout ):
486485 source = open (filename ).read ()
487- tree = compile (source , filename , "exec" , 0x400 )
486+ tree = compile (source , filename , "exec" , _ast . PyCF_ONLY_AST )
488487 Unparser (tree , output )
489488
490489
0 commit comments