Skip to Content
Course content

192: Class Decorators for Modifying Classes

Click on the "Edit" button in the top corner of the screen to edit your slide content.

You've probably spent a lot of time using function decorators to wrap logic around a specific call. But when you start building larger systems—think of a framework for API endpoints or a set of data models—you'll find yourself wanting to modify the class itself. Maybe you want to inject a specific method into every model, or perhaps you need to register every class that handles a certain data type into a global registry.

The Rigidity of Inheritance Mixins

The first instinct most developers have when they want to share behavior across classes is to use a Mixin. It feels clean at first. You create a base class with the shared logic and inherit from it. Let's say we're building a system where several different classes need a metadata attribute and a log_status() method for debugging.

class AuditMixin:
    def log_status(self):
        print(f"Checking status of {self.__class__.__name__}...")

class UserProfile(AuditMixin):
    def __init__(self, username):
        self.username = username

class PaymentTransaction(AuditMixin):
    def __init__(self, amount):
        self.amount = amount

This works fine until your project grows. What happens when UserProfile already needs to inherit from a complex DatabaseModel class? You're now diving into multiple inheritance and the dreaded Method Resolution Order (MRO). I've seen projects where the inheritance tree becomes a spiderweb just because someone wanted to add a simple logging utility to ten different classes. It's overkill. You're forcing a "is-a" relationship when what you actually want is a "has-this-behavior" capability.

Decoupling via Class Decorators

This is where class decorators come in. Unlike function decorators, which usually wrap a function in another function, a class decorator takes a class as an argument and returns a class. It allows you to modify the class definition dynamically after it's been created but before it's ever used by the rest of your app.

Instead of forcing UserProfile to be a child of AuditMixin, we can just "decorate" it. Here is how I'd handle the same auditing logic without touching the inheritance chain:

def add_audit_logging(cls):
    # We can inject new methods directly into the class dictionary
    def log_status(self):
        print(f"Checking status of {self.__class__.__name__}...")
    
    cls.log_status = log_status
    cls._is_audited = True # Add some metadata for the system to find later
    return cls

@add_audit_logging
class UserProfile:
    def __init__(self, username):
        self.username = username

@add_audit_logging
class PaymentTransaction:
    def __init__(self, amount):
        self.amount = amount

Notice that UserProfile is now a plain class. It doesn't inherit from anything. But because of the decorator, it possesses the log_status method. I prefer this because it keeps the class hierarchy flat. If I decide later that PaymentTransaction shouldn't be audited, I just delete one line of code rather than restructuring the entire object model.

When This Breaks and How to Handle It

There is a catch. If your decorator returns a new class (like a wrapper) instead of modifying and returning the original class, you'll break isinstance() checks. If the decorator returns a wrapper, isinstance(user, UserProfile) will return False because the object is now an instance of the wrapper, not the original class.

For 90% of class modification tasks, modifying the class in place (as I did above) is the right move. It's fast, it preserves the identity of the class, and it's transparent. You only need to return a wrapper if you intend to completely intercept how the class is instantiated—but that's a different architectural problem entirely. For modifying attributes or adding methods, stick to the "mutate and return" pattern.




📋 Practical Task

Building a Plugin Registry Decorator

Imagine you are building a plugin system where different classes handle different file formats (e.g., JsonHandler, XmlHandler, CsvHandler). Instead of manually adding these classes to a list in your main script, you should create a class decorator called register_plugin.

Your Requirements:

  • Create a global list called PLUGIN_REGISTRY.
  • Write a class decorator register_plugin that adds the decorated class to the PLUGIN_REGISTRY list.
  • The decorator should also add a class attribute plugin_enabled = True to the class it decorates.
  • Create three different handler classes (e.g., JsonHandler, XmlHandler, CsvHandler) and apply the decorator to each.
  • Finally, print the PLUGIN_REGISTRY to verify that all three classes were automatically registered.
Rating
0 0

There are no comments for now.

to be the first to leave a comment.