In [1]:
import tokenize
from StringIO import StringIO

In [2]:
a = "x = 3"
b = "y = x + 52.817"
hw = "print 'Hello, World!'"

In [3]:
def print_token(token_type, token_str, (srow, scol), (erow, ecol), line):
    if srow == erow:
        print "{}: {}, {}, {}".format(srow, tokenize.tok_name[token_type], repr(token_str), line)
    else:
        print "{}-{}: {}, {}, {}".format(srow, erow, tokenize.tok_name[token_type], repr(token_str), repr(line))

def print_token_verbose(token_type, token_str, (srow, scol), (erow, ecol), line):
    print "({},{})--({},{}), {}, {}, {}".format(
        srow, scol, erow, ecol, tokenize.tok_name[token_type], repr(token_str), repr(line))

def display_tokens(string, verbose=False):
    s = StringIO(string).readline
    f = print_token
    if verbose:
        f = print_token_verbose
    tokenize.tokenize(s, f)

In [4]:
display_tokens(a, verbose=True)


(1,0)--(1,1), NAME, 'x', 'x = 3'
(1,2)--(1,3), OP, '=', 'x = 3'
(1,4)--(1,5), NUMBER, '3', 'x = 3'
(2,0)--(2,0), ENDMARKER, '', ''

In [5]:
display_tokens(b, verbose=True)


(1,0)--(1,1), NAME, 'y', 'y = x + 52.817'
(1,2)--(1,3), OP, '=', 'y = x + 52.817'
(1,4)--(1,5), NAME, 'x', 'y = x + 52.817'
(1,6)--(1,7), OP, '+', 'y = x + 52.817'
(1,8)--(1,14), NUMBER, '52.817', 'y = x + 52.817'
(2,0)--(2,0), ENDMARKER, '', ''

In [6]:
display_tokens(hw, verbose=True)


(1,0)--(1,5), NAME, 'print', "print 'Hello, World!'"
(1,6)--(1,21), STRING, "'Hello, World!'", "print 'Hello, World!'"
(2,0)--(2,0), ENDMARKER, '', ''

In [7]:
longer="""a = 3
b = a**4
b - a == 78"""

display_tokens(longer, verbose=True)


(1,0)--(1,1), NAME, 'a', 'a = 3\n'
(1,2)--(1,3), OP, '=', 'a = 3\n'
(1,4)--(1,5), NUMBER, '3', 'a = 3\n'
(1,5)--(1,6), NEWLINE, '\n', 'a = 3\n'
(2,0)--(2,1), NAME, 'b', 'b = a**4\n'
(2,2)--(2,3), OP, '=', 'b = a**4\n'
(2,4)--(2,5), NAME, 'a', 'b = a**4\n'
(2,5)--(2,7), OP, '**', 'b = a**4\n'
(2,7)--(2,8), NUMBER, '4', 'b = a**4\n'
(2,8)--(2,9), NEWLINE, '\n', 'b = a**4\n'
(3,0)--(3,1), NAME, 'b', 'b - a == 78'
(3,2)--(3,3), OP, '-', 'b - a == 78'
(3,4)--(3,5), NAME, 'a', 'b - a == 78'
(3,6)--(3,8), OP, '==', 'b - a == 78'
(3,9)--(3,11), NUMBER, '78', 'b - a == 78'
(4,0)--(4,0), ENDMARKER, '', ''

In [8]:
class_ex="""class Foo(object):
    def __init__(self, x):
        self.x = x
f = Foo(3)
print f.x"""

display_tokens(class_ex, verbose=True)


