Home > Software engineering >  Wrap all methods of class that contain specific argument in Python
Wrap all methods of class that contain specific argument in Python

Time:09-02

I have a class Stuff that has several methods, some of which have some argument, let's call it argument_x. For example:

class Stuff:
    def method_1(self, argument_x, **other_args):
        pass
    def method_2(self, argument_x, **other_args):
        pass
    def method_3(self, I_dont_have_argument_x):
        pass

Now I want to subclass this class wrapping all methods that have argument_x in the same way. For example if I were to proceed by hand I would do:

class StuffWithConstantX(Stuff):
    def __init__(self, argument_x_value):
        super().__init__()
        self._argument_x_value = argument_x_value
    def method_1(self, **other_args):
        super().method_1(argument_x=self._argument_x_value, **other_args)
    def method_2(self, **other_args):
        super().method_2(argument_x=self._argument_x_value, **other_args)

As method_3 does not have argument_x I leave it unchanged.

Is it possible to automate this? How?

CodePudding user response:

Here's how you might define this as a wrapper, rather than a subclass:

class Stuff:
    def method_1(self, argument_x, **other_args):
        print("method 1:", argument_x)
    def method_2(self, argument_x, **other_args):
        print("method 2:", argument_x)
    def method_3(self, i_dont_have_argument_x):
        print("method 3:", i_dont_have_argument_x)


class StuffWithConstantX:
    def __init__(self, argument_x_value) -> None:
        self._stuff = Stuff()
        self._argument_x = argument_x_value

    def __getattr__(self, __name: str):
        attr = getattr(self._stuff, __name)
        if not callable(attr):
            return attr
        def wrapped(*args, **kwargs):
            try:
                return attr(argument_x=self._argument_x, *args, **kwargs)
            except TypeError:
                # Beware -- if there's a TypeError raised from attr itself,
                # it will get run twice before the caller sees the exception.
                # You can potentially work around this by closely inspecting
                # either the exception or the attr object itself.
                return attr(*args, **kwargs)
        return wrapped


stuff = StuffWithConstantX("foo")
stuff.method_1()
stuff.method_2()
stuff.method_3("bar")
method 1: foo
method 2: foo
method 3: bar

As noted in the comments, this code is more or less impossible to statically typecheck, and I would not recommend actually using this pattern unless you have a really good reason.

CodePudding user response:

Here's another way you could do it.

import inspect
import functools

class StuffWithConstantX(Stuff):
    def __init__(self, argument_x_value):
        super().__init__()
        self._argument_x_value = argument_x_value

        for func_name, func in inspect.getmembers(Stuff, inspect.isfunction):
            arg_names = inspect.getfullargspec(func).args
            if 'argument_x' in arg_names:
                setattr(self, func_name, functools.partial(func, self=self, argument_x=self._argument_x_value))
  • Related