12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182 |
- """Filters used to touch up the not-quite-perfect ASTs that we allow macros
- to return."""
- from ast import *
- from macropy.core.util import register
- from macros import filters
- from walkers import Walker
- @register(filters)
- def fix_ctx(tree, **kw):
- return ast_ctx_fixer.recurse(tree, Load())
- @Walker
- def ast_ctx_fixer(tree, ctx, stop, set_ctx, **kw):
- """Fix any missing `ctx` attributes within an AST; allows you to build
- your ASTs without caring about that stuff and just filling it in later."""
- if "ctx" in type(tree)._fields and (not hasattr(tree, "ctx") or tree.ctx is None):
- tree.ctx = ctx
- if type(tree) is arguments:
- for arg in tree.args:
- ast_ctx_fixer.recurse(arg, Param())
- for default in tree.defaults:
- ast_ctx_fixer.recurse(default, Load())
- stop()
- return tree
- if type(tree) is AugAssign:
- ast_ctx_fixer.recurse(tree.target, AugStore())
- ast_ctx_fixer.recurse(tree.value, AugLoad())
- stop()
- return tree
- if type(tree) is Attribute:
- set_ctx(Load())
- return tree
- if type(tree) is Assign:
- for target in tree.targets:
- ast_ctx_fixer.recurse(target, Store())
- ast_ctx_fixer.recurse(tree.value, Load())
- stop()
- return tree
- if type(tree) is Delete:
- for target in tree.targets:
- ast_ctx_fixer.recurse(target, Del())
- stop()
- return tree
- @register(filters)
- def fill_line_numbers(tree, lineno, col_offset, **kw):
- """Fill in line numbers somewhat more cleverly than the
- ast.fix_missing_locations method, which doesn't take into account the
- fact that line numbers are monotonically increasing down lists of AST
- nodes."""
- if type(tree) is list:
- for sub in tree:
- if isinstance(sub, AST) \
- and hasattr(sub, "lineno") \
- and hasattr(sub, "col_offset") \
- and (sub.lineno, sub.col_offset) > (lineno, col_offset):
- lineno = sub.lineno
- col_offset = sub.col_offset
- fill_line_numbers(sub, lineno, col_offset)
- elif isinstance(tree, AST):
- if not (hasattr(tree, "lineno") and hasattr(tree, "col_offset")):
- tree.lineno = lineno
- tree.col_offset = col_offset
- for name, sub in iter_fields(tree):
- fill_line_numbers(sub, tree.lineno, tree.col_offset)
- return tree
|