Pass self to decorator object

Question

Having such decorator object

class wait_for_page_load(object):

    def __init__(self, driver, time_to_wait=20):
        self.driver = driver
        self.time_to_wait = time_to_wait

    def __call__(self, function):
        @functools.wraps(function)
        def wrapper(*args):
            old_page = self.driver.find_element_by_tag_name('html')
            function(*args)
            WebDriverWait(self.driver, self.time_to_wait).until(staleness_of(old_page))
        return wrapper

I want to apply it to method of another class, like this:

class VehiclePage(object):

    def __init__(self, driver):
        self.driver = driver

    @wait_for_page_load(self.driver)
    def open(self):
        self.driver.get('%s/vehicles/' % BASE_URL)

This gives me an error. Is there a way to pass self.driver to decorator?

Answer 1

You don't have to pass self to a decorator object. If the decorator returns a function, then that function will get access to self when it is called. eg.

def pass_value(function):
    def wrapper(self):
        function(self, self.value)
    return wrapper

class Printer(object):
    def __init__(self, value):
        self.value = value

    @pass_value
    def print_(self, v):
        print v

Printer("blah").print_()

The one problem with this method is that it requires self to implement a specific interface (such as having a field called driver, rather than directly passing the driver to the decorator).

Your decorator would become:

def wait_for_page_load(time_to_wait=20):
    def decorator(function):
        @functools.wraps(function)
        def wrapper(self, *args):
            old_page = self.driver.find_element_by_tag_name('html')
            function(self, *args)
            WebDriverWait(self.driver, time_to_wait).until(staleness_of(old_page))
        return wrapper
    return decorator

Used as:

@wait_for_page_load() # brackets are needed
def open(self):
    ...

Answer 2

Short answer: No there is not.

Long answer: The driver attribute is set when you instantiate the class. However, the decorator is run when the class is interpreted. That is, when the interpreter first reads it when loading the module. At this point, you don't have any instance ready. To do this kind of stuff you will have to refactor your code.

Also, even if that worked, you would end up using a single instance of your decorator class for all your objects. Probably not what you expected.

A simple workaround, though, could be to apply the decorator in __init__. Though not very elegant, that would work if you really need to apply the decorator.

def __init__(self, driver):
    self.driver = driver
    self.open = wait_for_page_load(self.driver)(self.open)

But then I believe you need to bind the wrapper to the class yourself by calling types.MethodType - honestly, it's probably better you just reorganize your code.