python decorate class method sending a parameter to the decorator-CodePudding

I have the following class in python:

class MatrixOfCarriers(DataLoader):
    def load_data(self):
        brands = self.brands
        try:
            data = self._load_data()
            print("'Carrier Matrix' data loaded successfully.\n")
        except Exception as e:
            print(
                "\nCouldn't load 'Carier Matrix' data due to the following "
                  "error: '{}'".format(str(e))
            )
            raise e

        return data

I want to decorate method MatrixOfCarriers.load() with the following decorator:

def cast_pct_buffer_columns(brands):
    def inner(func):
        def wrapper(*args, **kwargs):
            data = func(*args, **kwargs)

            for brand in func.brands.values():
                if brand["pct_buffer_column"] in data.columns:
                    data[brand["pct_buffer_column"]] = (
                        pd.to_numeric(
                            data[brand["pct_buffer_column"]].str.replace(
                                "[,%]", "", regex=True
                            )
                        )
                        / 100
                    )

            return data

        return wrapper

    return inner

Thr problem is that such method requires a parameter brand that is available as an instance variable, but I can't send it using:

@cast_pct_buffer_columns(self.brands)
def load_data(self):

but self.brands is not in scope out of the body of an instance method.

I also tried to set brands = self.brands in the body of method load_data() and then called brands = func.brands from the decorator, but it didn't work neither.

How can I do this?

CodePudding user response：

If you are always decorating a method, and will always need to get the .brands attribute, there are two things that can be changed: first, Python will pass the instance itself as the first argument to your wrapper - just as it pass the self parameter - so you can just use the .brands attribute in that first parameter.

Since the wrapper method can have access to the needed data, there is no need for an intermediate level of the decorator, to pass it the parameters, so it could be simply:

def cast_pct_buffer_columns(func):
    def wrapper(instance, *args, **kwargs):
        data = func(instance, *args, **kwargs)

        for brand in instance.brands.values():
            if brand["pct_buffer_column"] in data.columns:
                data[brand["pct_buffer_column"]] = (
                    pd.to_numeric(
                        data[brand["pct_buffer_column"]].str.replace(
                            "[,%]", "", regex=True
                        )
                    )
                    / 100
                )

        return data
    return wrapper

(here, I wrote "instance" as the parameter that will take the value that is usually called self inside a method. You have to add it explicitly to the inner call of func)

Now, if the decorator will be used in other situations and classes, and the desired values won't always be in the .brands attribute - let's suppose in another class, the same decorator were applied, but it'd need to pick the vendor attribute - then, that can be passed as a string to the decorator - and you can use the getattr built-in function to get to the values themselves:

def cast_pct_buffer_columns(attrname):
        def wrapper(instance, *args, **kwargs):
            data = func(*args, **kwargs)

            for item in getattr(instance, attrname).values():
                # do stuff
                ....
                
            return data
        return wrapper
    return inner


class MatrixOfVendorCarriers(DataLoader):
    @cast_pct_buffer_columns("vendors")
    def load_data(self):
        vendors self.vendors
        ...
        return data