CODEWITHSUNDEEP
pythonwrapperpython-decorators
Kurt Peek
Kurt Peek

Reputation: 57421

In Python, how to define a function wrapper which validates an argument with a certain name?

I'm writing several functions which accept an argument called policy, which is allowed only to have certain values (namely, 'allow' or 'deny'). If it doesn't, I would like a ValueError to be raised.

For brevity, I would like to define a decorator for this. So far, I have come up with the following:

def validate_policy(function):
    '''Wrapper which ensures that if the function accepts a 'policy' argument, that argument is either 'allow' or 'deny'.'''
    def wrapped_function(policy, *args, **kwargs):
        if policy not in ['allow', 'deny']:
            raise ValueError("The policy must be either 'allow' or 'deny'.")
        return function(policy, *args, **kwargs)
    return wrapped_function

The problem is that this only works if policy is the first positional argument of the function. However, I would like to allow for policy to appear at any position.

To be specific, here are some (dummy) functions called make_decision and make_informed_decision which accept an argument policy at different positions, and some test cases to go with them:

import pytest

@validate_policy
def make_decision(policy):      # The 'policy' might be the first positional argument
    if policy == 'allow':
        print "Allowed."
    elif policy == 'deny':
        print "Denied."

@validate_policy
def make_informed_decision(data, policy):   # It also might be the second one
    if policy == 'allow':
        print "Based on the data {data} it is allowed.".format(data=data)
    elif policy == 'deny':
        print "Based on the data {data} it is denied.".format(data=data)


'''Tests'''
def test_make_decision_with_invalid_policy_as_positional_argument():
    with pytest.raises(ValueError):
        make_decision('foobar')

def test_make_decision_with_invalid_policy_as_keyword_argument():
    with pytest.raises(ValueError):
        make_decision(policy='foobar')

def test_make_informed_decision_with_invalid_policy_as_positional_argument():
    with pytest.raises(ValueError):
        make_informed_decision("allow", "foobar")

def test_make_informed_decision_with_invalid_policy_as_keyword_argument():
    with pytest.raises(ValueError):
        make_informed_decision(data="allow", policy="foobar")


if __name__ == "__main__":
    pytest.main([__file__])

Currently all the tests pass except the third one, because the first positional argument 'allow' is interpreted as the policy rather than as the data as it should be.

How can I adapt the validate_policy decorator such that all the tests pass?

Upvotes: 3

Views: 136

Answers (2)

Kurt Peek
Kurt Peek

Reputation: 57421

Here is another solution using inspect.getcallargs:

def validate_policy(function):
    '''Wrapper which ensures that if the function accepts a 'policy' argument, that argument is either 'allow' or 'deny'.'''
    def wrapped_function(*args, **kwargs):
        call_args = inspect.getcallargs(function, *args, **kwargs)
        if 'policy' in call_args:
            if call_args['policy'] not in ['allow', 'deny']:
                raise ValueError("The policy must be either 'allow' or 'deny'.")
        return function(*args, **kwargs)
    return wrapped_function

It makes all the tests pass.

Upvotes: 1

Aran-Fey
Aran-Fey

Reputation: 43136

You can use the inspect module's Signature.bind function:

import inspect

def validate_policy(function):
    '''Wrapper which ensures that if the function accepts a 'policy' argument, that argument is either 'allow' or 'deny'.'''
    signature= inspect.signature(function)
    def wrapped_function(*args, **kwargs):
        bound_args= signature.bind(*args, **kwargs)
        bound_args.apply_defaults()
        if bound_args.arguments.get('policy') not in ['allow', 'deny']:
            raise ValueError("The policy must be either 'allow' or 'deny'.")
        return function(*args, **kwargs)
    return wrapped_function

Upvotes: 2

Related Questions