159 lines
5.9 KiB
Python
159 lines
5.9 KiB
Python
# -*- coding: utf-8; -*-
|
|
'''Fix missing `ctx` attributes and source location info in an AST.'''
|
|
|
|
__all__ = ['fix_missing_ctx', 'fix_missing_locations']
|
|
|
|
from ast import (Load, Store, Del,
|
|
Assign, AnnAssign, AugAssign,
|
|
Attribute, Subscript,
|
|
comprehension,
|
|
For, AsyncFor,
|
|
withitem,
|
|
Delete,
|
|
iter_child_nodes)
|
|
|
|
from .walker import Walker
|
|
|
|
try: # Python 3.8+
|
|
from ast import NamedExpr
|
|
except ImportError:
|
|
class _NoSuchNodeType:
|
|
pass
|
|
NamedExpr = _NoSuchNodeType
|
|
|
|
|
|
class _CtxFixer(Walker):
|
|
def __init__(self):
|
|
super().__init__(ctxclass=Load)
|
|
|
|
def transform(self, tree):
|
|
self._fix_one(tree)
|
|
self._setup_subtree_contexts(tree)
|
|
return self.generic_visit(tree)
|
|
|
|
def _fix_one(self, tree):
|
|
'''Fix one missing `ctx` attribute, using the currently active ctx class.'''
|
|
if ("ctx" in type(tree)._fields and (not hasattr(tree, "ctx") or tree.ctx is None)):
|
|
tree.ctx = self.state.ctxclass()
|
|
|
|
def _setup_subtree_contexts(self, tree):
|
|
'''Autoselect correct `ctx` class for subtrees of `tree`.'''
|
|
# The default ctx class is `Load`. We set up any `Store` and `Del`, as
|
|
# well as any `Load` for trees that may appear inside others that are
|
|
# set up as `Store` or `Del` (that mainly concerns expressions).
|
|
tt = type(tree)
|
|
if tt is Assign:
|
|
self.withstate(tree.targets, ctxclass=Store)
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
elif tt is AnnAssign:
|
|
self.withstate(tree.target, ctxclass=Store)
|
|
self.withstate(tree.annotation, ctxclass=Load)
|
|
if tree.value:
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
elif tt is NamedExpr:
|
|
self.withstate(tree.target, ctxclass=Store)
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
elif tt is AugAssign:
|
|
# `AugStore` and `AugLoad` are for internal use only, not even
|
|
# meant to be exposed to the user; the compiler expects `Store`
|
|
# and `Load` here. https://bugs.python.org/issue39988
|
|
self.withstate(tree.target, ctxclass=Store)
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
|
|
elif tt is Attribute:
|
|
# The tree's own `ctx` can be whatever, but `value` always has `Load`.
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
elif tt is Subscript:
|
|
# The tree's own `ctx` can be whatever, but `value` and `slice` always have `Load`.
|
|
self.withstate(tree.value, ctxclass=Load)
|
|
self.withstate(tree.slice, ctxclass=Load)
|
|
|
|
elif tt is comprehension:
|
|
self.withstate(tree.target, ctxclass=Store)
|
|
self.withstate(tree.iter, ctxclass=Load)
|
|
self.withstate(tree.ifs, ctxclass=Load)
|
|
|
|
elif tt in (For, AsyncFor):
|
|
self.withstate(tree.target, ctxclass=Store)
|
|
self.withstate(tree.iter, ctxclass=Load)
|
|
elif tt is withitem:
|
|
self.withstate(tree.context_expr, ctxclass=Load)
|
|
self.withstate(tree.optional_vars, ctxclass=Store)
|
|
|
|
elif tt is Delete:
|
|
self.withstate(tree.targets, ctxclass=Del)
|
|
|
|
|
|
def fix_missing_ctx(tree):
|
|
'''Fix any missing `ctx` attributes in `tree`.
|
|
|
|
Modifies `tree` in-place. For convenience, returns the modified `tree`.
|
|
'''
|
|
return _CtxFixer().visit(tree)
|
|
|
|
|
|
def fix_missing_locations(tree, reference_node, *, mode):
|
|
'''Like `ast.fix_missing_locations`, but customized for a macro expander.
|
|
|
|
Differences:
|
|
|
|
- If `reference_node` has no location info, return immediately (no-op).
|
|
- If `tree is None`, return immediately (no-op).
|
|
- If `tree` is a `list` of AST nodes, loop over it.
|
|
|
|
The `mode` parameter:
|
|
|
|
- If `mode="reference"`, populate any missing location info by
|
|
copying it from `reference_node`. Always use the same values.
|
|
|
|
Good for a macro expander.
|
|
|
|
- If `mode="update"`, behave exactly like `ast.fix_missing_locations`,
|
|
except that at the top level of `tree`, initialize `lineno` and
|
|
`col_offset` from `reference_node` (instead of using `1` and `0`
|
|
like `ast.fix_missing_locations` does).
|
|
|
|
So if a node is missing location info, copy the current reference info
|
|
in, but if it has location info, then update the reference info.
|
|
|
|
Good for general use.
|
|
|
|
- If `mode="overwrite"`, copy location info from `reference_node`,
|
|
regardless of if the target node already has it.
|
|
|
|
Good when `tree` is a code template that comes from another file,
|
|
so that any line numbers already in the AST would be misleading
|
|
at the use site.
|
|
|
|
Modifies `tree` in-place. For convenience, returns the modified `tree`.
|
|
'''
|
|
if not (hasattr(reference_node, "lineno") and hasattr(reference_node, "col_offset")):
|
|
return tree
|
|
def _fix(tree, lineno, col_offset):
|
|
if tree is None:
|
|
return
|
|
if isinstance(tree, list):
|
|
for elt in tree:
|
|
_fix(elt, lineno, col_offset)
|
|
return
|
|
if 'lineno' in tree._attributes:
|
|
if mode == "overwrite":
|
|
tree.lineno = lineno
|
|
else:
|
|
if not hasattr(tree, 'lineno'):
|
|
tree.lineno = lineno
|
|
elif mode == "update":
|
|
lineno = tree.lineno
|
|
if 'col_offset' in tree._attributes:
|
|
if mode == "overwrite":
|
|
tree.col_offset = col_offset
|
|
else:
|
|
if not hasattr(tree, 'col_offset'):
|
|
tree.col_offset = col_offset
|
|
elif mode == "update":
|
|
col_offset = tree.col_offset
|
|
for child in iter_child_nodes(tree):
|
|
_fix(child, lineno, col_offset)
|
|
_fix(tree, reference_node.lineno, reference_node.col_offset)
|
|
return tree
|