CODEWITHSUNDEEP
pythonreturn
Misha
Misha

Reputation: 127

Python Check if function has return statement

E.g.

def f1():
    return 1

def f2():
    return None

def f3():
    print("Hello")

Functions f1() and f2() returns something but f3() not.

a = f2()
b = f3()

And here a equals b so I can't just compare the result of functions to check if one has return or not.

Upvotes: 6

Views: 17839

Answers (3)

T. Smith
T. Smith

Reputation: 154

Inspired directly by @kevin's answer. Some rearranging to reduce the long-line-itis and a fix for one of the issues indicated in his answer, namely nested functions.

There are likely better ways to do this, I just figured an easy way is to override the behavior of ast.walk to have it skip the inner functions.

import ast
import inspect

def walk_special(node):
    from collections import deque
    todo = deque([node])
    top_level_node = node
    while todo:
        node = todo.popleft()

        # Added this conditional
        if isinstance(node, ast.FunctionDef):
            if node not in ast.iter_child_nodes(top_level_node):
                continue

        todo.extend(ast.iter_child_nodes(node))
        yield node

def contains_explicit_return(f):
    source = inspect.getsource(f)
    parsed = ast.parse(source)
    for node in walk_special(parsed):
        if isinstance(node, ast.Return):
            return True
    return False

def f1():
    return 1

def f2():
    return None

def f3():
    print("return")

def f4():
    def inner_a():
        return 2

    print("return")

def f5():
    def inner_b():
        return 3

    return "test"

for f in (f1, f2, f3, f4, f5):
    print(f, contains_explicit_return(f))

edit:

I realized this might not work in the case of function definitions nested inside other nodes. I imagine this is unlikely to be a common occurrence so, presumably, not a big deal.

Example:

def f6():
    if True:
        def foo():
            return 4
    print("done")

I would not have expected contains_explicit_return(f6) == False which is correct but unexpected from my current understanding.

Upvotes: 2

Kevin
Kevin

Reputation: 76194

I like @st0le's idea of inspecting the source, but you can take it a step further and parse the source into a source tree, which eliminates the possibility of false positives.

import ast
import inspect

def contains_explicit_return(f):
    return any(isinstance(node, ast.Return) for node in ast.walk(ast.parse(inspect.getsource(f))))

def f1():
    return 1

def f2():
    return None

def f3():
    print("return")

for f in (f1, f2, f3):
    print(f, contains_explicit_return(f))

Result:

<function f1 at 0x01D151E0> True
<function f2 at 0x01D15AE0> True
<function f3 at 0x0386E108> False

Of course, this only works for functions that have source code written in Python, and not all functions do. For instance, contains_explicit_return(math.sqrt) will give you a TypeError.

Furthermore, this won't tell you anything about whether any particular execution of a function hit a return statement or not. Consider the functions:

def f():
    if random.choice((True, False)):
        return 1

def g():
    if False:
        return 1

contains_explicit_return will give True on both of these, despite f not encountering a return in half of its executions, and g not encountering a return ever.

Last, it won't distinguish between whether the function has a return statement or a function defined inside it has a return statement:

def f():
    def g():
        return 5

In this case, contains_explicit_return(f) will return True, even though f will never return anything, because the function g inside contains an explicit return.

Upvotes: 13

st0le
st0le

Reputation: 33545

A function by definition always returns something. Even if you don't specify it, there is an implicit return None at the end of a python function.

You can check for a "return" with the inspect module.

EDIT: I just realized. This is horribly wrong because, it'll return True if there's a string literal in the function that has "return" in it. I suppose a robust a regex will help here.

from inspect import getsourcelines


def f(n):
    return 2 * n


def g(n):
    print(n)


def does_function_have_return(func):
    lines, _  = getsourcelines(func)
    return any("return" in line for line in lines) # might give false positives, use regex for better checks


print(does_function_have_return(f))
print(does_function_have_return(g))

Upvotes: 9

Related Questions