Python tracking sub-classes that are in scope

Question

I'm trying to write a tracker class where the instances of the tracker class track the sub-classes of another class that are in the scope of the tracker instance.

More concretely, the following is an example of what I am trying to achieve:

class Foo(object): pass

class FooTracker(object):
     def __init__(self):

          # use Foo.__subclasses__() or a metaclass to track subclasses 
          # - but how do I filter this to only get the ones in scope?

          self.inscope = <something magic goes here>

ft1 = FooTracker()
assert ft1.inscope == []

class Bar(Foo): pass
ft2 = FooTracker()
assert ft2.inscope == [<class '__main__.Bar'>]

def afunction():
    class Baz(Foo): pass    # the global definition of Bar is now hidden
    class Bar(Foo): pass
    ft3 = FooTracker()

    assert (set(ft3.inscope) == set([<class '__main__.afunction.<locals>.Baz'>,
                                     <class '__main__.afunction.<locals>.Bar'>])

ft4 = FooTracker()   # afunction.Baz and afunction.Bar are no longer in scope
assert ft4.inscope == [<class '__main__.Bar'>]

So I want the instances of FooTracker to track the sub-classes of Foo that were in scope at the time the FooTracker object was created.

I've tried a few different things, such as parsing the qualified names of the Foo sub-classes and using exec() to do the name resolution but the fundamental problem is that it always works out the sub-classes relative to the scope within FooTracker.__init__() and not where it was called.

My only other thought was to try something with inspect.currentframe() but even if this were possible it would probably be too much of a hack and would make the code too brittle (e.g., there is a comment in the docs that not all Python implementations will have frame support in the interpreter").

Answer 1

There's no easy way to do exactly what you're asking for. But you might be able to use some Python features to get something with a roughly similar API, without as much hassle.

One option would be to require each subclass to be decorated with a method of your Tracker class. This would make it really easy to keep track of them, since you'd just append each caller of the method to a list:

class Tracker:
    def __init__(self):
        self.subclasses = []

    def register(self, cls):
        self.subclasses.append(cls)
        return cls

class Foo(): pass

foo_tracker = Tracker()

@foo_tracker.register
class FooSubclass1(Foo): pass

@foo_tracker.register
class FooSubclass2(Foo): pass

print(foo_tracker.subclasses)

This doesn't actually require that the classes being tracked are subclasses of Foo, all classes (and even non-class objects) can be tracked if you pass them to the register method. Decorator syntax makes it a little nicer than just appending each class to a list after you define it, but not by a whole lot (you still repeat yourself a fair amount, which may be annoying unless you make the tracker and method names very short).

A slightly trickier version might get passed the base class, so that it would detect subclasses automatically (via Foo.__subclasses__). To limit the subclasses it detects (rather than getting all subclasses of the base that have ever existed), you could make it behave as a context manager, and only track new subclasses defined within a with block:

class Tracker:
    def __init__(self, base):
        self.base = base
        self._exclude = set()
        self.subclasses = set()

    def __enter__(self):
        self._exclude = set(self.base.__subclasses__())
        return self

    def __exit__(self, *args):
        self.subclasses = set(self.base.__subclasses__()) - self._exclude
        return False

class Foo(): pass
class UntrackedSubclass1(Foo): pass

with Tracker(Foo) as foo_tracker:
    class TrackedSubclass1(Foo): pass
    class TrackedSubclass2(Foo): pass

class UntrackedSubclass2(Foo): pass

print(foo_tracker.subclasses)

If you're using Python 3.6 or later, you can do the tracking a different way by injecting an __init_subclass__ class method into the tracked base class, rather than relying upon __subclasses__. If you don't need to support class hierarchies that are already using __init_subclass__ for their own purposes (and you don't need to support nested trackers), it can be quite elegant:

class Tracker:
    def __init__(self, base):
        self.base = base
        self.subclasses = []

    def __enter__(self):
        @classmethod
        def __init_subclass__(cls, **kwargs):
            self.subclasses.append(cls)

        self.base.__init_subclass__ = __init_subclass__
        return self

    def __exit__(self, *args):
        del self.base.__init_subclass__
        return False

class Foo(): pass
class UntrackedSubclass1(Foo): pass

with Tracker(Foo) as foo_tracker:
    class TrackedSubclass1(Foo): pass
    class TrackedSubclass2(Foo): pass

class UntrackedSubclass2(Foo): pass

print(foo_tracker.subclasses)

One nice feature of this version is that it automatically tracks deeper inheritance hierarchies. If a subclass of a subclass is created within the with block, that "grandchild" class will still be tracked. We could make the previous __subclasses__ based version work this way too, if you wanted, by adding another function to recursively expand out the subclasses of each class we find.

If you do want to play nice with existing __init_subclass__ methods, or want to be able to nest trackers, you need to make the code a bit more complicated. Injecting a well behaved classmethod in a reversible way is tricky since you need handle both the case where the base class has its own method, and the case where it's inheriting a version from its parents.

class Tracker:
    def __init__(self, base):
        self.base = base
        self.subclasses = []

    def __enter__(self):
        if '__init_subclass__' in self.base.__dict__:
            self.old_init_subclass = self.base.__dict__['__init_subclass__']
        else:
            self.old_init_subclass = None

        @classmethod
        def __init_subclass__(cls, **kwargs):
            if self.old_init_subclass is not None:
                self.old_init_subclass.__get__(None, cls)(**kwargs)
            else:
                super(self.base, cls).__init_subclass__(**kwargs)
            self.subclasses.append(cls)

        self.base.__init_subclass__ = __init_subclass__
        return self

    def __exit__(self, *args):
        if self.old_init_subclass is not None:
            self.base.__init_subclass__ = self.old_init_subclass
        else:
            del self.base.__init_subclass__
        return False

class Foo:
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        print("Foo!")

class Bar(Foo): pass   # every class definition from here on prints "Foo!" when it runs

with Tracker(Bar) as tracker1:
    class Baz(Bar): pass

    with Tracker(Foo) as tracker2:
        class Quux(Foo): pass

        with Tracker(Bar) as tracker3:
            class Plop(Bar): pass

# four Foo! lines will have be printed by now by Foo.__init_subclass__
print(tracker1.subclasses) # will describe Baz and Plop, but not Quux
print(tracker2.subclasses) # will describe Quux and Plop
print(tracker3.subclasses) # will describe only Plop