"""
codetransformer.utils.pretty
----------------------------
Utilities for pretty-printing ASTs and code objects.
"""
from ast import iter_fields, AST, Name, Num, parse
import dis
from functools import partial, singledispatch
from io import StringIO
from itertools import chain
from operator import attrgetter
import sys
from types import CodeType
from codetransformer.code import Flag
INCLUDE_ATTRIBUTES_DEFAULT = False
INDENT_DEFAULT = ' '
__all__ = [
'a',
'd',
'display',
'pformat_ast',
'pprint_ast',
]
def _extend_name(prev, parent_co):
return prev + (
'.<locals>.' if parent_co.co_flags & Flag.CO_NEWLOCALS else '.'
)
[docs]def pprint_ast(node,
include_attributes=INCLUDE_ATTRIBUTES_DEFAULT,
indent=INDENT_DEFAULT,
file=None):
"""
Pretty-print an AST tree.
Parameters
----------
node : ast.AST
Top-level node to render.
include_attributes : bool, optional
Whether to include node attributes. Default False.
indent : str, optional.
Indentation string for nested expressions. Default is two spaces.
file : None or file-like object, optional
File to use to print output. If the default of `None` is passed, we
use sys.stdout.
"""
if file is None:
file = sys.stdout
print(
pformat_ast(
node,
include_attributes=include_attributes,
indent=indent
),
file=file,
)
def walk_code(co, _prefix=''):
"""
Traverse a code object, finding all consts which are also code objects.
Yields pairs of (name, code object).
"""
name = _prefix + co.co_name
yield name, co
yield from chain.from_iterable(
walk_code(c, _prefix=_extend_name(name, co))
for c in co.co_consts
if isinstance(c, CodeType)
)
def iter_attributes(node):
attrs = node._attributes
if not attrs:
return
yield from zip(attrs, attrgetter(*attrs)(node))
[docs]def a(text, mode='exec', indent=' ', file=None):
"""
Interactive convenience for displaying the AST of a code string.
Writes a pretty-formatted AST-tree to `file`.
Parameters
----------
text : str
Text of Python code to render as AST.
mode : {'exec', 'eval'}, optional
Mode for `ast.parse`. Default is 'exec'.
indent : str, optional
String to use for indenting nested expressions. Default is two spaces.
file : None or file-like object, optional
File to use to print output. If the default of `None` is passed, we
use sys.stdout.
"""
pprint_ast(parse(text, mode=mode), indent=indent, file=file)
[docs]def d(obj, mode='exec', file=None):
"""
Interactive convenience for displaying the disassembly of a function,
module, or code string.
Compiles `text` and recursively traverses the result looking for `code`
objects to render with `dis.dis`.
Parameters
----------
obj : str, CodeType, or object with __code__ attribute
Object to disassemble.
If `obj` is an instance of CodeType, we use it unchanged.
If `obj` is a string, we compile it with `mode` and then disassemble.
Otherwise, we look for a `__code__` attribute on `obj`.
mode : {'exec', 'eval'}, optional
Mode for `compile`. Default is 'exec'.
file : None or file-like object, optional
File to use to print output. If the default of `None` is passed, we
use sys.stdout.
"""
if file is None:
file = sys.stdout
for name, co in walk_code(extract_code(obj, compile_mode=mode)):
print(name, file=file)
print('-' * len(name), file=file)
dis.dis(co, file=file)
print('', file=file)
@singledispatch
def extract_code(obj, compile_mode):
"""
Generic function for converting objects into instances of `CodeType`.
"""
try:
code = obj.__code__
if isinstance(code, CodeType):
return code
raise ValueError(
"{obj} has a `__code__` attribute, "
"but it's an instance of {notcode!r}, not CodeType.".format(
obj=obj,
notcode=type(code).__name__,
)
)
except AttributeError:
raise ValueError("Don't know how to extract code from %s." % obj)
@extract_code.register(CodeType)
def _(obj, compile_mode):
return obj
@extract_code.register(str) # noqa
def _(obj, compile_mode):
return compile(obj, '<show>', compile_mode)
_DISPLAY_TEMPLATE = """\
====
Text
====
{text}
====================
Abstract Syntax Tree
====================
{ast}
===========
Disassembly
===========
{code}
"""
[docs]def display(text, mode='exec', file=None):
"""
Show `text`, rendered as AST and as Bytecode.
Parameters
----------
text : str
Text of Python code to render.
mode : {'exec', 'eval'}, optional
Mode for `ast.parse` and `compile`. Default is 'exec'.
file : None or file-like object, optional
File to use to print output. If the default of `None` is passed, we
use sys.stdout.
"""
if file is None:
file = sys.stdout
ast_section = StringIO()
a(text, mode=mode, file=ast_section)
code_section = StringIO()
d(text, mode=mode, file=code_section)
rendered = _DISPLAY_TEMPLATE.format(
text=text,
ast=ast_section.getvalue(),
code=code_section.getvalue(),
)
print(rendered, file=file)