We Keep Coding

ast python how to find connections between methods at different files

I want to index all methods and thee connections between them in an entire application (A directory with sub directories and files eventually). I'm using ast, looping over directories till individual files and then loads them into an ast object like so ast.parse(self.file_content)
The index that i'm trying to create is this
connection
Here is my code, if it's relevant.
def scan(self):
'''
scans a file line by line while keeping context of location and classes
indexes a file buffer content into a ast, Abstract Syntax Trees, https://en.wikipedia.org/wiki/AST.
Then, iterate over the elements, find relevant ones to index, and index them into the db.
'''
parsed_result = ast.parse(self.file_content)
for element in parsed_result.body:
results = self.index_element(element)
def index_element(self, element, class_name=None):
'''
if element is relevant, meaning method -> index
if element is Class -> recursively call it self
:param element:
:param class_name:
:return: [{insert_result: <db_insert_result>, 'structured_data': <method> object}, ...]
'''
# find classes
# find methods inside classes
# find hanging functions
# validation on element type
if self.should_index_element(element):
if self.is_class_definition(element):
class_element = element
indexed_items = []
for inner_element in class_element.body:
# recursive call
results = self.index_element(inner_element, class_name=class_element.name)
indexed_items += results
return indexed_items
else:
structured_data = self.structure_method_into_an_object(element, class_name=class_name)
result_graph = self.dal_client.methods_graph.insert_or_update(structured_data)
return "WhatEver"
return "WhatEver"
My question is, is it possible to create this graph using ast. If yes, how?
From my understanding, I currently can't since I'm loading one file at a time to the ast object and it is not aware of outside methods.
here is an example for 2 files that I want to link between them:
sample_a.py
from sample_class_b import SampleClassB
sample_b = SampleClassB()
class SampleClassA(object):
def __init__(self):
self.a = 1
def test_call_to_another_function(self):
return sample_b.test()
sample_b.py
class SampleClassB(object):
def __init__(self):
self.b = 1
def test(self):
return True

You can traverse the ast.Ast tree and at each recursive call do one of four things:
If the tree is a class definition, store the class name with its associated methods, and then apply Connections.walk to each of the methods, storing the class and method name in the scope.
If the tree is an import statement, load the module and recursively run Connections.walk on it.
If an attribute lookup is being made and Connections.walk is within a method, check if the attribute name is a method of any classes currently loaded. If so, add an edge to edges that links the current scope with this new method discovered.
If none of the above occurs, continue to traverse the tree.
import ast, itertools
import re, importlib
class Connections:
def __init__(self):
self._classes, self.edges = {}, []
def walk(self, tree, scope=None):
t_obj = None
if isinstance(tree, ast.ClassDef):
self._classes[tree.name] = [i for i in tree.body if isinstance(i, ast.FunctionDef) and not re.findall('__[a-z]+__', i.name)]
_ = [self.walk(i, [tree.name, i.name]) for i in self._classes[tree.name]]
t_obj = [i for i in tree.body if i not in self._classes[tree.name]]
elif isinstance(tree, (ast.Import, ast.ImportFrom)):
for p in [tree.module] if hasattr(tree, 'module') else [i.name for i in tree.names]:
with open(importlib.import_module(p).__file__) as f:
t_obj = ast.parse(f.read())
elif isinstance(tree, ast.Attribute) and scope is not None:
if (c:=[a for a, b in self._classes.items() if any(i.name == tree.attr for i in b)]):
self.edges.append((scope, [c[0], tree.attr]))
t_obj = tree.value
if isinstance(t_obj:=(tree if t_obj is None else t_obj), list):
for i in t_obj:
self.walk(i, scope = scope)
else:
for i in getattr(t_obj, '_fields', []):
self.walk(getattr(t_obj, i), scope=scope)
with open('sample_a.py') as f:
c = Connections()
c.walk(ast.parse(f.read()))
print(c.edges)
Output:
[(['SampleClassA', 'test_call_to_another_function'], ['SampleClassB', 'test'])]
Important note: depending on the complexity of the files you are running Connections.walk on, a RecursionError might occur. To circumvent this, here is a Gist that contains an iterative version of Connections.walk.
Creating a graph from edges:
import networkx as nx
import matplotlib.pyplot as plt
g, labels, c1 = nx.DiGraph(), {}, itertools.count(1)
for m1, m2 in c.edges:
if (p1:='.'.join(m1)) not in labels:
labels[p1] = next(c1)
if (p2:='.'.join(m2)) not in labels:
labels[p2] = next(c1)
g.add_node(labels[p1])
g.add_node(labels[p2])
g.add_edge(labels[p1], labels[p2])
nx.draw(g, pos, labels={b:a for a, b in labels.items()}, with_labels = True)
plt.show()
Output:

