Python - Choose which class to inherit

Question

I want to make two classes A and B, in which B is a slight - but significant - variation of A, and then make a third class C that can inherit either A or B and add functionality to them. The problem is, how do I tell C to inherit A or B based on my preference?

To make things more clear, suppose I have this code:

class A:

   def __init__(self, x, y):
       self.x = x
       self.y = y

   def first(self):
       return do_something(1)

   def second(self):
       return do_something(2)

   def third(self):
       return do_something(3)

   def start(self):
       self.first()
       self.second()
       self.third()

class B(A):

    def __init__(self, x, y, z):
        super().__init__(x, y)
        self.z = z

    def second(self):
        super().second()
        do_stuff()

    def third(self):
        do_other_stuff()

That is a very simplified version of the code I used. In particular, A represents a simulator of a manufacturing system, while B represents a simulator of the same manufacturing system with a modification of the behaviour of the main machine-tool.

Now, what I want is to add code to compute some statistics. What it does is something like this:

class C(A):

    def __init__(self, *args):
        super().__init__(*args)
        self.stat = 0

    def second(self):
        super().second()
        self.stat += 1

    def third(self):
        super().third()
        self.stat *= 3

The problem is that the class C works the exactly same way whether if I inherit class A (as in the previous listing) or class B (exact same code, with as first line class C(B):

How can I do that? Or am I using a non-feasible way? I think an ideal solution is to be able to choose which class to inherit, A or B, when I initialize C. Or, maybe, to be able to pass to class C the class to inherit.

I made some researches, and I found also the possibility of aggregation (that I didn't know before), but I don't see it really useful. As a last note, be aware that class A might have up to 20-30 methods, and when I use class C I want class A (or B, depending on which it inherits) to work exactly as before with the added chunks of C inbetween.

P.S. I'm looking for a possibly elegant, no code-heavy, "pythonic" way of doing this. I'm also really looking forward on advices on everything you think could be done better. Finally, I can totally modify class C, but class A and B must remain (apart from small changes) the same.

Answer 1

You can use new-style classes and their method resolution order.

Considering these definitions:

class A(object):
    def __init__(self, x):
        pass

    def foo(self):
        print "A"

class B(object):
    def __init__(self, x, y):
        pass

    def foo(self):
        print "B"

you can build a mixin intended to add functionality to A or B:

class Cmix(object):
    def foo(self):
        super(Cmix, self).foo()
        print "mix"

and inherit from both Cmix and A (or B, respectively):

class CA(Cmix, A):
    pass

class CB(Cmix, B):
    pass

Finally, you can write a convenience function to choose between CA and CB based on the number of parameters:

def C(*args):
    if len(args) == 1:
        return CA(*args)
    else:
        return CB(*args)

Now we have

C(1).foo()
# A
# mix
C(1, 2).foo()
# B
# mix

Note that C is not a real class and you cannot use it as a second argument in isinstance, for example.