cleanup.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. """Filters used to touch up the not-quite-perfect ASTs that we allow macros
  2. to return."""
  3. from ast import *
  4. from macropy.core.util import register
  5. from macros import filters
  6. from walkers import Walker
  7. @register(filters)
  8. def fix_ctx(tree, **kw):
  9. return ast_ctx_fixer.recurse(tree, Load())
  10. @Walker
  11. def ast_ctx_fixer(tree, ctx, stop, set_ctx, **kw):
  12. """Fix any missing `ctx` attributes within an AST; allows you to build
  13. your ASTs without caring about that stuff and just filling it in later."""
  14. if "ctx" in type(tree)._fields and (not hasattr(tree, "ctx") or tree.ctx is None):
  15. tree.ctx = ctx
  16. if type(tree) is arguments:
  17. for arg in tree.args:
  18. ast_ctx_fixer.recurse(arg, Param())
  19. for default in tree.defaults:
  20. ast_ctx_fixer.recurse(default, Load())
  21. stop()
  22. return tree
  23. if type(tree) is AugAssign:
  24. ast_ctx_fixer.recurse(tree.target, AugStore())
  25. ast_ctx_fixer.recurse(tree.value, AugLoad())
  26. stop()
  27. return tree
  28. if type(tree) is Attribute:
  29. set_ctx(Load())
  30. return tree
  31. if type(tree) is Assign:
  32. for target in tree.targets:
  33. ast_ctx_fixer.recurse(target, Store())
  34. ast_ctx_fixer.recurse(tree.value, Load())
  35. stop()
  36. return tree
  37. if type(tree) is Delete:
  38. for target in tree.targets:
  39. ast_ctx_fixer.recurse(target, Del())
  40. stop()
  41. return tree
  42. @register(filters)
  43. def fill_line_numbers(tree, lineno, col_offset, **kw):
  44. """Fill in line numbers somewhat more cleverly than the
  45. ast.fix_missing_locations method, which doesn't take into account the
  46. fact that line numbers are monotonically increasing down lists of AST
  47. nodes."""
  48. if type(tree) is list:
  49. for sub in tree:
  50. if isinstance(sub, AST) \
  51. and hasattr(sub, "lineno") \
  52. and hasattr(sub, "col_offset") \
  53. and (sub.lineno, sub.col_offset) > (lineno, col_offset):
  54. lineno = sub.lineno
  55. col_offset = sub.col_offset
  56. fill_line_numbers(sub, lineno, col_offset)
  57. elif isinstance(tree, AST):
  58. if not (hasattr(tree, "lineno") and hasattr(tree, "col_offset")):
  59. tree.lineno = lineno
  60. tree.col_offset = col_offset
  61. for name, sub in iter_fields(tree):
  62. fill_line_numbers(sub, tree.lineno, tree.col_offset)
  63. return tree