Search code examples
pythonpython-3.xclassabc

How to define two identical classes with different method parameter names?


I'm developing a scientific library where I would define vector functions in the time and frequency domain (linked by FFT). I created a class for vector formulas in the freq domain, and now I'd want to define an identical class for the time domain.

I want that in the time domain, the class functions - although being identical to their frequency-domain twin - have one parameter named t instead of omega. Is there an easier way of achieving this instead of repeated definition of every single method, while maintaining readibility?

My code:
(Note: my classes are much more complicated, and one can't just use the functions as formula.x_func(...) - some checking and etc are included. Also, there are actually 6 components.)

class VecFormula(object):
    pass

class FreqFormula(VecFormula):

    def __init__(self, x_func, y_func, z_func):
        self.x_func = x_func
        self.y_func = y_func
        self.z_func = z_func

    def x(self, x, y, z, omega, params):
        return self.x_func(x, y, z, omega, params)

    def y(self, x, y, z, omega, params):
        return self.y_func(x, y, z, omega, params)

    def z(self, x, y, z, omega, params):
        return self.z_func(x, y, z, omega, params)

    def component(self, comp, x, y, z, omega, params):
        if comp == 'x':
            return self.x(x, y, z, omega, params)
        elif comp == 'y':
            return self.y(x, y, z, omega, params)
        elif comp == 'z':
            return self.z(x, y, z, omega, params)
        else:
            raise ValueError(f'invalid component: {comp}')


class TimeFormula(FreqFormula):
    "same as FreqFormula, but the omega parameter is renamed to t"
    
    def x(self, x, y, z, t, params):
        return super(TimeFormula, self).x(x, y, z, t, params)

    def y(self, x, y, z, t, params):
        return super(TimeFormula, self).y(x, y, z, t, params)

    def z(self, x, y, z, t, params):
        return super(TimeFormula, self).z(x, y, z, t, params)

    def component(self, comp, x, y, z, t, params):
        return super(TimeFormula, self).component(x, y, z, t, params)

Solution

  • It is easy to add methods to a class after creation they are just class attributes. The hard part here, is that you need to dynamically create new functions to clone the methods from the original class to be able to change their signature.

    It is not the most clear part of Python, and dynamic function creation is not documented in the official reference documentations but can be found on SO: Python: dynamically create function at runtime

    So here is a possible way:

    # have the new class derive from the common base
    class TimeFormula(VecFormula):
        "same as FreqFormula, but the omega parameter is renamed to t"
        pass
    
    # loop over methods of origina class
    for i,j in inspect.getmembers(FreqFormula, inspect.isfunction):
        # copy the __init__ special method
        if i == '__init__':
            setattr(TimeFormula, i, j)
        elif i.startswith('__'): continue  # ignore all other special attributes
        if not j.__qualname__.endswith('.'.join((FreqFormula.__name__, i))):
            continue   # ignore methods defined in parent classes
        # clone the method from the original class
        spec = inspect.getfullargspec(j)
        newspec = inspect.FullArgSpec(['t' if i == 'omega' else i
                                       for i in spec.args], *spec[1:])
        f = types.FunctionType(j.__code__, j.__globals__, i, newspec, j.__closure__)
        f.__qualname__ = '.'.join((TimeFormula.__qualname__, i))
        # adjust the signature
        sig = inspect.signature(j)
        if ('omega' in sig.parameters):
            f.__signature__ = sig.replace(
                parameters = [p.replace(name='t') if name == 'omega' else p
                              for name, p in sig.parameters.items()])
        # and finally insert the new method in the class
        setattr(TimeFormula, i, f)