How to programmatically ensure that a function includes a return statement?

Viewed 250

How can I validate that a function includes a return keyword? I frequently forget the return line, so I am worried that the users of my package will too when they provide a function-based input.

def appler():
    a = "apple"
    # `return` is missing

def bananer():
    b = "banana"
    return b

I could parse the actual code string of the function for a final line that includes "return" but that isn't very robust (it could be triggered by comments).

def validate_funk(funk):
    if condition_to_check_that_it_contains_rtrn:
        pass
    else:
        raise ValueError(f"Yikes - The function you provided not contain a `return` statement:\n\n{funk}")
>>> validate_funk(appler)
#triggers ValueError

>>> validate_funk(bananer)
# passes

EDIT: ideally without running the function.

5 Answers

What you actually care about is probably not the return statement itself, but that the function returns something of a certain type. This you can most easily accomplish by using type hints (PEP 484):

def appler() -> str:
    a = "apple"
    # `return` is missing

def bananer() -> str:
    b = "banana"
    return b

Now, running a static analysis tool like mypy or Pyre (or many others): will emit a warning about a wrong return type in function appler (expected str, got NoneType).

Look for sabik's answer for a more general answer. Writing (unit) tests is another good practice that catches many more issues and - if done well - is an invest in code maintainability.

A function without return statement returns None by default.

>>> def abc():
    pass
>>> print(abc())
None
>>> 

You can add a check using this:

def validate_func(function):
    if function() == None:
        raise ValueError("Yikes - Does not contain a `return` statement")

There are few cons though.

  1. You have to execute the function
  2. It wont work if you are returning None in a function

Not much practical but yea, that is one way. You can also get a list of local functions or a list of methods in a class and loop through them without having to check each function individually.

For the question as asked, the ast module will let you check that.

However, it doesn't seem very useful just by itself - as others have pointed out, a function without a return is valid (it returns None), and just because a function does have a return doesn't mean that it returns the correct value, or even any value (the return could be in an if statement).

There are a couple of standard ways of dealing with this:

  • Unit tests - separate code that calls your function with various combinations of inputs (possibly just one, possibly hundreds) and checks that the answers match the ones you calculated manually, or otherwise satisfy requirements.

  • A more general implementation of the idea of checking for a return statement is "lint", in the case of Python pylint; that looks through your code and checks for various patterns that look like they could be mistakes. A side benefit is that it already exists and it checks dozens of common patterns.

  • Another, different more general implementation is the mypy type checker; that not only checks that there's a return statement, but also that it returns the correct type, as annotated in the header of the function.

Typically, these would be used together with a "gated trunk" development process; manual changes to the main version are forbidden, and only changes which pass the tests, lint and/or mypy are accepted into the main version.

As others have mentioned, simply calling the function is not enough: a return statement might only be present in a conditional, and thus, specific input would need to be passed in order to execute the return statement. That, too, is not a good indicator of the presence of a return, since it could return None, causing greater ambiguity. Instead, the inspect and ast module can be used:

Test functions:

def appler():
   a = "apple"
   # `return` is missing

def bananer():
   b = "banana"
   return b

def deeper_test(val, val1):
  if val and val1:
     if val+val1 == 10:
        return 

def gen_func(v):
   for i in v:
      if isinstance(i, list):
         yield from gen_func(i)
      else:
         yield i

inspect.getsource returns the entire source of the function as a string, which can then be passed to ast.parse. From there, the syntax tree can be recursively traversed, searching for the presence of a return statement:

import inspect, ast
fs = [appler, bananer, deeper_test, gen_func]
def has_return(f_obj):
   return isinstance(f_obj, ast.Return) or \
          any(has_return(i) for i in getattr(f_obj, 'body', []))

result = {i.__name__:has_return(ast.parse(inspect.getsource(i))) for i in fs}

Output:

{'appler': False, 'bananer': True, 'deeper_test': True, 'gen_func': False}

With a defined validate_funk:

def validate_funk(f):
   if not has_return(ast.parse(inspect.getsource(f))):
      raise ValueError(f"function '{f.__name__}' does not contain a `return` statement")
   return True

Notes:

  1. This solution does not require the test functions to be called.
  2. The solution must be run in a file. If it is run in the shell, an OSError will be raised. For the file, see this Github Gist.

You can simplify return checking with a decorator:

def ensure_return(func):
    def wrapper(*args, **kwargs):
        res = func(*args, **kwargs)
        if res is None:
            raise ValueError(f'{func} did not return a value')
        return res
    return wrapper

@ensure_return
def appler():
    a = "apple"
    # `return` is missing

@ensure_return
def bananer():
    b = "banana"
    return b

then:

>>> appler()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 5, in wrapper
ValueError: <function appler at 0x7f99d1a01160> did not return a value
>>> bananer()
'banana'
Related