1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
|
"""Helpers for AST (Abstract Syntax Tree)."""
from __future__ import annotations
import ast
from typing import overload
OPERATORS: dict[type[ast.AST], str] = {
ast.Add: "+",
ast.And: "and",
ast.BitAnd: "&",
ast.BitOr: "|",
ast.BitXor: "^",
ast.Div: "/",
ast.FloorDiv: "//",
ast.Invert: "~",
ast.LShift: "<<",
ast.MatMult: "@",
ast.Mult: "*",
ast.Mod: "%",
ast.Not: "not",
ast.Pow: "**",
ast.Or: "or",
ast.RShift: ">>",
ast.Sub: "-",
ast.UAdd: "+",
ast.USub: "-",
}
@overload
def unparse(node: None, code: str = '') -> None:
...
@overload
def unparse(node: ast.AST, code: str = '') -> str:
...
def unparse(node: ast.AST | None, code: str = '') -> str | None:
"""Unparse an AST to string."""
if node is None:
return None
elif isinstance(node, str):
return node
return _UnparseVisitor(code).visit(node)
# a greatly cut-down version of `ast._Unparser`
class _UnparseVisitor(ast.NodeVisitor):
def __init__(self, code: str = '') -> None:
self.code = code
def _visit_op(self, node: ast.AST) -> str:
return OPERATORS[node.__class__]
for _op in OPERATORS:
locals()[f'visit_{_op.__name__}'] = _visit_op
def visit_arg(self, node: ast.arg) -> str:
if node.annotation:
return f"{node.arg}: {self.visit(node.annotation)}"
else:
return node.arg
def _visit_arg_with_default(self, arg: ast.arg, default: ast.AST | None) -> str:
"""Unparse a single argument to a string."""
name = self.visit(arg)
if default:
if arg.annotation:
name += " = %s" % self.visit(default)
else:
name += "=%s" % self.visit(default)
return name
def visit_arguments(self, node: ast.arguments) -> str:
defaults: list[ast.expr | None] = list(node.defaults)
positionals = len(node.args)
posonlyargs = len(node.posonlyargs)
positionals += posonlyargs
for _ in range(len(defaults), positionals):
defaults.insert(0, None)
kw_defaults: list[ast.expr | None] = list(node.kw_defaults)
for _ in range(len(kw_defaults), len(node.kwonlyargs)):
kw_defaults.insert(0, None)
args: list[str] = []
for i, arg in enumerate(node.posonlyargs):
args.append(self._visit_arg_with_default(arg, defaults[i]))
if node.posonlyargs:
args.append('/')
for i, arg in enumerate(node.args):
args.append(self._visit_arg_with_default(arg, defaults[i + posonlyargs]))
if node.vararg:
args.append("*" + self.visit(node.vararg))
if node.kwonlyargs and not node.vararg:
args.append('*')
for i, arg in enumerate(node.kwonlyargs):
args.append(self._visit_arg_with_default(arg, kw_defaults[i]))
if node.kwarg:
args.append("**" + self.visit(node.kwarg))
return ", ".join(args)
def visit_Attribute(self, node: ast.Attribute) -> str:
return f"{self.visit(node.value)}.{node.attr}"
def visit_BinOp(self, node: ast.BinOp) -> str:
# Special case ``**`` to not have surrounding spaces.
if isinstance(node.op, ast.Pow):
return "".join(map(self.visit, (node.left, node.op, node.right)))
return " ".join(self.visit(e) for e in [node.left, node.op, node.right])
def visit_BoolOp(self, node: ast.BoolOp) -> str:
op = " %s " % self.visit(node.op)
return op.join(self.visit(e) for e in node.values)
def visit_Call(self, node: ast.Call) -> str:
args = ', '.join([self.visit(e) for e in node.args]
+ [f"{k.arg}={self.visit(k.value)}" for k in node.keywords])
return f"{self.visit(node.func)}({args})"
def visit_Constant(self, node: ast.Constant) -> str:
if node.value is Ellipsis:
return "..."
elif isinstance(node.value, (int, float, complex)):
if self.code:
return ast.get_source_segment(self.code, node) or repr(node.value)
else:
return repr(node.value)
else:
return repr(node.value)
def visit_Dict(self, node: ast.Dict) -> str:
keys = (self.visit(k) for k in node.keys if k is not None)
values = (self.visit(v) for v in node.values)
items = (k + ": " + v for k, v in zip(keys, values))
return "{" + ", ".join(items) + "}"
def visit_Index(self, node: ast.Index) -> str:
return self.visit(node.value)
def visit_Lambda(self, node: ast.Lambda) -> str:
return "lambda %s: ..." % self.visit(node.args)
def visit_List(self, node: ast.List) -> str:
return "[" + ", ".join(self.visit(e) for e in node.elts) + "]"
def visit_Name(self, node: ast.Name) -> str:
return node.id
def visit_Set(self, node: ast.Set) -> str:
return "{" + ", ".join(self.visit(e) for e in node.elts) + "}"
def visit_Subscript(self, node: ast.Subscript) -> str:
def is_simple_tuple(value: ast.AST) -> bool:
return (
isinstance(value, ast.Tuple) and
bool(value.elts) and
not any(isinstance(elt, ast.Starred) for elt in value.elts)
)
if is_simple_tuple(node.slice):
elts = ", ".join(self.visit(e) for e in node.slice.elts) # type: ignore
return f"{self.visit(node.value)}[{elts}]"
elif isinstance(node.slice, ast.Index) and is_simple_tuple(node.slice.value):
elts = ", ".join(self.visit(e) for e in node.slice.value.elts) # type: ignore
return f"{self.visit(node.value)}[{elts}]"
else:
return f"{self.visit(node.value)}[{self.visit(node.slice)}]"
def visit_UnaryOp(self, node: ast.UnaryOp) -> str:
# UnaryOp is one of {UAdd, USub, Invert, Not}, which refer to ``+x``,
# ``-x``, ``~x``, and ``not x``. Only Not needs a space.
if isinstance(node.op, ast.Not):
return f"{self.visit(node.op)} {self.visit(node.operand)}"
return f"{self.visit(node.op)}{self.visit(node.operand)}"
def visit_Tuple(self, node: ast.Tuple) -> str:
if len(node.elts) == 0:
return "()"
elif len(node.elts) == 1:
return "(%s,)" % self.visit(node.elts[0])
else:
return "(" + ", ".join(self.visit(e) for e in node.elts) + ")"
def generic_visit(self, node):
raise NotImplementedError('Unable to parse %s object' % type(node).__name__)
|
