Transform as much of a python ast as possible into python literals

astpp.py

"""
A pretty-printing dump function for the ast module.  The code was copied from
the ast.dump function and modified slightly to pretty-print.

Alex Leone (acleone ~AT~ gmail.com), 2010-01-30
"""

from ast import *

def dump(node, annotate_fields=True, include_attributes=False, indent='  '):
    """
    Return a formatted dump of the tree in *node*.  This is mainly useful for
    debugging purposes.  The returned string will show the names and the values
    for fields.  This makes the code impossible to evaluate, so if evaluation is
    wanted *annotate_fields* must be set to False.  Attributes such as line
    numbers and column offsets are not dumped by default.  If this is wanted,
    *include_attributes* can be set to True.
    """
    def _format(node, level=0):
        if isinstance(node, AST):
            fields = [(a, _format(b, level)) for a, b in iter_fields(node)]
            if include_attributes and node._attributes:
                fields.extend([(a, _format(getattr(node, a), level))
                               for a in node._attributes])
            return ''.join([
                node.__class__.__name__,
                '(',
                ', '.join(('%s=%s' % field for field in fields)
                           if annotate_fields else
                           (b for a, b in fields)),
                ')'])
        elif isinstance(node, list):
            lines = ['[']
            lines.extend((indent * (level + 2) + _format(x, level + 2) + ','
                         for x in node))
            if len(lines) > 1:
                lines.append(indent * (level + 1) + ']')
            else:
                lines[-1] += ']'
            return '\n'.join(lines)
        return repr(node)
    if not isinstance(node, AST):
        raise TypeError('expected AST, got %r' % node.__class__.__name__)
    return _format(node)
    
if __name__ == '__main__':
    import sys
    for filename in sys.argv[1:]:
        print '=' * 50
        print 'AST tree for', filename
        print '=' * 50
        f = open(filename, 'r')
        fstr = f.read()
        f.close()
        print dump(parse(fstr, filename=filename), include_attributes=True)
        print

mako-setup-snippet.py

setup(name='Mako',
      version=VERSION,
      description="A super-fast templating language that borrows the \
 best ideas from the existing templating languages.",
      long_description=readme,
      classifiers=[
      'Development Status :: 5 - Production/Stable',
      'Environment :: Web Environment',
      'Intended Audience :: Developers',
      'Programming Language :: Python',
      'Programming Language :: Python :: 3',
      "Programming Language :: Python :: Implementation :: CPython",
      "Programming Language :: Python :: Implementation :: PyPy",
      'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
      ],
      keywords='templates',
      author='Mike Bayer',
      author_email='mike@zzzcomputing.com',
      url='http://www.makotemplates.org/',
      license='MIT',
      packages=find_packages('.', exclude=['examples*', 'test*']),
      scripts=['scripts/mako-render'],
      tests_require=['nose >= 0.11'],
      test_suite="nose.collector",
      zip_safe=False,
      install_requires=install_requires,
      extras_require={'beaker': ['Beaker>=1.1']},
      entry_points="""
      [python.templating.engines]
      mako = mako.ext.turbogears:TGPlugin

      [pygments.lexers]
      mako = mako.ext.pygmentplugin:MakoLexer
      html+mako = mako.ext.pygmentplugin:MakoHtmlLexer
      xml+mako = mako.ext.pygmentplugin:MakoXmlLexer
      js+mako = mako.ext.pygmentplugin:MakoJavascriptLexer
      css+mako = mako.ext.pygmentplugin:MakoCssLexer

      [babel.extractors]
      mako = mako.ext.babelplugin:extract
      """
)

test-py-output.xt

{
    'author': 'Mike Bayer',
    'author_email': 'mike@zzzcomputing.com',
    'classifiers': [
        'Development Status :: 5 - Production/Stable',
        'Environment :: Web Environment',
        'Intended Audience :: Developers',
        'Programming Language :: Python',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: Implementation :: CPython',
        'Programming Language :: Python :: Implementation :: PyPy',
        'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
    ],
    'description': 'A super-fast templating language that borrows the  best ideas from the existing templating languages.',
    'entry_points': '\n      [python.templating.engines]\n      mako = mako.ext.turbogears:TGPlugin\n\n      [pygments.lexers]\n      mako = mako.ext.pygmentplugin:MakoLexer\n      html+mako = mako.ext.pygmentplugin:MakoHtmlLexer\n      xml+mako = mako.ext.pygmentplugin:MakoXmlLexer\n      js+mako = mako.ext.pygmentplugin:MakoJavascriptLexer\n      css+mako = mako.ext.pygmentplugin:MakoCssLexer\n\n      [babel.extractors]\n      mako = mako.ext.babelplugin:extract\n      ',
    'extras_require': {'beaker': ['Beaker>=1.1']},
    'install_requires': Name('install_requires', Load()),
    'keywords': 'templates',
    'license': 'MIT',
    'long_description': Name('readme', Load()),
    'name': 'Mako',
    'packages': Call(Name('find_packages', Load()), [
        '.',
      ], {'exclude': ['examples*', 'test*']}, None, None),
    'scripts': ['scripts/mako-render'],
    'test_suite': 'nose.collector',
    'tests_require': ['nose >= 0.11'],
    'url': 'http://www.makotemplates.org/',
    'version': Name('VERSION', Load()),
    'zip_safe': False,
}

