Source

astdump / astdump.py

Full commit
#!/usr/bin/env python
"""
Module to play with Abstract Syntax Tree (AST). Can be
used to get version of another module (along with other
distutils metadata) without importing it.
"""

__author__ = 'anatoly techtonik <techtonik@gmail.com>'
__version__ = '1.0'
__license__ = 'Public Domain'
__description__ = 'Extract information from Python module without importing it.'


import ast


def propnames(node):
  """return names of attributes specific for the current node"""
  return [x for x in dir(node) if not x.startswith('_')
                                  and x not in ['col_offset', 'lineno']]

def dumpattrs(node, indent=0, oneline=False):
  """indented print of attributes of the given node
  """
  outlines = []
  for n in propnames(node):
    outlines.append("%s: %s" % (n, node.__dict__[n]))
  if not oneline:
    print " "*indent + ("\n"+" "*indent).join(outlines)
  else:
    print "[%s]" % ", ".join(outlines)


class TreeDumper(ast.NodeVisitor):
  def dump(self, node, types=[], level=None, callback=None):
    """pretty-print AST tree

       if `types` is set, process only types in the list
       if `level` is set, limit output to the given depth
       `callback` (if set) will be called to process filtered node
    """
    self.depth = 0
    self.types = types
    self.level = level
    self.callback = callback
    self.visit(node)

  def visit(self, node):
    """this function is called automatically"""
    nodetype = type(node)
    if not self.types or nodetype in self.types:
      if self.callback:
        self.callback(node, self.depth)
      else:
        print ' '*self.depth*2 + nodetype.__name__
    self.depth += 1
    if self.level == None or self.depth <= self.level:
      self.generic_visit(node)
    self.depth -= 1


# --- Callbacks ---
def printcb(node, level):
  nodename = node.__class__.__name__
  print ' '*(max(level,1)-1)*2 + nodename

def printassign(node, level):
  nodename = node.__class__.__name__
  print nodename
  dumpattrs(node, 1)
#--- /Callbacks ---


def top_level_vars(filename):
  """Return name/value pairs for top level variables for the script specified as `filename`.
     Only string and int values are supported.
  """
  root = read_ast(filename)
  return node_top_level_vars(root)

def node_top_level_vars(root):
  """
  Return dict with top level variables for the given node. Only string and
  int values are supported.
  """
  variables = {}
  def get_vars_cb(node, level):
    if type(node) != ast.Assign:
      return
    if level != 1:
      return
    if type(node.value) not in [ast.Str, ast.Num]:
      return
    for t in node.targets:  # all targets are of Name type
      #print " " + type(t).__name__
      #dumpattrs(t, 2)
      name = t.id
      #dumpattrs(node.value, 2)
      #print type(t), type(t) == ast.Name, type(t).__name__
      if type(node.value) == ast.Str:
        variables[name] = node.value.s
      elif type(node.value) == ast.Num:
        variables[name] = node.value.n

  TreeDumper().dump(root, types=[ast.Assign], level=1, callback=get_vars_cb)
  return variables


def read_ast(filename):
  """Read filename and return root node of the AST"""
  return ast.parse(open(filename).read(), filename)


def get_setup_py(filename):
  props = top_level_vars(filename)
  
  modname = filename[:-3] if filename.endswith(".py") else filename
  setup = """\
#!/usr/bin/env python
from distutils.core import setup

setup(
    name = '%s',
""" % modname

  if '__version__' in props:
    setup += "    version = '%s',\n" % props['__version__']
  author = props.get('__author__', None)
  if author:
    email = None
    if '@' in author and '<' in author:
      email = author.split('<', 1)[1].strip('>')
      author = author.split('<', 1)[0].strip()
    setup += "    author = '%s',\n" % author
    if email:
      setup += "    author_email = '%s',\n" % email
  for propname in ['__description__', '__license__']:
    if propname in props:
      setup += "    %s = '%s',\n" % (
                 propname.strip('_'), props[propname])

  setup += "\n    py_modules=['%s'],\n" % modname
  setup += ")\n"
  return setup


if __name__ == '__main__':
  import sys
  import optparse  

  parser = optparse.OptionParser(usage="%prog [options] <filename.py>",
             description="Read top level variables from the module without "
                         "importing it. Additional keys allow to generate "
                         "setup.py automatically or view tree structure of "
                         "an Abstract Syntax Tree.\n")
  parser.add_option('--dump', action='store_true',
                              help='dump abstract syntax tree')
  parser.add_option('--generate', action='store_true',
                              help='generate setup.py for a given filename')
  opts, args = parser.parse_args()
  
  if len(args) == 0:
    parser.print_help()
    sys.exit(-1)
  filename = args[0]

  if opts.dump:
    root = read_ast(filename)
    #print ast.dump(root, annotate_fields=False)
    TreeDumper().dump(root)
    sys.exit(0)

  if opts.generate:
    print get_setup_py(filename)
    sys.exit(0)

  #root = read_ast(filename)
  #TreeDumper().dump(root, types=[ast.Assign])
  #TreeDumper().dump(root, types=[ast.Assign], level=1)
  #TreeDumper().dump(root, types=[ast.Assign], level=1, callback=printcb)
  #TreeDumper().dump(root, types=[ast.Assign], level=1, callback=printassign)

  topvars = top_level_vars(filename)
  for name in sorted(topvars):
    print name, '=', repr(topvars[name])