Using AST to change function names from CamelCase to snake_case

Here in this question, I was asking for a way to convert function names from CamelCase to snake_case, one of the comments suggested using AST.
I found a code snippet to find all function calls in a script
import ast
from collections import deque
class FuncCallVisitor(ast.NodeVisitor):
def __init__(self):
self._name = deque()
#property
def name(self):
return '.'.join(self._name)
#name.deleter
def name(self):
self._name.clear()
def visit_Name(self, node):
self._name.appendleft(node.id)
def visit_Attribute(self, node):
try:
self._name.appendleft(node.attr)
self._name.appendleft(node.value.id)
except AttributeError:
self.generic_visit(node)
def get_func_calls(tree):
func_calls = []
for node in ast.walk(tree):
if isinstance(node, ast.Call):
callvisitor = FuncCallVisitor()
callvisitor.visit(node.func)
func_calls.append(callvisitor.name)
return func_calls
if __name__ == '__main__':
tree = ast.parse(open("some_dir").read())
print(get_func_calls(tree))
using this code I have all function calls in my script, now I want to write a code that converts this name to snake_case.
I found this code snippet to modify a node in AST tree
class RewriteName(ast.NodeTransformer):
def visit_Name(self, node):
return ast.copy_location(ast.Subscript(
value=ast.Name(id='data', ctx=ast.Load()),
slice=ast.Index(value=ast.Str(s=node.id)),
ctx=node.ctx
), node)
tree = RewriteName().visit(tree)
I didn't understand how to use it to serve my purpose. Any explanation or other pieces of advice?

I am kind of late to this, but maybe it will be found in the future.
Anyway, here is a quick hack at it. Actually, you were almost there with your solution. The name method returns your name, then you can arbitrarily change that. So in your def get_func_calls(tree) call you can manipulate the string and re-assign the new name to the Call object.
ccName = callvisitor.name # work with some local var
new_name = '' # the new func name
for char_i in range(len(ccName)): # go over the name
if ccName[char_i].isupper(): # check if the current char is with uppercase
if ccName[char_i - 1] == '.': # check if the previous character is a dot
new_name += ccName[char_i].lower() # if it is, make the char to lowercase
else:
new_name += '_' + ccName[char_i].lower() # otherwise add the snake_
else:
new_name += ccName[char_i] # just add the rest of the lower chars
callvisitor._name = new_name # just re-asign the new name
func_calls.append(callvisitor._name)
This is definitely not a pretty solution and it also depends if you want to change only function definitions or every single function call in a file, but this should give you an idea on how to change the ast.

Identifying pure functions in python

I have a decorator #pure that registers a function as pure, for example:
#pure
def rectangle_area(a,b):
return a*b
#pure
def triangle_area(a,b,c):
return ((a+(b+c))(c-(a-b))(c+(a-b))(a+(b-c)))**0.5/4
Next, I want to identify a newly defined pure function
def house_area(a,b,c):
return rectangle_area(a,b) + triangle_area(a,b,c)
Obviously house_area is pure, since it only calls pure functions.
How can I discover all pure functions automatically (perhaps by using ast)