(1,0)--(1,5), NAME, 'class', 'class Foo(object):\n'
(1,6)--(1,9), NAME, 'Foo', 'class Foo(object):\n'
(1,9)--(1,10), OP, '(', 'class Foo(object):\n'
(1,10)--(1,16), NAME, 'object', 'class Foo(object):\n'
(1,16)--(1,17), OP, ')', 'class Foo(object):\n'
(1,17)--(1,18), OP, ':', 'class Foo(object):\n'
(1,18)--(1,19), NEWLINE, '\n', 'class Foo(object):\n'
(2,0)--(2,4), INDENT, '    ', '    def __init__(self, x):\n'
(2,4)--(2,7), NAME, 'def', '    def __init__(self, x):\n'
(2,8)--(2,16), NAME, '__init__', '    def __init__(self, x):\n'
(2,16)--(2,17), OP, '(', '    def __init__(self, x):\n'
(2,17)--(2,21), NAME, 'self', '    def __init__(self, x):\n'
(2,21)--(2,22), OP, ',', '    def __init__(self, x):\n'
(2,23)--(2,24), NAME, 'x', '    def __init__(self, x):\n'
(2,24)--(2,25), OP, ')', '    def __init__(self, x):\n'
(2,25)--(2,26), OP, ':', '    def __init__(self, x):\n'
(2,26)--(2,27), NEWLINE, '\n', '    def __init__(self, x):\n'
(3,0)--(3,8), INDENT, '        ', '        self.x = x\n'
(3,8)--(3,12), NAME, 'self', '        self.x = x\n'
(3,12)--(3,13), OP, '.', '        self.x = x\n'
(3,13)--(3,14), NAME, 'x', '        self.x = x\n'
(3,15)--(3,16), OP, '=', '        self.x = x\n'
(3,17)--(3,18), NAME, 'x', '        self.x = x\n'
(3,18)--(3,19), NEWLINE, '\n', '        self.x = x\n'
(4,0)--(4,0), DEDENT, '', 'f = Foo(3)\n'
(4,0)--(4,0), DEDENT, '', 'f = Foo(3)\n'
(4,0)--(4,1), NAME, 'f', 'f = Foo(3)\n'
(4,2)--(4,3), OP, '=', 'f = Foo(3)\n'
(4,4)--(4,7), NAME, 'Foo', 'f = Foo(3)\n'
(4,7)--(4,8), OP, '(', 'f = Foo(3)\n'
(4,8)--(4,9), NUMBER, '3', 'f = Foo(3)\n'
(4,9)--(4,10), OP, ')', 'f = Foo(3)\n'
(4,10)--(4,11), NEWLINE, '\n', 'f = Foo(3)\n'
(5,0)--(5,5), NAME, 'print', 'print f.x'
(5,6)--(5,7), NAME, 'f', 'print f.x'
(5,7)--(5,8), OP, '.', 'print f.x'
(5,8)--(5,9), NAME, 'x', 'print f.x'
(6,0)--(6,0), ENDMARKER, '', ''

In [9]:
import ast
import copy

In [10]:
ast.dump(ast.parse(a))


Out[10]:
"Module(body=[Assign(targets=[Name(id='x', ctx=Store())], value=Num(n=3))])"

In [11]:
exec(compile(ast.parse(hw, mode='exec'), '', 'exec'))


Hello, World!

In [12]:
ast.dump(ast.parse(a), include_attributes=True)


Out[12]:
"Module(body=[Assign(targets=[Name(id='x', ctx=Store(), lineno=1, col_offset=0)], value=Num(n=3, lineno=1, col_offset=4), lineno=1, col_offset=0)])"

In [13]:
print ast.dump(ast.parse(a))


Module(body=[Assign(targets=[Name(id='x', ctx=Store())], value=Num(n=3))])

In [14]:
# http://05d2db1380b6504cc981-8cbed8cf7e3a131cd8f1c3e383d10041.r93.cf2.rackcdn.com/pycon-us-2011/419_what-would-you-do-with-an-ast.mp4

use_x = ast.Name(id="x", ctx=ast.Load())
assign_x = ast.Name(id="x", ctx=ast.Store())
one = ast.Num(n=1)
two = ast.Num(n=2)
four = ast.Num(n=4)

