Shortening a oft-used code segment for testing a return value in Python

Question

Consider this Python segment:

def someTestFunction():
    if someTest:
        return value1
    elif someOtherTest:
        return value2
    elif yetSomeOtherTest:
        return value3
    return None

def SomeCallingFunction():
    a = someTestFunction()
    if a != None:
        return a

    ... normal execution continues

Now, the question: the three-line segment in the beginning of SomeCallingFunction to get the value of the test function and bail out if it's not None, is repeated very often in many other functions. Three lines is too long. I want to shorten it to one. How do I do that?

I can freely restructure this code and the contents of someTestFunction however needed. I thought of using exceptions, but those don't seem to help in cutting down the calling code length.

(I've read a bit about Python decorators, but haven't used them. Would this be the place? How would it work?)

Answer 1

I often use a hash table in place of a series of elifs:

def someTestFunction(decorated_test):
    options = {
        'val1': return_val_1,
        'val2': return_val_2
    }
    return options[decorated_test]

You can set up options as a defaultdict(None) to default to None if a key isn't found.

If you can't get your tests in that form, then a series of if statements might actually be the best thing to do.

One small thing you can do to shorten your code is to use this:

if a: return a

There may be other ways to shorten your code, but these are the ones I can come up with on the spot.

Answer 2

If you want to use a decorator, it would look like this:

def testDecorator(f):
    def _testDecorator():
        a = someTestFunction()
        if a is None:
            return f()
        else: return a
    return _testDecorator

@testDecorator
def SomeCallingFunction():
    ... normal execution

When the module is first imported, it runs testDecorator, passing it your original SomeCallingFunction as a parameter. A new function is returned, and that gets bound to the SomeCallingFunction name. Now, whenever you call SomeCallingFunction, it runs that other function, which does the check, and returns either a, or the result of the original SomeCallingFunction.

Answer 3

I think this would do it:

UPDATE Fixed!
Sorry for yesterday, I rushed and didn't test the code!

def test_decorator( test_func ):
    def tester( normal_function ):
        def tester_inner():
            a = test_func()
            if a is not None: 
                return a
            return normal_function()
        return tester_inner    
    return tester 


#usage:
@test_decorator( my_test_function )
def my_normal_function():
    #.... normal execution continue ...

It's similar to DNS's answer but allows you to specify which test function you want to use