180 lines
6.9 KiB
Python
180 lines
6.9 KiB
Python
# -*- coding: utf-8; -*-
|
|
"""Fix `ctx` attributes and source location info in an AST."""
|
|
|
|
__all__ = ["fix_ctx", "fix_locations"]
|
|
|
|
from ast import (Load, Store, Del,
|
|
Assign, AnnAssign, AugAssign,
|
|
Attribute, Subscript,
|
|
comprehension,
|
|
For, AsyncFor,
|
|
withitem,
|
|
Delete,
|
|
iter_child_nodes)
|
|
from copy import copy
|
|
|
|
from . import walkers
|
|
|
|
try: # Python 3.8+
|
|
from ast import NamedExpr
|
|
except ImportError:
|
|
class _NoSuchNodeType:
|
|
pass
|
|
NamedExpr = _NoSuchNodeType
|
|
|
|
|
|
class _CtxFixer(walkers.ASTTransformer):
|
|
def __init__(self, *, copy_seen_nodes):
|
|
super().__init__(ctxclass=Load)
|
|
self.copy_seen_nodes = copy_seen_nodes
|
|
|
|
def reset(self, **bindings):
|
|
super().reset(**bindings)
|
|
self._seen = set()
|
|
|
|
def transform(self, tree):
|
|
tree = self._fix_one(tree)
|
|
self._setup_subtree_contexts(tree)
|
|
return self.generic_visit(tree)
|
|
|
|
def _fix_one(self, tree):
|
|
"""Fix one `ctx` attribute, using the currently active ctx class."""
|
|
if "ctx" in type(tree)._fields:
|
|
if self.copy_seen_nodes:
|
|
# Shallow-copy `tree` if already seen. This mode is used in the
|
|
# global postprocess pass. Note only nodes whose `ctx` we have
|
|
# already updated are considered seen.
|
|
if id(tree) in self._seen:
|
|
tree = copy(tree)
|
|
tree.ctx = self.state.ctxclass()
|
|
self._seen.add(id(tree))
|
|
return tree
|
|
|
|
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_ctx(tree, *, copy_seen_nodes):
|
|
"""Fix `ctx` attributes in `tree`.
|
|
|
|
If `copy_seen_nodes=True`, then, if the same AST node instance appears
|
|
multiple times, and the node requires a `ctx`, shallow-copy it the second
|
|
and further times it is encountered. This prevents problems if the same
|
|
node instance has been spliced into two or more positions that require
|
|
different `ctx`.
|
|
|
|
Modifies `tree` in-place. For convenience, returns the modified `tree`.
|
|
"""
|
|
return _CtxFixer(copy_seen_nodes=copy_seen_nodes).visit(tree)
|
|
|
|
|
|
def fix_locations(tree, reference_node, *, mode):
|
|
"""Like `ast.fix_missing_locations`, but customized for a macro expander.
|
|
|
|
Differences:
|
|
|
|
- If `reference_node` has source 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 reference info.
|
|
|
|
Good when expanding a macro invocation, to set the source location
|
|
of any macro-generated nodes to that of the macro invocation node.
|
|
|
|
- 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 (because they point to lines in that other file).
|
|
|
|
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
|