static analysis - How can I find Python methods without return statements?

Question

I really like it when methods of objects, which modify the objects property, return self so that you can chain method calls. For example:

boundingBox.grow(0.05).shift(x=1.3)

instead of

boundingBox.grow(0.05)
boundingBox.shift(x=1.3)

I would like to search the code of my old projects to adjust this pattern. How can I find methods which don't have a return statement?

Ideally, I would like to let a program run over a folder. The program searches Python files, looks for classes, examines their methods and searches return statements. If no return statement is there, it outputs the filename, the name of the class and the name of the method.

Answer 1

You can get the names with ast, I will work on getting the line numbers:

import inspect
import importlib
import ast

class FindReturn(ast.NodeVisitor):
    def __init__(self):
        self.data = []

    def visit_ClassDef(self,node):
        self.data.append(node.name)
        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        if not any(isinstance(n, ast.Return) for n in node.body):
            self.data.append(node.name)
        self.generic_visit(node)

mod = "test"
mod = importlib.import_module(mod)
p = ast.parse(inspect.getsource(mod))

f = FindReturn()
f.visit(p)

print(f.data)

Input:

class Foo(object):
    def __init__(self):
        self.foo = "foo"

    def meth1(self):
        self.bar = "bar"

    def meth2(self):
        self.foobar = "foobar"


    def meth3(self):
        self.returns = "foobar"
        return self.returns

class Bar(object):
    def __init__(self):
        self.foo = "foo"

    def meth1(self):
        self.bar = "bar"

    def meth2(self):
        self.foobar = "foobar"


    def meth3(self):
        self.returns = "foobar"
        return self.returns

Output:

['Foo', '__init__', 'meth1', 'meth2', 'Bar', '__init__', 'meth1', 'meth2']

The filename is obviously "test.py" here.

This is probably a nicer way to group the data:

import inspect
import importlib
import ast
from collections import defaultdict

mod = "test"
mod = importlib.import_module(mod)
p = ast.parse(inspect.getsource(mod))



data = defaultdict(defaultdict)
classes = [cls for cls in p.body if isinstance(cls, ast.ClassDef)]
for cls in classes:
    name = "class_{}".format(cls.name)
    data[mod][name] = {"methods": []}
    for node in cls.body:
        if not any(isinstance(n, ast.Return) for n in node.body):
            if node.name != "__init__":
                data[mod][name]["methods"].append(node.name)

Output:

{<module 'test' from '/home/padraic/test.pyc'>: defaultdict(None, {'class_Foo': {'methods': ['meth1', 'meth2']}, 'class_Bar': {'methods': ['meth1', 'meth2']}})}

To go through a directory:

data = defaultdict(defaultdict)
import os
path = "/home/padraic/tests"
for py in os.listdir(path):
    with open(os.path.join(path,py)) as f:
        p = ast.parse(f.read(), "", "exec")

    classes = [cls for cls in p.body if isinstance(cls, ast.ClassDef)]
    for cls in classes:
        name = "class_{}".format(cls.name)
        data[py][name] = {"methods": []}
        for node in cls.body:
            if not any(isinstance(n, ast.Return) for n in node.body):
                if node.name != "__init__":
                    data[py][name]["methods"].append(node.name)


from pprint import pprint as pp

pp(dict(data))

{'test.py': defaultdict(None, {'class_Foo': {'methods': ['meth1', 'meth2']}, 
'class_Bar': {'methods': ['meth1', 'meth2']}}),'test2.py': 
defaultdict(None, {'class_Test2': {'methods': ['test1', 'test2']}})}

Where test2 contains:

class Test2:
    def test1(self):
        pass

    def test2(self):
        self.f=4
        s = self.test_return()
        i = 3

    def test_return(self):
        return "Test2"

You can get the line before the method definition with node.lineno:

classes = [cls for cls in p.body if isinstance(cls, ast.ClassDef)]
    for cls in classes:
        name = "class_{}".format(cls.name)
        data[py][name] = {"methods": []}
        for node in cls.body:
            if not any(isinstance(n, ast.Return) for n in node.body):
                if node.name != "__init__":
                    data[py][name]["methods"].append({"meth":node.name,"line":node.lineno})

Output:

{'test.py': defaultdict(None, {'class_Foo': {'methods': [{'meth': 'meth1', 'line': 6}, {'meth': 'meth2', 'line': 9}]}, 'class_Bar': {'methods': [{'meth': 'meth1', 'line': 21}, {'meth': 'meth2', 'line': 24}]}}),
 'test2.py': defaultdict(None, {'class_Test2': {'methods': [{'meth': 'test1', 'line': 2}, {'meth': 'test2', 'line': 5}]}})}

Or we can guesstimate where the return is missing by getting the line number from the last arg in the body:

data[py][name]["methods"].append({"meth":node.name,"line": node.body[-1].lineno})

Output:

{'test.py': defaultdict(None, {'class_Foo': {'methods': [{'meth': 'meth1', 'line': 7},
 {'meth': 'meth2', 'line': 10}]}, 'class_Bar': {'methods': [{'meth': 'meth1', 'line': 22}, {'meth': 'meth2', 'line': 25}]}}),
 'test2.py': defaultdict(None, {'class_Test2': {'methods': [{'meth': 'test1', 'line': 3}, {'meth': 'test2', 'line': 8}]}})}

It might also be better to use iglob to ignore other files:

import glob
for py in glob.iglob(os.path.join(path,"*.py")):
    with open(os.path.join(path, py)) as f:
        p = ast.parse(f.read(), "", "exec")