1from astn import AstToGAst, GAstToAst
2import ast
3import gast
4
5
6class Ast2ToGAst(AstToGAst):
7
8    # mod
9    def visit_Module(self, node):
10        new_node = gast.Module(
11            self._visit(node.body),
12            []  # type_ignores
13        )
14        return new_node
15
16    # stmt
17    def visit_FunctionDef(self, node):
18        new_node = gast.FunctionDef(
19            self._visit(node.name),
20            self._visit(node.args),
21            self._visit(node.body),
22            self._visit(node.decorator_list),
23            None,  # returns
24            None,  # type_comment
25        )
26        gast.copy_location(new_node, node)
27        new_node.end_lineno = new_node.end_col_offset = None
28        return new_node
29
30    def visit_ClassDef(self, node):
31        new_node = gast.ClassDef(
32            self._visit(node.name),
33            self._visit(node.bases),
34            [],  # keywords
35            self._visit(node.body),
36            self._visit(node.decorator_list),
37        )
38
39        gast.copy_location(new_node, node)
40        new_node.end_lineno = new_node.end_col_offset = None
41        return new_node
42
43    def visit_Assign(self, node):
44        new_node = gast.Assign(
45            self._visit(node.targets),
46            self._visit(node.value),
47            None,  # type_comment
48        )
49
50        gast.copy_location(new_node, node)
51        new_node.end_lineno = new_node.end_col_offset = None
52        return new_node
53
54    def visit_For(self, node):
55        new_node = gast.For(
56            self._visit(node.target),
57            self._visit(node.iter),
58            self._visit(node.body),
59            self._visit(node.orelse),
60            []  # type_comment
61        )
62        gast.copy_location(new_node, node)
63        new_node.end_lineno = new_node.end_col_offset = None
64        return new_node
65
66    def visit_With(self, node):
67        new_node = gast.With(
68            [gast.withitem(
69                self._visit(node.context_expr),
70                self._visit(node.optional_vars)
71            )],
72            self._visit(node.body),
73            None,  # type_comment
74        )
75        gast.copy_location(new_node, node)
76        new_node.end_lineno = new_node.end_col_offset = None
77        return new_node
78
79    def visit_Raise(self, node):
80        ntype = self._visit(node.type)
81        ninst = self._visit(node.inst)
82        ntback = self._visit(node.tback)
83
84        what = ntype
85
86        if ninst is not None:
87            what = gast.Call(ntype, [ninst], [])
88            gast.copy_location(what, node)
89            what.end_lineno = what.end_col_offset = None
90
91        if ntback is not None:
92            attr = gast.Attribute(what, 'with_traceback', gast.Load())
93            gast.copy_location(attr, node)
94            attr.end_lineno = attr.end_col_offset = None
95
96            what = gast.Call(
97                attr,
98                [ntback],
99                []
100            )
101            gast.copy_location(what, node)
102            what.end_lineno = what.end_col_offset = None
103
104        new_node = gast.Raise(what, None)
105
106        gast.copy_location(new_node, node)
107        new_node.end_lineno = new_node.end_col_offset = None
108        return new_node
109
110    def visit_TryExcept(self, node):
111        new_node = gast.Try(
112            self._visit(node.body),
113            self._visit(node.handlers),
114            self._visit(node.orelse),
115            []  # finalbody
116        )
117        gast.copy_location(new_node, node)
118        new_node.end_lineno = new_node.end_col_offset = None
119        return new_node
120
121    def visit_TryFinally(self, node):
122        new_node = gast.Try(
123            self._visit(node.body),
124            [],  # handlers
125            [],  # orelse
126            self._visit(node.finalbody)
127        )
128        gast.copy_location(new_node, node)
129        new_node.end_lineno = new_node.end_col_offset = None
130        return new_node
131
132    # expr
133
134    def visit_Name(self, node):
135        new_node = gast.Name(
136            self._visit(node.id),
137            self._visit(node.ctx),
138            None,
139            None,
140        )
141        gast.copy_location(new_node, node)
142        new_node.end_lineno = new_node.end_col_offset = None
143        return new_node
144
145    def visit_Num(self, node):
146        new_node = gast.Constant(
147            node.n,
148            None,
149        )
150        gast.copy_location(new_node, node)
151        new_node.end_lineno = new_node.end_col_offset = None
152        return new_node
153
154    def visit_Subscript(self, node):
155        new_slice = self._visit(node.slice)
156        new_node = gast.Subscript(
157            self._visit(node.value),
158            new_slice,
159            self._visit(node.ctx),
160        )
161        gast.copy_location(new_node, node)
162        new_node.end_lineno = new_node.end_col_offset = None
163        return new_node
164
165    def visit_Ellipsis(self, node):
166        new_node = gast.Constant(
167            Ellipsis,
168            None,
169        )
170        gast.copy_location(new_node, node)
171        new_node.end_lineno = new_node.end_col_offset = None
172        return new_node
173
174    def visit_Index(self, node):
175        return self._visit(node.value)
176
177    def visit_ExtSlice(self, node):
178        new_dims = self._visit(node.dims)
179        new_node = gast.Tuple(new_dims, gast.Load())
180        gast.copy_location(new_node, node)
181        new_node.end_lineno = new_node.end_col_offset = None
182        return new_node
183
184    def visit_Str(self, node):
185        new_node = gast.Constant(
186            node.s,
187            None,
188        )
189        gast.copy_location(new_node, node)
190        new_node.end_lineno = new_node.end_col_offset = None
191        return new_node
192
193    def visit_Call(self, node):
194        if node.starargs:
195            star = gast.Starred(self._visit(node.starargs), gast.Load())
196            gast.copy_location(star, node)
197            star.end_lineno = star.end_col_offset = None
198            starred = [star]
199        else:
200            starred = []
201
202        if node.kwargs:
203            kwargs = [gast.keyword(None, self._visit(node.kwargs))]
204        else:
205            kwargs = []
206
207        new_node = gast.Call(
208            self._visit(node.func),
209            self._visit(node.args) + starred,
210            self._visit(node.keywords) + kwargs,
211        )
212        gast.copy_location(new_node, node)
213        new_node.end_lineno = new_node.end_col_offset = None
214        return new_node
215
216    def visit_comprehension(self, node):
217        new_node = gast.comprehension(
218            target=self._visit(node.target),
219            iter=self._visit(node.iter),
220            ifs=self._visit(node.ifs),
221            is_async=0,
222        )
223        gast.copy_location(new_node, node)
224        new_node.end_lineno = new_node.end_col_offset = None
225        return new_node
226
227    # arguments
228    def visit_arguments(self, node):
229        # missing locations for vararg and kwarg set at function level
230        if node.vararg:
231            vararg = ast.Name(node.vararg, ast.Param())
232        else:
233            vararg = None
234
235        if node.kwarg:
236            kwarg = ast.Name(node.kwarg, ast.Param())
237        else:
238            kwarg = None
239
240        if node.vararg:
241            vararg = ast.Name(node.vararg, ast.Param())
242        else:
243            vararg = None
244
245        new_node = gast.arguments(
246            self._visit(node.args),
247            [],  # posonlyargs
248            self._visit(vararg),
249            [],  # kwonlyargs
250            [],  # kw_defaults
251            self._visit(kwarg),
252            self._visit(node.defaults),
253        )
254        return new_node
255
256    def visit_alias(self, node):
257        new_node = gast.alias(
258            self._visit(node.name),
259            self._visit(node.asname),
260        )
261        new_node.lineno = new_node.col_offset = None
262        new_node.end_lineno = new_node.end_col_offset = None
263        return new_node
264
265
266class GAstToAst2(GAstToAst):
267
268    # mod
269    def visit_Module(self, node):
270        new_node = ast.Module(self._visit(node.body))
271        return new_node
272
273    # stmt
274    def visit_FunctionDef(self, node):
275        new_node = ast.FunctionDef(
276            self._visit(node.name),
277            self._visit(node.args),
278            self._visit(node.body),
279            self._visit(node.decorator_list),
280        )
281        # because node.args doesn't have any location to copy from
282        if node.args.vararg:
283            ast.copy_location(node.args.vararg, node)
284        if node.args.kwarg:
285            ast.copy_location(node.args.kwarg, node)
286
287        ast.copy_location(new_node, node)
288        return new_node
289
290    def visit_ClassDef(self, node):
291        new_node = ast.ClassDef(
292            self._visit(node.name),
293            self._visit(node.bases),
294            self._visit(node.body),
295            self._visit(node.decorator_list),
296        )
297
298        ast.copy_location(new_node, node)
299        return new_node
300
301    def visit_Assign(self, node):
302        new_node = ast.Assign(
303            self._visit(node.targets),
304            self._visit(node.value),
305        )
306
307        ast.copy_location(new_node, node)
308        return new_node
309
310    def visit_For(self, node):
311        new_node = ast.For(
312            self._visit(node.target),
313            self._visit(node.iter),
314            self._visit(node.body),
315            self._visit(node.orelse),
316        )
317
318        ast.copy_location(new_node, node)
319        return new_node
320
321    def visit_With(self, node):
322        new_node = ast.With(
323            self._visit(node.items[0].context_expr),
324            self._visit(node.items[0].optional_vars),
325            self._visit(node.body)
326        )
327        ast.copy_location(new_node, node)
328        return new_node
329
330    def visit_Raise(self, node):
331        if isinstance(node.exc, gast.Call) and \
332           isinstance(node.exc.func, gast.Attribute) and \
333           node.exc.func.attr == 'with_traceback':
334            raised = self._visit(node.exc.func.value)
335            traceback = self._visit(node.exc.args[0])
336        else:
337            raised = self._visit(node.exc)
338            traceback = None
339        new_node = ast.Raise(raised, None, traceback)
340        ast.copy_location(new_node, node)
341        return new_node
342
343    def visit_Try(self, node):
344        if node.finalbody:
345            new_node = ast.TryFinally(
346                self._visit(node.body),
347                self._visit(node.finalbody)
348            )
349        else:
350            new_node = ast.TryExcept(
351                self._visit(node.body),
352                self._visit(node.handlers),
353                self._visit(node.orelse),
354            )
355        ast.copy_location(new_node, node)
356        return new_node
357
358    # expr
359
360    def visit_Name(self, node):
361        new_node = ast.Name(
362            self._visit(node.id),
363            self._visit(node.ctx),
364        )
365        ast.copy_location(new_node, node)
366        return new_node
367
368    def visit_Constant(self, node):
369        if isinstance(node.value, (bool, int, long, float, complex)):
370            new_node = ast.Num(node.value)
371        elif node.value is Ellipsis:
372            new_node = ast.Ellipsis()
373        else:
374            new_node = ast.Str(node.value)
375        ast.copy_location(new_node, node)
376        return new_node
377
378    def visit_Subscript(self, node):
379        def adjust_slice(s):
380            if isinstance(s, (ast.Slice, ast.Ellipsis)):
381                return s
382            else:
383                return ast.Index(s)
384        if isinstance(node.slice, gast.Tuple):
385            new_slice = ast.ExtSlice([adjust_slice(self._visit(elt))
386                                      for elt in node.slice.elts])
387        else:
388            new_slice = adjust_slice(self._visit(node.slice))
389        ast.copy_location(new_slice, node.slice)
390
391        new_node = ast.Subscript(
392            self._visit(node.value),
393            new_slice,
394            self._visit(node.ctx),
395        )
396        ast.copy_location(new_node, node)
397        new_node.end_lineno = new_node.end_col_offset = None
398        return new_node
399
400    def visit_Call(self, node):
401        if node.args and isinstance(node.args[-1], gast.Starred):
402            args = node.args[:-1]
403            starargs = node.args[-1].value
404        else:
405            args = node.args
406            starargs = None
407
408        if node.keywords and node.keywords[-1].arg is None:
409            keywords = node.keywords[:-1]
410            kwargs = node.keywords[-1].value
411        else:
412            keywords = node.keywords
413            kwargs = None
414
415        new_node = ast.Call(
416            self._visit(node.func),
417            self._visit(args),
418            self._visit(keywords),
419            self._visit(starargs),
420            self._visit(kwargs),
421        )
422        ast.copy_location(new_node, node)
423        return new_node
424
425    def visit_arg(self, node):
426        new_node = ast.Name(node.arg, ast.Param())
427        ast.copy_location(new_node, node)
428        return new_node
429
430    # arguments
431    def visit_arguments(self, node):
432        vararg = node.vararg and node.vararg.id
433        kwarg = node.kwarg and node.kwarg.id
434
435        new_node = ast.arguments(
436            self._visit(node.args),
437            self._visit(vararg),
438            self._visit(kwarg),
439            self._visit(node.defaults),
440        )
441        return new_node
442
443    def visit_alias(self, node):
444        new_node = ast.alias(
445            self._visit(node.name),
446            self._visit(node.asname)
447        )
448        return new_node
449
450
451def ast_to_gast(node):
452    return Ast2ToGAst().visit(node)
453
454
455def gast_to_ast(node):
456    return GAstToAst2().visit(node)
457