Custom python decorator checking function after each line

Question

i have an example function below:

def test_function(test_df):
    test_df = do_something1(test_df)
    test_df = do_something2(test_df)
    test_df = do_something3(test_df)
    return test_df

where I input a pandas dataframe, call some other functions on the inputted dataframe and output it at the end.

My question is, is there a way to define a python wrapper (or any other method) that checks after each line if your inputted dataframe (test_df) is empty, and if so, return an empty dataframe? My purpose for this is, if any of these functions filter the data such that test_df is empty, the code will break in any of the subsequent function calls. The equivalent function to the wrapper I'm looking for would be:

def test_function(test_df):
    test_df = do_something1(test_df)
    if test_df.empty:
        return pd.DataFrame()
    test_df = do_something2(test_df)
    if test_df.empty:
        return pd.DataFrame()
    test_df = do_something3(test_df)
    return test_df

I think for a few function calls, the empty checks are fine, but think this detracts from readability and as the number of function calls you make increased, it gets more and more cluttered.

Ive looked at defining custom wrappers but they seem to look at the function as a whole, not line by line. Also looked into context managers but it doesnt quite seem to fit my purpose. Not really looking to change the code all too much, hence why a nice wrapper would be ideal. Any other clean solution with the same output would be just as nice.

Answer 1

I would just use a loop.

def test_function(test_df):
    for f in [do_something1, do_something2, do_something3]:
        test_df = f(test_df)
        if test_df.empty:
            return pd.DataFrame()
    return test_df

Answer 2

Decorators aren't intended to be used to modify the contents of the function being decorated. So, in the spirit of your question, you can't (or at least shouldn't) create a decorator for test_function. They work the best when you want to add a little behavior before and/or after a function's call.

If you want a decorator, you would use it on your called functions and it would look something like this:

# the decorator is handed the function being decorated
def decorator(func):
    # define a function that will be called instead
    # of the decorated function
    def wrapper(df):
        # call the function being decorated
        result = func(df)
        if result.empty:
            return pd.DataFrame()
        return result
    # return the function that will replace
    # the function being decorated
    return wrapper

@decorator
def do_something1(df):
    # some code
    return df

@decorator
def do_something2(df):
    # some code
    return df

@decorator
def do_something3(df):
    # some code
    return df

def test_function(test_df):
    test_df = do_something1(test_df)
    test_df = do_something2(test_df)
    test_df = do_something3(test_df)
    return test_df

Decorating a function is equivalent to this:

def decorator(func):
    def wrapper(df):
        result = func(df)
        if result.empty:
            return pd.DataFrame()
        return result
    return wrapper

# The function to decorate
def do_something2(df):
    return df

do_something2 = decorator(do_something2)

You effectively replace the function with a wrapped version of the function.

Alternatively, you could just create a wrapper for the result, which would mean no extra code added to the definitions.

def workable_result(df):
    if df.empty:
        return pd.DataFrame()
    return df

def test_function(test_df):
    test_df = workable_result(do_something1(test_df))
    test_df = workable_result(do_something2(test_df))
    test_df = do_something3(test_df)
    return test_df