Class with a registry of methods based on decorators

Question

I have a class that has several methods which each have certain properties (in the sense of quality). I'd like these methods to be available in a list inside the class so they can be executed at once. Note that the properties can be interchangeable so this can't be solved by using further classes that would inherit from the original one. In an ideal world it would look something like this:

class MyClass:
    def __init__():
        red_rules = set()
        blue_rules = set()
        hard_rules = set()
        soft_rules = set()

    @red
    def rule_one(self):
        return 1

    @blue
    @hard
    def rule_two(self):
        return 2

    @hard
    def rule_three(self):
        return 3

    @blue
    @soft
    def rule_four(self):
        return 4

When the class is instantiated, it should be easy to simply execute all red and soft rules by combining the sets and executing the methods. The decorators for this are tricky though since a regular registering decorator can fill out a global object but not the class attribute:

def red(fn):
    red_rules.add(fn)
    return fn

How do I go about implementing something like this?

Answer 1

Decorators are applied when the function is defined ; in a class that's when the class is defined. At this point in time there are no instances yet!

You have three options:

1. Register your decorators at the class level. This is not as clean as it may sound; you either have to explicitly pass additional objects to your decorators ( red_rules = set() , then @red(red_rules) so the decorator factory can then add the function to the right location), or you have to use some kind of class initialiser to pick up specially marked functions; you could do this with a base class that defines the __init_subclass__ class method , at which point you can iterate over the namespace and find those markers (attributes set by the decorators).

2. Have your __init__ method (or a __new__ method) loop over all the methods on the class and look for special attributes the decorators have put there.

   The decorator would only need to add a _rule_name or similar attribute to decorated methods, and {getattr(self, name) for for name in dir(self) if getattr(getattr(self, name), '_rule_name', None) == rule_name} would pick up any method that has the right rule name defined in rule_name .

3. Make your decorators produce new descriptor objects ; descriptors have their __set_name__() method called when the class object is created. This gives you access to the class, and thus you can add attributes to that class.

Note that __init_subclass__ and __set_name__ require Python 3.6 or newer; you'd have to resort to a metaclass to achieve similar functionality in earlier versions.

Also note that when you register functions at the class level, that you need to then explicitly bind them with function.__get__(self, type(cls)) to turn them into methods, or you can explicitly pass in self when calling them. You could automate this by making a dedicated class to hold the rule sets, and make this class a descriptor too:

import types
from collections.abc import MutableSet

class RulesSet(MutableSet):
    def __init__(self, values=(), rules=None, instance=None, owner=None):
        self._rules = rules or set()  # can be a shared set!
        self._instance = instance
        self._owner = owner
        self |= values

    def __repr__(self):
        bound = ''
        if self._owner is not None:
            bound = f', instance={self._instance!r}, owner={self._owner!r}'
        rules = ', '.join([repr(v) for v in iter(self)])
        return f'{type(self).__name__}({{{rules}}}{bound})'

    def __contains__(self, ob):
        try:
            if ob.__self__ is self._instance or ob.__self__ is self._owner:
                # test for the unbound function instead when both are bound, this requires staticmethod and classmethod to be unwrapped!
                ob = ob.__func__
                return any(ob is getattr(f, '__func__', f) for f in self._rules)
        except AttributeError:
            # not a method-like object
            pass
        return ob in self._rules

    def __iter__(self):
        if self._owner is not None:
            return (f.__get__(self._instance, self._owner) for f in self._rules)
        return iter(self._rules)

    def __len__(self):
        return len(self._rules)

    def add(self, ob):
        while isinstance(ob, Rule):
            # remove any rule wrappers
            ob = ob._function
        assert isinstance(ob, (types.FunctionType, classmethod, staticmethod))
        self._rules.add(ob)

    def discard(self, ob):
        self._rules.discard(ob)

    def __get__(self, instance, owner):
        # share the set with a new, bound instance.
        return type(self)(rules=self._rules, instance=instance, owner=owner)

class Rule:
    @classmethod
    def make_decorator(cls, rulename):
        ruleset_name = f'{rulename}_rules'
        def decorator(f):
            return cls(f, ruleset_name)
        decorator.__name__ = rulename
        return decorator

    def __init__(self, function, ruleset_name):
        self._function = function
        self._ruleset_name = ruleset_name

    def __get__(self, *args):
        # this is mostly here just to make Python call __set_name__
        return self._function.__get__(*args)

    def __set_name__(self, owner, name):
        # register, then replace the name with the original function
        # to avoid being a performance bottleneck
        ruleset = getattr(owner, self._ruleset_name, None)
        if ruleset is None:
            ruleset = RulesSet()
            setattr(owner, self._ruleset_name, ruleset)
        ruleset.add(self)
        # transfer controrol to any further rule objects
        if isinstance(self._function, Rule):
            self._function.__set_name__(owner, name)
        else:
            setattr(owner, name, self._function)