In [15]:
assign_2 = ast.Assign(targets=[copy.copy(assign_x)], value=copy.copy(two))
assign_2.lineno = 1
assign_2.col_offset = 0
ast.fix_missing_locations(assign_2) # <-- Applies lineno and col_offset to nodes in subtree (wrong, but ok for now)


Out[15]:
<_ast.Assign at 0x10cdb3d90>

In [16]:
if_ = ast.If()
if_.test = ast.Compare()
if_.test.left = copy.copy(use_x)
if_.test.ops = [ast.Eq()]
if_.test.comparators = [copy.copy(four)]

In [17]:
if_.body = [ast.Print(dest=None, values=[copy.copy(use_x)], nl=True)]

In [18]:
x_plus_1 = ast.BinOp(left=copy.copy(use_x), op=ast.Add(), right=copy.copy(one))
assign = ast.Assign(targets=[copy.copy(assign_x)], value=x_plus_1)
if_.orelse = [assign]

In [19]:
if_.lineno = 2
if_.col_offset = 0
ast.fix_missing_locations(if_)


Out[19]:
<_ast.If at 0x10cdd8250>

In [20]:
print_x = ast.Print(dest=None, values=[copy.copy(use_x)], nl=True)

In [21]:
print_x.lineno = 3
print_x.col_offset = 0
ast.fix_missing_locations(print_x)


Out[21]:
<_ast.Print at 0x10cdb3b10>

In [22]:
mod = ast.Module(body=[assign_2, if_, print_x])

In [23]:
code = compile(mod, '', 'exec')

In [24]:
exec code
print
print (ast.dump(mod, include_attributes=True))
print


3

Module(body=[Assign(targets=[Name(id='x', ctx=Store(), lineno=1, col_offset=0)], value=Num(n=2, lineno=1, col_offset=0), lineno=1, col_offset=0), If(test=Compare(left=Name(id='x', ctx=Load(), lineno=2, col_offset=0), ops=[Eq()], comparators=[Num(n=4, lineno=2, col_offset=0)], lineno=2, col_offset=0), body=[Print(dest=None, values=[Name(id='x', ctx=Load(), lineno=2, col_offset=0)], nl=True, lineno=2, col_offset=0)], orelse=[Assign(targets=[Name(id='x', ctx=Store(), lineno=2, col_offset=0)], value=BinOp(left=Name(id='x', ctx=Load(), lineno=2, col_offset=0), op=Add(), right=Num(n=1, lineno=2, col_offset=0), lineno=2, col_offset=0), lineno=2, col_offset=0)], lineno=2, col_offset=0), Print(dest=None, values=[Name(id='x', ctx=Load(), lineno=3, col_offset=0)], nl=True, lineno=3, col_offset=0)])


In [25]:
fcn_str = """def f(a, b=3, c=None):
    print a
    print "b={}, c={}".format(b,c)"""

In [26]:
fcn_parsed = ast.parse(fcn_str)
print ast.dump(fcn_parsed)


Module(body=[FunctionDef(name='f', args=arguments(args=[Name(id='a', ctx=Param()), Name(id='b', ctx=Param()), Name(id='c', ctx=Param())], vararg=None, kwarg=None, defaults=[Num(n=3), Name(id='None', ctx=Load())]), body=[Print(dest=None, values=[Name(id='a', ctx=Load())], nl=True), Print(dest=None, values=[Call(func=Attribute(value=Str(s='b={}, c={}'), attr='format', ctx=Load()), args=[Name(id='b', ctx=Load()), Name(id='c', ctx=Load())], keywords=[], starargs=None, kwargs=None)], nl=True)], decorator_list=[])])

In [27]:
for arg in fcn_parsed.body[0].args.args:
    print arg.id


a
b
c

In [28]:
ast.dump(fcn_parsed.body[0].args)


Out[28]:
"arguments(args=[Name(id='a', ctx=Param()), Name(id='b', ctx=Param()), Name(id='c', ctx=Param())], vararg=None, kwarg=None, defaults=[Num(n=3), Name(id='None', ctx=Load())])"

