ast.dump implementation using list of generated strings

ast_compare.py

import ast
from _ast import *

def dump(node, annotate_fields=True, include_attributes=False, indent=None):

    def _format(node, level=0):
        if indent is not None:
            level += 1
            prefix = '\n' + indent * level
            sep = ',\n' + indent * level
        else:
            prefix = ''
            sep = ', '
        if isinstance(node, AST):
            args = []
            keywords = annotate_fields
            for field in node._fields:
                try:
                    value = getattr(node, field)
                except AttributeError:
                    keywords = True
                else:
                    if keywords:
                        args.append('%s=%s' % (field, _format(value, level)))
                    else:
                        args.append(_format(value, level))
            if include_attributes and node._attributes:
                for a in node._attributes:
                    try:
                        args.append('%s=%s' % (a, _format(getattr(node, a), level)))
                    except AttributeError:
                        pass
            if not args:
                return '%s()' % (node.__class__.__name__,)
            return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args))
        elif isinstance(node, list):
            if not node:
                return '[]'
            return '[%s%s]' % (prefix, sep.join(_format(x, level) for x in node))
        return repr(node)
    if not isinstance(node, AST):
        raise TypeError('expected AST, got %r' % node.__class__.__name__)

    if indent is not None and not isinstance(indent, str):
        indent = ' ' * indent
    return _format(node)

def dump_cached(node, annotate_fields=True, include_attributes=False, indent=None):

    def _format(node, level=0):
        if indent is not None:
            level += 1
            try:
                prefix = indentations[level]
            except IndexError:
                child_indentation = indentations[level - 1] + indent
                prefix = '\n' + indent * level
                indentations.append(prefix)
            sep = ',\n' + indent * level
        else:
            prefix = ''
            sep = ', '
        if isinstance(node, AST):
            args = []
            keywords = annotate_fields
            for field in node._fields:
                try:
                    value = getattr(node, field)
                except AttributeError:
                    keywords = True
                else:
                    if keywords:
                        args.append('%s=%s' % (field, _format(value, level)))
                    else:
                        args.append(_format(value, level))
            if include_attributes and node._attributes:
                for a in node._attributes:
                    try:
                        args.append('%s=%s' % (a, _format(getattr(node, a), level)))
                    except AttributeError:
                        pass
            if not args:
                return '%s()' % (node.__class__.__name__,)
            return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args))
        elif isinstance(node, list):
            if not node:
                return '[]'
            return '[%s%s]' % (prefix, sep.join(_format(x, level) for x in node))
        return repr(node)
    if not isinstance(node, AST):
        raise TypeError('expected AST, got %r' % node.__class__.__name__)

    if indent is not None:
        if not isinstance(indent, str):
            indent = ' ' * indent
        indentations = ["\n" + indent]
    return _format(node)


obj = ast.parse("1 + 1 == 2" * 100)
assert dump(obj) == dump_cached(obj)