red = Rule.make_decorator('red')
blue = Rule.make_decorator('blue')
hard = Rule.make_decorator('hard')
soft = Rule.make_decorator('soft')

Then just use:

class MyClass:
    @red
    def rule_one(self):
        return 1

    @blue
    @hard
    def rule_two(self):
        return 2

    @hard
    def rule_three(self):
        return 3

    @blue
    @soft
    def rule_four(self):
        return 4

and you can access self.red_rules , etc. as a set with bound methods:

>>> inst = MyClass()
>>> inst.red_rules
RulesSet({<bound method MyClass.rule_one of <__main__.MyClass object at 0x106fe7550>>}, instance=<__main__.MyClass object at 0x106fe7550>, owner=<class '__main__.MyClass'>)
>>> inst.blue_rules
RulesSet({<bound method MyClass.rule_two of <__main__.MyClass object at 0x106fe7550>>, <bound method MyClass.rule_four of <__main__.MyClass object at 0x106fe7550>>}, instance=<__main__.MyClass object at 0x106fe7550>, owner=<class '__main__.MyClass'>)
>>> inst.hard_rules
RulesSet({<bound method MyClass.rule_three of <__main__.MyClass object at 0x106fe7550>>, <bound method MyClass.rule_two of <__main__.MyClass object at 0x106fe7550>>}, instance=<__main__.MyClass object at 0x106fe7550>, owner=<class '__main__.MyClass'>)
>>> inst.soft_rules
RulesSet({<bound method MyClass.rule_four of <__main__.MyClass object at 0x106fe7550>>}, instance=<__main__.MyClass object at 0x106fe7550>, owner=<class '__main__.MyClass'>)
>>> for rule in inst.hard_rules:
...     rule()
...
2
3

The same rules are accessible on the class; normal functions remain unbound:

>>> MyClass.blue_rules
RulesSet({<function MyClass.rule_two at 0x107077a60>, <function MyClass.rule_four at 0x107077b70>}, instance=None, owner=<class '__main__.MyClass'>)
>>> next(iter(MyClass.blue_rules))
<function MyClass.rule_two at 0x107077a60>

Containment testing works as expected:

>>> inst.rule_two in inst.hard_rules
True
>>> inst.rule_two in inst.soft_rules
False
>>> MyClass.rule_two in MyClass.hard_rules
True
>>> MyClass.rule_two in inst.hard_rules
True

You can use these rules to register classmethod and staticmethod objects too:

>>> class Foo:
...     @hard
...     @classmethod
...     def rule_class(cls):
...         return f'rule_class of {cls!r}'
...
>>> Foo.hard_rules
RulesSet({<bound method Foo.rule_class of <class '__main__.Foo'>>}, instance=None, owner=<class '__main__.Foo'>)
>>> next(iter(Foo.hard_rules))()
"rule_class of <class '__main__.Foo'>"
>>> Foo.rule_class in Foo.hard_rules
True

Answer 2

You can subclass set and give it a decorator method:

class MySet(set):
    def register(self, method):
        self.add(method)
        return method

class MyClass:
    red_rules = MySet()
    blue_rules = MySet()
    hard_rules = MySet()
    soft_rules = MySet()

    @red_rules.register
    def rule_one(self):
        return 1

    @blue_rules.register
    @hard_rules.register
    def rule_two(self):
        return 2

    @hard_rules.register
    def rule_three(self):
        return 3

    @blue_rules.register
    @soft_rules.register
    def rule_four(self):
        return 4

Or if you find using the .register method ugly, you can always define the __call__ method to use the set itself as a decorator:

class MySet(set):
    def __call__(self, method):
        """Use set as a decorator to add elements to it."""
        self.add(method)
        return method

class MyClass:
    red_rules = MySet()
    ...

    @red_rules
    def rule_one(self):
        return 1

    ...

This looks better, but it's less explicit, so for other collaborators (or future yourself) it might be harder to grasp what's happening here.

---

To call the stored functions, you can just loop over the set you want and pass in the instance as the self argument:

my_instance = MyClass()
for rule in MyClass.red_rules:
    rule(my_instance)

You can also create an utility function to do this for you, for example you can create a MySet.invoke() method:

class MySet(set):
    ...
    def invoke(self, obj):
        for rule in self:
            rule(obj)

And now just call:

MyClass.red_rules.invoke(my_instance)

Or you could have MyClass handle this instead:

class MyClass:
    ...
    def invoke_rules(self, rules):
        for rule in rules:
            rule(self)

And then call this on an instance of MyClass :

my_instance.invoke_rules(MyClass.red_rules)