In [29]:
for x in fcn_parsed.body[0].args.defaults:
    if isinstance(x, ast.Name):
        print x.id
    elif isinstance(x, ast.Num):
        print x.n
    else:
        print x


3
None

In [30]:
call = "foo(a, b, x=3, y='meh')"
ast.dump(ast.parse(call))


Out[30]:
"Module(body=[Expr(value=Call(func=Name(id='foo', ctx=Load()), args=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())], keywords=[keyword(arg='x', value=Num(n=3)), keyword(arg='y', value=Str(s='meh'))], starargs=None, kwargs=None))])"

In [31]:
ast.dump(ast.parse("foo.bar(a,b,x=y)"))


Out[31]:
"Module(body=[Expr(value=Call(func=Attribute(value=Name(id='foo', ctx=Load()), attr='bar', ctx=Load()), args=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())], keywords=[keyword(arg='x', value=Name(id='y', ctx=Load()))], starargs=None, kwargs=None))])"

In [32]:
ast.dump(ast.parse("""
import foo
foo.bar(a,b,x=y)"""))


Out[32]:
"Module(body=[Import(names=[alias(name='foo', asname=None)]), Expr(value=Call(func=Attribute(value=Name(id='foo', ctx=Load()), attr='bar', ctx=Load()), args=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())], keywords=[keyword(arg='x', value=Name(id='y', ctx=Load()))], starargs=None, kwargs=None))])"

In [36]:
ast.dump(ast.parse("""kw = {"b": 3}
foo(2, [3,4,"bah"], **kw)"""))


Out[36]:
"Module(body=[Assign(targets=[Name(id='kw', ctx=Store())], value=Dict(keys=[Str(s='b')], values=[Num(n=3)])), Expr(value=Call(func=Name(id='foo', ctx=Load()), args=[Num(n=2), List(elts=[Num(n=3), Num(n=4), Str(s='bah')], ctx=Load())], keywords=[], starargs=None, kwargs=Name(id='kw', ctx=Load())))])"

In [69]:
x = ast.parse("foo(2, 3, a=6)")
print ast.dump(x.body[0].value)
x.body[0].value.func.id


Call(func=Name(id='foo', ctx=Load()), args=[Num(n=2), Num(n=3)], keywords=[keyword(arg='a', value=Num(n=6))], starargs=None, kwargs=None)
Out[69]:
'foo'

In [38]:
with open("../pip-create/setup.py") as f:
    setup = f.read()

In [93]:
functions = {}

def grab_crap(x):
    if isinstance(x, ast.Str):
        return x.s
    elif isinstance(x, ast.Name):
        return x.id
    elif isinstance(x, ast.List):
        return [grab_crap(y) for y in x.elts]
    elif isinstance(x, ast.Dict):
        return {grab_crap(k): grab_crap(v) for k,v in zip(x.keys, x.values)}
    else:
        return x

class FunctionGrabber(ast.NodeVisitor):
    def visit_Call(self, node):
        
        if isinstance(node.func, ast.Name) and node.func.id == "setup":
            #print ast.dump(node)
            
            gah = {kw.arg: grab_crap(kw.value) for kw in node.keywords}
            
            print gah

In [94]:
FunctionGrabber().visit(ast.parse(setup))


{'entry_points': {'console_scripts': ['pip-create=pip_create:main']}, 'name': 'pip-create', 'license': 'MIT', 'author': 'Jack Maney', 'author_email': 'jackmaney@gmail.com', 'include_package_data': 'True', 'classifiers': ['Programming Language :: Python :: 2.7'], 'url': 'https://github.com/jackmaney/pip-create', 'version': '__version__', 'install_requires': 'requirements', 'packages': <_ast.Call object at 0x10d2a0e50>, 'long_description': 'long_description', 'description': 'Tools to create a Python distribution for use with pip'}

In [72]:
functions


Out[72]:
{}

In [89]:
zip(["a", "b"], [1, 2])


Out[89]:
[('a', 1), ('b', 2)]

In [ ]: