Skip to content

Commit

Permalink
Cleanup and optimize the way gast builds its AST
Browse files Browse the repository at this point in the history 
This gives a 5 to 10% speedup in ast generation depending on the input. As gast
relies on the tree generated by the ast module and does not create any sharing
with the original tree, we're slower, but that's what make it a clean and simple
module :-)
  • Loading branch information
@serge-sans-paille
serge-sans-paille committed Jul 23, 2021
1 parent 80cf94c commit 2271bc2
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 82 deletions.
94 changes: 33 additions & 61 deletions gast/ast3.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,7 @@ def visit_alias(self, node):

def visit_ExtSlice(self, node):
new_node = gast.Tuple(self._visit(node.dims), gast.Load())
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_Index(self, node):
return self._visit(node.value)
Expand Down Expand Up @@ -50,8 +49,7 @@ def visit_Num(self, node):
node.n,
None,
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_Ellipsis(self, node):
new_node = gast.Constant(
Expand All @@ -67,16 +65,14 @@ def visit_Str(self, node):
node.s,
None,
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_Bytes(self, node):
new_node = gast.Constant(
node.s,
None,
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_FunctionDef(self, node):
new_node = gast.FunctionDef(
Expand All @@ -87,8 +83,7 @@ def visit_FunctionDef(self, node):
self._visit(node.returns),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_AsyncFunctionDef(self, node):
new_node = gast.AsyncFunctionDef(
Expand All @@ -99,8 +94,7 @@ def visit_AsyncFunctionDef(self, node):
self._visit(node.returns),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_For(self, node):
new_node = gast.For(
Expand All @@ -110,8 +104,7 @@ def visit_For(self, node):
self._visit(node.orelse),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_AsyncFor(self, node):
new_node = gast.AsyncFor(
Expand All @@ -121,26 +114,23 @@ def visit_AsyncFor(self, node):
self._visit(node.orelse),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_With(self, node):
new_node = gast.With(
self._visit(node.items),
self._visit(node.body),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_AsyncWith(self, node):
new_node = gast.AsyncWith(
self._visit(node.items),
self._visit(node.body),
None, # type_comment
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_Call(self, node):
if sys.version_info.minor < 5:
Expand All @@ -166,8 +156,7 @@ def visit_Call(self, node):
self._visit(node.args) + starred,
self._visit(node.keywords) + kwargs,
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_NameConstant(self, node):
if node.value is None:
Expand All @@ -176,8 +165,7 @@ def visit_NameConstant(self, node):
new_node = gast.Constant(True, None)
elif node.value is False:
new_node = gast.Constant(False, None)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_arguments(self, node):
new_node = gast.arguments(
Expand All @@ -189,42 +177,38 @@ def visit_arguments(self, node):
self._visit(node.kwarg),
self._visit(node.defaults),
)
gast.copy_location(new_node, node)
return new_node
return gast.copy_location(new_node, node)

def visit_Name(self, node):
new_node = gast.Name(
self._visit(node.id),
node.id, # micro-optimization here, don't call self._visit
self._visit(node.ctx),
None,
None,
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_arg(self, node):
if sys.version_info.minor < 8:
extra_args = [None]
extra_arg = None
else:
extra_args = [self._visit(node.type_comment)]
extra_arg = self._visit(node.type_comment)

new_node = gast.Name(
self._visit(node.arg),
node.arg, # micro-optimization here, don't call self._visit
gast.Param(),
self._visit(node.annotation),
*extra_args # type_comment
extra_arg # type_comment
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_ExceptHandler(self, node):
if node.name:
new_node = gast.ExceptHandler(
self._visit(node.type),
gast.Name(node.name, gast.Store(), None, None),
self._visit(node.body))
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)
else:
return self.generic_visit(node)

Expand Down Expand Up @@ -274,17 +258,15 @@ def adjust_slice(s):
new_slice,
self._visit(node.ctx),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_Assign(self, node):
new_node = ast.Assign(
self._visit(node.targets),
self._visit(node.value),
)

ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

if sys.version_info.minor < 8:

Expand All @@ -305,8 +287,7 @@ def visit_Constant(self, node):
new_node = ast.Str(node.value)
else:
new_node = ast.Bytes(node.value)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def _make_arg(self, node):
if node is None:
Expand All @@ -329,8 +310,7 @@ def visit_Name(self, node):
self._visit(node.id),
self._visit(node.ctx),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_ExceptHandler(self, node):
if node.name:
Expand Down Expand Up @@ -366,8 +346,7 @@ def visit_Call(self, node):
self._visit(starargs),
self._visit(kwargs),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_ClassDef(self, node):
self.generic_visit(node)
Expand All @@ -392,8 +371,7 @@ def visit_FunctionDef(self, node):
self._visit(node.decorator_list),
self._visit(node.returns),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_AsyncFunctionDef(self, node):
new_node = ast.AsyncFunctionDef(
Expand All @@ -403,8 +381,7 @@ def visit_AsyncFunctionDef(self, node):
self._visit(node.decorator_list),
self._visit(node.returns),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_For(self, node):
new_node = ast.For(
Expand All @@ -413,8 +390,7 @@ def visit_For(self, node):
self._visit(node.body),
self._visit(node.orelse),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_AsyncFor(self, node):
new_node = ast.AsyncFor(
Expand All @@ -424,33 +400,29 @@ def visit_AsyncFor(self, node):
self._visit(node.orelse),
None, # type_comment
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_With(self, node):
new_node = ast.With(
self._visit(node.items),
self._visit(node.body),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_AsyncWith(self, node):
new_node = ast.AsyncWith(
self._visit(node.items),
self._visit(node.body),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_Call(self, node):
new_node = ast.Call(
self._visit(node.func),
self._visit(node.args),
self._visit(node.keywords),
)
ast.copy_location(new_node, node)
return new_node
return ast.copy_location(new_node, node)

def visit_arguments(self, node):
extra_args = [self._make_arg(node.vararg),
Expand Down
21 changes: 13 additions & 8 deletions gast/astn.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,24 +7,29 @@ def _generate_translators(to):
class Translator(ast.NodeTransformer):

def _visit(self, node):
if isinstance(node, list):
return [self._visit(n) for n in node]
elif isinstance(node, ast.AST):
if isinstance(node, ast.AST):
return self.visit(node)
elif isinstance(node, list):
return [self._visit(n) for n in node]
else:
return node

def generic_visit(self, node):
cls = type(node).__name__
# handle nodes that are not part of the AST
if not hasattr(to, cls):
try:
new_node = getattr(to, cls)()
except AttributeError:
# handle nodes that are not part of the AST
return
new_node = getattr(to, cls)()

for field in node._fields:
setattr(new_node, field, self._visit(getattr(node, field)))
for attr in getattr(node, '_attributes'):
if hasattr(node, attr):

for attr in node._attributes:
try:
setattr(new_node, attr, getattr(node, attr))
except AttributeError:
pass
return new_node

return Translator
Expand Down
28 changes: 15 additions & 13 deletions gast/gast.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,25 +12,27 @@ class TypeIgnore(AST):


def _make_node(Name, Fields, Attributes, Bases):
NBFields = len(Fields)

def create_node(self, *args, **kwargs):
nbparam = len(args) + len(kwargs)
assert nbparam in (0, len(Fields)), \
"Bad argument number for {}: {}, expecting {}".\
format(Name, nbparam, len(Fields))
self._fields = Fields
self._attributes = Attributes
for argname, argval in zip(self._fields, args):
setattr(self, argname, argval)
for argname, argval in kwargs.items():
assert argname in Fields, \
"Invalid Keyword argument for {}: {}".format(Name, argname)
setattr(self, argname, argval)
if args:
if len(args) != NBFields:
raise TypeError(
"{} constructor takes either 0 or {} positional arguments".
format(Name, NBFields))
for argname, argval in zip(Fields, args):
setattr(self, argname, argval)
if kwargs:
for argname, argval in kwargs.items():
setattr(self, argname, argval)

setattr(_sys.modules[__name__],
Name,
type(Name,
Bases,
{'__init__': create_node}))
{'__init__': create_node,
'_fields': Fields,
'_attributes': Attributes}))


_nodes = (
Expand Down

0 comments on commit 2271bc2

Please sign in to comment.