Home > Back-end >  Python decorator to check all parameters of a function which may include a list
Python decorator to check all parameters of a function which may include a list

Time:12-28

I'm trying to implement a decorator function to sanitize their parameters for mongoDB. This is what I have so far:

def sanitize(function_to_decorate):
    def wrapper(*args):
        for query in args:
            if any(elem in r'${/}\\' for elem in query):
                raise Exception('String contains invalid characters')
        function_to_decorate(*args)
    return wrapper

@sanitize
def test(arg, arg2, list):
    print (arg, arg2, list) 

test('Hi', 'Me', '1') # Passes
test('Hi', 'Me', '{1') # Fails
test('Hi',  'Hey', ['Me', '{1']) # Passes

This makes sense as I'm only iterating over all parameters once, but since some functions may require lists as parameters (say, for a $in check) this isn't universally useful.

How can I drill down on an arbitrary number of lists in a list? Say

def a(str, str, [] )
def b([], [] ,[])

I also don't know how to implement this for keyword arguments, but that's a decorator specific issue.

CodePudding user response:

You can pass off the responsibility of detecting parameters to a function which checks them recursively:

def rsanitize(args):
    for query in args:
        if isinstance(query, list):
            rsanitize(query)
        elif any(elem in r'${/}\\' for elem in query):
            raise Exception('String contains invalid characters')
    
def sanitize(function_to_decorate):
    def wrapper(*args):
        rsanitize(args)
        function_to_decorate(*args)
    return wrapper

Now your final test case will fail properly:

Traceback (most recent call last):
  File "D:\xxx\soRecursivelyTestLists.py", line 22, in <module>
    test('Hi',  'Hey', ['Me', '{1']) # Passes
  File "D:\xxx\soRecursivelyTestLists.py", line 12, in wrapper
    rsanitize(args)
  File "D:\xxx\soRecursivelyTestLists.py", line 6, in rsanitize
    rsanitize(query)
  File "D:\xxx\soRecursivelyTestLists.py", line 8, in rsanitize
    raise Exception('String contains invalid characters')
Exception: String contains invalid characters
  • Related