test.py

#!/usr/bin/env python

import ast
import codecs
try:
    from pprintpp import pprint
except ImportError:
    from pprint import pprint


def gather_setup_info(fileobj):
    parsed = ast.parse(fileobj.read(), fileobj.name)
    visitor = SetupScriptVisitor()
    visitor.visit(parsed)

    non_literals, extensions = {}, []
    for key, value in visitor.keywords.items():
        if key == 'ext_modules':
            if isinstance(value, list):
                for ext in value:
                    if  (isinstance(ext, ast.Call) and
                         isinstance(ext.func, ast.Name) and
                         ext.func.id == 'Extension' and
                         not has_non_literals(ext.args)):
                        extensions.append(ext.args[0])
        elif has_non_literals(value):
            non_literals[key] = value
            del visitor.keywords[key]

    return visitor.keywords, visitor.imported_modules, non_literals, extensions


class SetupScriptVisitor(ast.NodeVisitor):
    def __init__(self):
        ast.NodeVisitor.__init__(self)
        self.keywords = {}
        self.non_literals = []
        self.imported_modules = set()

    def visit_Expr(self, node):
        if isinstance(node.value, ast.Call) and \
           isinstance(node.value.func, ast.Name) and \
           node.value.func.id == 'setup':
            self.visit_setup(node.value)

    def visit_setup(self, node):
        call = LiteralAstTransform().visit(node)
        self.keywords = call.keywords
        for k, v in self.keywords.iteritems():
            if has_non_literals(v):
               self.non_literals.append(k)

    def visit_Import(self, node):
        for alias in node.names:
            self.imported_modules.add(alias.name)

    def visit_ImportFrom(self, node):
        self.imported_modules.add(node.module)


class LiteralAstTransform(ast.NodeTransformer):
    """Simplify the ast through evaluation of literals."""
    excluded_fields = ['ctx']

    def visit(self, node):
        if not isinstance(node, ast.AST):
            return node
        else:
            return ast.NodeTransformer.visit(self, node)

    def generic_visit(self, node):
        try:
            return ast.literal_eval(node)
        except ValueError:
            for field, value in ast.iter_fields(node):
                if field in self.excluded_fields:
                    delattr(node, field)
                if value is None:
                    continue

                if isinstance(value, list):
                    if field in ('keywords', 'kwargs'):
                        new_value = dict((kw.arg, self.visit(kw.value)) for kw in value)
                    else:
                        new_value = [self.visit(i) for i in value]
                else:
                    new_value = self.visit(value)
                setattr(node, field, new_value)
            return node

    def visit_Name(self, node):
        if hasattr('__builtins__', node.id):
            return getattr(__builtins__, node.id)
        else:
            return self.generic_visit(node)

    def visit_Tuple(self, node):
        return tuple(self.visit(v) for v in node.elts)

    def visit_List(self, node):
        return [self.visit(v) for v in node.elts]

    def visit_Set(self, node):
        return set(self.visit(v) for v in node.elts)

    def visit_Dict(self, node):
        keys = (self.visit(k) for k in node.keys)
        values = (self.visit(v) for v in node.values)
        return dict(zip(keys, values))


def has_non_literals(value):
    if isinstance(value, ast.AST):
        return True
    elif isinstance(value, basestring):
        return False
    elif hasattr(value, 'itervalues'):
        return any(has_non_literals(v) for v in value.itervalues())
    elif hasattr(value, '__iter__'):
        return any(has_non_literals(v) for v in value)


class Wrap(object):
    def __init__(self, astnode):
        self.astnode = astnode

    def __repr__(self):
        import astpp
        return astpp.dump(self.astnode, annotate_fields=False)


with codecs.open('mako-setup-snippet.py') as f:
    parsed = ast.parse(f.read(), f.name)
    visitor = SetupScriptVisitor()
    visitor.visit(parsed)
    data = dict(visitor.keywords)
    for key, value in data.iteritems():
        if isinstance(value, ast.AST):
            data[key] = Wrap(value)
    pprint(data)