En Python, comment définir un wrapper de fonction qui valide un argument avec un certain nom ?

Question

J'écris plusieurs fonctions qui acceptent un argument appelé policy qui ne peut avoir que certaines valeurs (à savoir, 'allow' o 'deny' ). Si ce n'est pas le cas, j'aimerais recevoir un ValueError à soulever.

Par souci de concision, j'aimerais définir un décorateur pour cela. Jusqu'à présent, j'ai trouvé ce qui suit :

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

Le problème est que cela ne fonctionne que si policy est le premier argument de position de la fonction. Cependant, j'aimerais permettre à policy pour apparaître à n'importe quel endroit.

Pour être plus précis, voici quelques fonctions (factices) appelées make_decision y make_informed_decision qui acceptent un argument policy à différentes positions, et quelques cas de test pour les accompagner :

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__])

Actuellement, tous les tests sont réussis, sauf le troisième, car le premier argument positionnel 'allow' est interprété comme le policy plutôt qu'en tant que data comme il se doit.

Comment puis-je adapter le validate_policy de manière à ce que tous les tests soient réussis ?

Answer 1

Vous pouvez utiliser le inspect du module Signature.bind fonction :

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

Answer 2

Voici une autre solution utilisant 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

Il permet de réussir tous les tests.