Assuming operators are all pure, then essentially you only need to check all the functions calls. This can indeed be done with the ast module.
First I defined the pure decorator as:
def pure(f):
f.pure = True
return f
Adding an attribute telling that it's pure, allows skipping early or "forcing" a function to identify as pure. This is useful if you'd need a function like math.sin to identify as pure. Additionally since you can't add attributes to builtin functions.
#pure
def sin(x):
return math.sin(x)
All in all. Use the ast module to visit all the nodes. Then for each Call node check whether the function being called is pure.
import ast
class PureVisitor(ast.NodeVisitor):
def __init__(self, visited):
super().__init__()
self.pure = True
self.visited = visited
def visit_Name(self, node):
return node.id
def visit_Attribute(self, node):
name = [node.attr]
child = node.value
while child is not None:
if isinstance(child, ast.Attribute):
name.append(child.attr)
child = child.value
else:
name.append(child.id)
break
name = ".".join(reversed(name))
return name
def visit_Call(self, node):
if not self.pure:
return
name = self.visit(node.func)
if name not in self.visited:
self.visited.append(name)
try:
callee = eval(name)
if not is_pure(callee, self.visited):
self.pure = False
except NameError:
self.pure = False
Then check whether the function has the pure attribute. If not get code and check if all the functions calls can be classified as pure.
import inspect, textwrap
def is_pure(f, _visited=None):
try:
return f.pure
except AttributeError:
pass
try:
code = inspect.getsource(f.__code__)
except AttributeError:
return False
code = textwrap.dedent(code)
node = compile(code, "<unknown>", "exec", ast.PyCF_ONLY_AST)
if _visited is None:
_visited = []
visitor = PureVisitor(_visited)
visitor.visit(node)
return visitor.pure
Note that print(is_pure(lambda x: math.sin(x))) doesn't work since inspect.getsource(f.__code__) returns code on a line by line basis. So the source returned by getsource would include the print and is_pure call, thus yielding False. Unless those functions are overridden.
To verify that it works, test it by doing:
print(house_area) # Prints: True
To list through all the functions in the current module:
import sys, types
for k in dir(sys.modules[__name__]):
v = globals()[k]
if isinstance(v, types.FunctionType):
print(k, is_pure(v))
The visited list keeps track of which functions have already been verified pure. This help circumvent problems related to recursion. Since the code isn't executed, the evaluation would recursively visit factorial.
#pure
def factorial(n):
return 1 if n == 1 else n * factorial(n - 1)
Note that you might need to revise the following code. Choosing another way to obtain a function from its name.
try:
callee = eval(name)
if not is_pure(callee, self.visited):
self.pure = False
except NameError:
self.pure = False

Good way to count number of functions of a python file, given path

I do not want to import my module. I have to count the number of functions given the .py file path. What is the best way to do so?
One thing I thought about was to count the number of "def" in my code, but it does not seem like the best way to go about this. Is there any better way to count the number of functions?

To count the top level definition, use ast module like this:
import ast
with open(filename) as f:
tree = ast.parse(f.read())
sum(isinstance(exp, ast.FunctionDef) for exp in tree.body)

You can use ast.NodeVisitor:
import inspect
import importlib
import ast
class CountFunc(ast.NodeVisitor):
func_count = 0
def visit_FunctionDef(self, node):
self.func_count += 1
mod = "/path/to/some.py"
p = ast.parse(open(mod).read())
f = CountFunc()
f.visit(p)
print(f.func_count)
If you wanted to include lambdas you would need to add a visit_Lambda:
def visit_Lambda(self, node):
self.func_count += 1
That will find all defs including methods of any classes, we could add more restrictions to disallow that:
class CountFunc(ast.NodeVisitor):
func_count = 0
def visit_ClassDef(self, node):
return
def visit_FunctionDef(self, node):
self.func_count += 1
def visit_Lambda(self, node):
self.func_count += 1
You can tailor the code however you like, all the nodes and their attributes are described in the greentreesnakes docs

Padraic beats me to it, but here is my code, which works with Python 2 and Python 3.
from __future__ import print_function
import ast
class FunctionCounter(ast.NodeVisitor):
def __init__(self, filename):
self.function_count = 0
with open(filename) as f:
module = ast.parse(f.read())
self.visit(module)
def visit_FunctionDef(self, node):
print('function: {}'.format(node.name))
self.function_count += 1
# Uncomment this to disable counting methods, properties within a
# class
# def visit_ClassDef(self, node):
# pass
if __name__ == '__main__':
counter = FunctionCounter('simple.py')
print('Number of functions: {}'.format(counter.function_count))
Discussion
This code not only count the function at the module level, it also count functions (method and properties) nested within class definitions.
To disable counting functions in the class definitions, uncomment the 2 lines for visit_ClassDef

You can use
len(dir(module))
Hope it helps

