Search code examples
pythonlistdecorator

Python decorator to check all parameters of a function which may include a list


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.


Solution

  • 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