CODEWITHSUNDEEP
pythonclassmethodsdecoratorwrapper
rockstiff
rockstiff

Reputation: 383

Class method as decorator in python

So, im writing a library for appium tests. I have a main class that look like this:

class APP():
    def __init__(self):
        self.variable1 = 1
        self.current_view = "main_screen"

    def do_operation_A(self):
        self.open_side_menu()
        do_something
        self.current_view = "side_menu"
    
    def do_operation_B(self):
        self.open_side_menu()
        do_something_else
        self.current_view = "side_menu"

    def set_landscape(self):
        self.open_settings_menu()
        configure_landscape
        self.current_view = "settings_menu"

The class has a lot of operations so i can do things like app.do_operation_A() or app.set_landscape() without having to first go to each menu manually (resolved inside the class)

To reduce this i want to implement a decorator to do something like this if possible:

class APP():
    def __init__(self):
        self.variable1 = 1
        self.current_view = "main_screen"

    #DEFINE_DECORATOR_HERE

    @side_menu
    def do_operation_A(self):
        do_something
    
    @side_menu
    def do_operation_B(self):
        do_something_else

    @settings_menu
    def set_landscape(self):
        configure_landscape

So i want to implement this decorators to navigate to the corresponding view and also change that variable that i use to check some things in other functions. I have seen some examples with functools.wraps but is not clear to me of how to implement the decorator inside the class to be able to modify this self variables.

Any help?

Upvotes: 2

Views: 4354

Answers (2)

talavis
talavis

Reputation: 66

Using a decorator means that you "wrap" your other function, i.e. you call the decorator and then call the function from inside the decorator.

E.g.:

import functools

def outer(func):
    @functools.wraps(func)
    def inner(*args, **kwargs):
        return func(*args, **kwargs)
    return inner

Upon defining the function, the decorator will be called, returning the inner function.

Whenever you call func, you will in reality call inner, which runs it's own code, including calling the original func function.

So for your use case, you should be able to create decorators similar to:

def settings_menu(func):
    @functools.wraps(func)
    def inner(self, *args, **kwargs):
        self.open_settings_menu()
        self.current_view = "settings_menu"
        return func(self, *args, **kwargs)
    return inner

Upvotes: 4

obeq
obeq

Reputation: 675

So a decorator is basically a function that returns another function, right?

def side_menu(func):
    def wrapper():
        return func()
    
    return wrapper

The wrapper, returned by side_menu, will be called whenever App().do_operationA is called. And whenever that method is called, self is always the first argument. Or rather, the first argument is the instance of App, but whatever. So we could do:

def side_menu(func):
    def wrapper(self, *args, **kwargs):
        self.open_menu()
        func(self, *args, **kwargs)

    return wrapper

Now, you don't want the method to present itself as wrapper - you like the name do_operationA. That's where @functools.wraps comes in, it makes things look and work right when decorating.

def side_menu(func):
    @functools.wraps
    def wrapper(self, *args, **kwargs):
        self.open_menu()
        func(self, *args, **kwargs)

    return wrapper

Upvotes: 2

Related Questions