you can use the pyclbr module to get results of the module, the only catch is that you need to use the name of the module as you would import it instead of the file path, this benefits from also recognizing from X import Y for python source based modules (not builtin ones like math)
from pyclbr import readmodule_ex, Function
#readmodule_ex is function 1
def test(): #2
pass
def other_func(): #3
pass
class Thing:
def method(self):
pass
result = readmodule_ex("test") #this would be it's own file if it is test.py
funcs = sum(isinstance(v,Function) for v in result.values())
print(funcs)

My answer is an edit to the answer of # Naveen Kumar because I can not edit his answer at the moment.
You can use
import module
len(dir(module))
i.e:
import math
print(len(dir(math)))
output:
63

How to get a list of all non imported names in a Python module?

Given a module containing :
import stuff
from foo import Foo
from bar import *
CST = True
def func(): pass
How can I define a function get_defined_objects so that I can do:
print(get_defined_objects('path.to.module'))
{'CST': True, 'func', <function path.to.module.func>}
Right now the only solution I can imagine is to read the original module file, extract defined names with re.search(r'^(?:def|class )?(\w+)(?:\s*=)?' then import the module, and find the intersection with __dict__.
Is there something cleaner ?

Here is something for you to start with using ast. Note that this code does not cover all possible cases, although it should handle e.g. multiple assignment properly. Consider investigating ast's data structures and API more closely if you would like to get access to compiled code, for example.
import ast
with open('module.py') as f:
data = f.read()
tree = ast.parse(data)
elements = [el for el in tree.body if type(el) in (ast.Assign, ast.FunctionDef, ast.ClassDef)]
result = {}
for el in elements:
if type(el) == ast.Assign:
for t in el.targets:
if type(el.value) == ast.Call:
result[t.id] = el.value.func.id + '()'
else:
for attr in ['id', 'i', 's']:
try:
result[t.id] = getattr(el.value, attr)
break
except Exception as e:
pass
elif type(el) == ast.FunctionDef:
result[el.name] = '<function %s>' % el.name
else:
result[el.name] = '<class %s>' % el.name
print result
#

mod = "foo"
import ast, inspect
import importlib
mod = importlib.import_module(mod)
p = ast.parse(inspect.getsource(mod))
from collections import defaultdict
data = defaultdict(defaultdict)
for node in p.body:
if isinstance(node, (ast.ImportFrom, ast.Import)):
continue
if isinstance(node, (ast.ClassDef, ast.FunctionDef)):
data["classes"][node.name] = mod.__dict__[node.name]
elif isinstance(node, ast.Assign):
for trg in node.targets:
if isinstance(node.value, ast.Num):
data["assignments"][trg.id] = node.value.n
elif isinstance(node.value, ast.Str):
data["assignments"][trg.id] = node.value.s
else:
data["assignments"][trg.id] = mod.__dict__[trg.id]
Output:
There is a nice explanation here that lists what the different types do and their attributes which this is based on:
class Nodes(ast.NodeVisitor):
def __init__(self):
self.data = defaultdict()
super(Nodes, self).__init__()
def visit_FunctionDef(self, node):
self.data[node.name] = mod.__dict__[node.name]
print("In FunctionDef with funcion {}".format(node.name))
def visit_ClassDef(self, node):
self.data[node.name] = mod.__dict__[node.name]
def visit_Assign(self, node):
for trg in node.targets:
if isinstance(node.value, (ast.Str, ast.Num, ast.Dict, ast.List, ast.ListComp, ast.NameConstant)):
self.data[trg.id] = mod.__dict__[trg.id]
self.generic_visit(node)
def visit_Name(self, node):
"""
class Name(idctx)
A variable name. id holds the name as a string
and ctx is either class Load class Store class Del.
"""
print("In Name with {}\n".format(node.id))
#
def visit_Dict(self, node):
"""
class Dict(keys, values)
A dictionary. keys and values
hold lists of nodes with matching order
"""
print("In Dict keys = {}, values = {}\n".format(node.keys,node.values))
def visit_Set(self,node):
"""
class Set(elts)
A set. elts holds a list of
nodes representing the elements.
"""
print("In Set elts = {}\n".format(node.elts))
def visit_List(self, node):
"""
class List(eltsctx)
lts holds a list of nodes representing the elements.
ctx is Store if the container
is an assignment target
(i.e. (x,y)=pt), and Load otherwise.
"""
print("In List elts = {}\nctx = {}\n".format(node.elts,node.ctx))
def visit_Tuple(self, node):
"""
class Tuple(eltsctx)
lts holds a list of nodes representing the elements.
ctx is Store if the container
is an assignment target
(i.e. (x,y)=pt), and Load otherwise.
"""
print("In Tuple elts = {}\nctx = {}\n".format(node.elts,node.ctx))
def visit_NameConstant(self, node):
"""
class NameConstant(value)
True, False or None. "value" holds one of those constants.
"""
print("In NameConstant getting value {}\n".format(node.value))
def visit_Load(self, node):
print("In Load with node {}\n".format(node.func))
def visit_Call(self, node):
"""
class Call(func, args, keywords, starargs, kwargs)
A function call. func is the function,
which will often be a Name or Attribute object. Of the arguments:
args holds a list of the arguments passed by position.
keywords holds a list of keyword objects representing arguments
passed by keyword.starargs and kwargs each hold a single node,
for arguments passed as *args and **kwargs.
"""
print("In Call with node {}\n".format(node.func))
def visit_Num(self, node):
print("In Num getting value {}\n".format(node.n))
def visit_Str(self, node):
print("In Str getting value {}\n".format(node.s))
f = Nodes()
f.visit(p)
print(f.data)

A bytecode hack for Python 3.4+. Possible due to dis.get_instructions.
import dis
import importlib
from itertools import islice
import marshal
import os
def consume_iterator(it, n=1):
next(islice(it, n, n), None)
def get_defined_names(module_path):
path, module_name = os.path.split(module_path)
module_name = module_name[:-3]
module_object = importlib.import_module(module_name)
pyc_name = '{}.cpython-34.pyc'.format(module_name)
pyc_path = os.path.join(path, '__pycache__/', pyc_name)
with open(pyc_path, 'rb') as f:
f.read(12) # drop the first 12 bytes
code = marshal.load(f)
# dis.disassemble(code) # see the byte code
instructions = dis.get_instructions(code)
objects = {}
for instruction in instructions:
if instruction.opname == 'STORE_NAME':
objects[instruction.argval] = getattr(module_object,
instruction.argval)
elif instruction.opname == 'IMPORT_NAME':
consume_iterator(instructions, 2)
elif instruction.opname == 'IMPORT_FROM':
consume_iterator(instructions, 1)
return objects
print(get_defined_names('/Users/ashwini/py/so.py'))
For a file like:
#/Users/ashwini/py/so.py
import os
from sys import argv, modules
from math import *
from itertools import product
CST = True
from itertools import permutations, combinations
from itertools import chain
E = 100
from itertools import starmap
def func(): pass
for x in range(10):
pass
class C:
a = 100
d = 1
The output will be:
{'d': 1, 'E': 100, 'CST': True, 'x': 9, 'func': <function func at 0x10efd0510>, 'C': <class 'so.C'>}
A much more better way as someone already mentioned in comments will be to parse the source code using ast module and find out the variable names from there.

While I accepted an answer, it can't hurt to post the solution I ended up using. It's a mix between the other proposals :
import ast
import inspect
import importlib
from types import ModuleType
def extract_definitions(module):
""" Returns the name and value of objects defined at the top level of the given module.
:param module: A module object or the name of the module to import.
:return: A dict {'classes': {}, 'functions': {}, 'assignments': {}} containing defined objects in the module.
"""
if not isinstance(module, ModuleType):
module = importlib.import_module(module)
tree = ast.parse(inspect.getsource(module))
definitions = {'classes': {}, 'functions': {}, 'assignments': {}}
for node in tree.body:
if isinstance(node, ast.ClassDef):
definitions["classes"][node.name] = getattr(module, node.name)
elif isinstance(node, ast.FunctionDef):
definitions["functions"][node.name] = getattr(module, node.name)
elif isinstance(node, ast.Assign):
# for unpacking, you need to loop on all names
for target in node.targets:
definitions["assignments"][target.id] = getattr(module, target.id)
return definitions
I added the ability to import from a string or a module object, then removed the parsing of values and replaced it by a simple getattr from the original module.

Untested
def unexported_names (module):
try:
return [name for name in module.__dict__ if name not in module.__all__]
except AttributeError:
return [name for name in module.__dict__ if name.startswith('_')]