Use Python Decorators to Flag Classes for Protocol Compliance

How Can I Use a Decorator to Flag Classes as Passing a Certain Protocol?

Hello, my dear friends! Welcome to this super-duper exciting blog post where we’re diving deep into the world of Python programming, but with a little bit of that Indian tadka to make it fun and easy to understand. Today, we’re talking about something very cool: how to use a decorator to flag classes as passing a certain protocol. Now, don’t worry if that sounds like a mouthful. By the end of this post, you’ll be like, “Arre, yeh toh badiya hai!” (Oh, this is awesome!)

What’s This Blog All About?

Before we jump into the code, let’s understand what we’re trying to do. In Python, a protocol is like a set of rules that a class needs to follow. Think of it like a checklist: if a class has certain methods, it “passes” the protocol. For example, if a class has a __len__ method, it follows the Sized protocol because it can tell you how big it is.

Now, how do we check if a class follows a protocol? And how do we flag it in a clean, reusable way? That’s where decorators come in. A decorator is like a rubber stamp that we can slap onto a class to add some extra functionality—like saying, “This class passes this protocol, yaar!”

In this blog, we’re going to:

• Learn what protocols are in Python.
• Understand what decorators are and how they work.
• Figure out how to create a decorator that checks if a class follows a protocol.
• Write code examples to flag classes that pass a protocol.
• Explore real-world use cases and why this is useful.
• Have some fun along the way with simple explanations and a bit of Indian flair!

By the end, you’ll know how to use decorators to mark classes as passing protocols, and you’ll feel like a Python rockstar. Chalo, shuru karte hain! (Let’s get started!)

---

What Are Protocols in Python?

To understand our main topic, we first need to know what a protocol is. In Python, a protocol is a way to define what a class should be able to do without forcing it to inherit from a specific base class. It’s like saying, “If you have these methods, you’re good to go!”

For example, let’s say you have a class that represents a list of students in a school. If this class has a __len__ method, it can tell you how many students are there. Python says, “Oh, this class follows the Sized protocol because it has __len__!” Similarly, if a class has __iter__, it follows the Iterable protocol, meaning you can loop over it.

Protocols are part of Python’s structural subtyping system, which is a fancy way of saying, “If it walks like a duck and quacks like a duck, it’s a duck.” This is different from nominal subtyping, where a class has to explicitly say, “I’m a subclass of Duck.” In Python, protocols are all about what a class does, not what it is.

Example of a Protocol

Let’s write a simple example to understand this better. Suppose we want a class to follow a protocol that says it should have a greet method. Here’s how it might look:

class Student:
    def greet(self):
        return "Hello, I'm a student!"

class Teacher:
    def greet(self):
        return "Hello, I'm a teacher!"

# Both classes have a `greet` method, so they follow the "Greeter" protocol
student = Student()
teacher = Teacher()

print(student.greet())  # Output: Hello, I'm a student!
print(teacher.greet())  # Output: Hello, I'm a teacher!

In this example, both Student and Teacher have a greet method, so they follow our imaginary “Greeter” protocol. But how do we check this programmatically? And how do we flag these classes as passing the protocol? That’s where decorators come in, but first, let’s learn about decorators.

---

What Are Decorators in Python?

Arre, decorators sound like some fancy interior designer thing, na? But in Python, they’re super useful and not that complicated. A decorator is a function that takes another function (or class) and adds some extra functionality to it. It’s like putting a shiny wrapper around a gift to make it look better.

In Python, decorators are used with the @ symbol. You’ve probably seen them before, like @staticmethod or @classmethod. But you can create your own decorators too!

How Do Decorators Work?

Let’s break it down with a simple example. Suppose we want to add some extra behavior to a function, like printing “Function is running!” before it runs. Here’s how we can do it with a decorator:

def my_decorator(func):
    def wrapper():
        print("Function is running!")
        func()
    return wrapper

@my_decorator
def say_hello():
    print("Hello, world!")

say_hello()

Output:

Function is running!
Hello, world!

What’s happening here?

1. my_decorator is a function that takes another function (func) as an argument.
2. Inside my_decorator, we define a wrapper function that adds extra behavior (printing “Function is running!”) and then calls the original function.
3. my_decorator returns the wrapper function.
4. When we put @my_decorator above say_hello, it’s like saying, “Take say_hello, wrap it with my_decorator, and replace say_hello with the wrapped version.”

This is how decorators work for functions. But guess what? Decorators can also work on classes! That’s what we’ll use to flag classes as passing a protocol.

---

Creating a Decorator to Flag Classes

Now that we know what protocols and decorators are, let’s get to the main dish: using a decorator to flag classes as passing a certain protocol. Our goal is to create a decorator that checks if a class has certain methods (defined by a protocol) and marks it as passing that protocol.

Step 1: Defining a Protocol

In Python, we can define protocols using the typing.Protocol class from the typing module (introduced in Python 3.8). This lets us formally specify what methods a class needs to have to follow a protocol. Let’s create a simple protocol called GreeterProtocol that requires a greet method.

from typing import Protocol

class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

Here, GreeterProtocol says, “Any class that has a greet method returning a string follows this protocol.” The ... (ellipsis) is a placeholder indicating that this method is not implemented in the protocol itself.

Step 2: Creating the Decorator

Now, let’s create a decorator that checks if a class follows GreeterProtocol and flags it by adding an attribute, like is_greeter = True. Here’s the code:

from typing import Protocol, Type, Any
from functools import wraps

class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

def greeter_protocol_decorator(cls: Type[Any]) -> Type[Any]:
    @wraps(cls)
    def wrapper(*args, **kwargs):
        return cls(*args, **kwargs)
    
    # Check if the class has a `greet` method
    if not hasattr(cls, 'greet'):
        raise TypeError(f"Class {cls.__name__} does not implement the GreeterProtocol (missing 'greet' method)")
    
    # Add a flag to indicate the class follows the protocol
    setattr(cls, 'is_greeter', True)
    
    return wrapper

# Example usage
@greeter_protocol_decorator
class Student:
    def greet(self):
        return "Hello, I'm a student!"

@greeter_protocol_decorator
class Teacher:
    def greet(self):
        return "Hello, I'm a teacher!"

# This class will raise an error
@greeter_protocol_decorator
class Robot:
    def move(self):
        return "Beep boop, I move!"

# Testing
student = Student()
teacher = Teacher()

print(student.greet())  # Output: Hello, I'm a student!
print(teacher.greet())  # Output: Hello, I'm a teacher!
print(Student.is_greeter)  # Output: True
print(Teacher.is_greeter)  # Output: True

What’s Happening Here?

1. We define GreeterProtocol as before.
2. We create a decorator called greeter_protocol_decorator that takes a class (cls) as an argument.
3. Inside the decorator, we:
   • Use @wraps(cls) to preserve the class’s metadata.
   • Define a wrapper function that creates instances of the class.
   • Check if the class has a greet method using hasattr.
   • If the method is missing, raise a TypeError.
   • If the method exists, add an is_greeter attribute to the class to flag it as passing the protocol.
4. We apply the decorator to Student and Teacher, which both have greet methods, so they pass.
5. We apply it to Robot, which doesn’t have a greet method, so it raises an error.

When you run this code, Student and Teacher will work fine, and you can check Student.is_greeter or Teacher.is_greeter to confirm they’re flagged. But Robot will throw an error because it doesn’t follow the protocol.

Step 3: Making the Decorator Reusable

The decorator above works for GreeterProtocol, but what if we want to check for other protocols? Let’s make a generic decorator that can work with any protocol. Here’s the improved version:

from typing import Protocol, Type, Any, get_type_hints
from functools import wraps
from inspect import getmembers, ismethod

def protocol_decorator(protocol: Type[Protocol]) -> callable:
    def decorator(cls: Type[Any]) -> Type[Any]:
        @wraps(cls)
        def wrapper(*args, **kwargs):
            return cls(*args, **kwargs)
        
        # Get the protocol's required methods
        required_methods = [
            name for name, method in getmembers(protocol, ismethod)
            if not name.startswith('__')
        ]
        
        # Check if the class implements all required methods
        for method_name in required_methods:
            if not hasattr(cls, method_name):
                raise TypeError(
                    f"Class {cls.__name__} does not implement the {protocol.__name__} "
                    f"(missing '{method_name}' method)"
                )
        
        # Add a flag to indicate the class follows the protocol
        flag_name = f"is_{protocol.__name__.lower()}"
        setattr(cls, flag_name, True)
        
        return wrapper
    return decorator

# Define a Greeter protocol
class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

# Define a Mover protocol
class MoverProtocol(Protocol):
    def move(self) -> str:
        ...

# Apply the decorator
@protocol_decorator(GreeterProtocol)
class Student:
    def greet(self):
        return "Hello, I'm a student!"

@protocol_decorator(MoverProtocol)
class Robot:
    def move(self):
        return "Beep boop, I move!"

# Testing
student = Student()
robot = Robot()

print(student.greet())  # Output: Hello, I'm a student!
print(robot.move())     # Output: Beep boop, I move!
print(Student.is_greeterprotocol)  # Output: True
print(Robot.is_moverprotocol)      # Output: True

What’s New Here?

1. We created a protocol_decorator that takes a protocol (like GreeterProtocol) as an argument.
2. It returns a decorator function that checks if the class implements all the methods required by the protocol.
3. We use getmembers and ismethod from the inspect module to find the protocol’s methods.
4. For each required method, we check if the class has it using hasattr.
5. We set a dynamic flag (e.g., is_greeterprotocol or is_moverprotocol) to mark the class.
6. The decorator is now reusable for any protocol, like GreeterProtocol or MoverProtocol.

This is super flexible, na? You can define any protocol and use the same decorator to flag classes that follow it.

---

Why Is This Useful?

Arre, you might be thinking, “Bhai, yeh sab kyun karna hai?” (Why do all this?) Well, using decorators to flag classes as passing protocols is super useful in real-world projects. Here are some reasons why:

1. Type Safety: By checking protocols at runtime, you ensure classes have the right methods, reducing bugs.
2. Code Readability: The decorator clearly shows which classes follow which protocols, making your code easier to understand.
3. Reusability: A generic decorator can be used for any protocol, saving you time.
4. Documentation: The flag (like is_greeterprotocol) acts like a label, telling other developers what the class can do.
5. Testing: You can check the flag in tests to verify that classes implement protocols correctly.

Real-World Example

Imagine you’re building a school management system. You have different classes like Student, Teacher, and Principal, and you want to ensure that certain classes can send notifications (e.g., via a send_notification method). You can define a NotifierProtocol and use a decorator to flag classes that implement it:

class NotifierProtocol(Protocol):
    def send_notification(self, message: str) -> None:
        ...

@protocol_decorator(NotifierProtocol)
class Teacher:
    def send_notification(self, message: str) -> None:
        print(f"Teacher sends: {message}")

@protocol_decorator(NotifierProtocol)
class Principal:
    def send_notification(self, message: str) -> None:
        print(f"Principal sends: {message}")

teacher = Teacher()
principal = Principal()

teacher.send_notification("Exam tomorrow!")  # Output: Teacher sends: Exam tomorrow!
principal.send_notification("School closed!")  # Output: Principal sends: School closed!
print(Teacher.is_notifierprotocol)  # Output: True
print(Principal.is_notifierprotocol)  # Output: True

This ensures that only classes with send_notification can be flagged as notifiers, keeping your system clean and safe.

---

Handling Edge Cases

To make our decorator robust, we need to handle some edge cases. Let’s discuss a few and update our code to deal with them.

Edge Case 1: Methods with Wrong Signatures

Our current decorator checks if a method exists but doesn’t verify its signature (e.g., parameters or return type). For example, a class might have a greet method that takes no arguments but returns an integer instead of a string. Let’s improve the decorator to check method signatures using type hints.

from typing import Protocol, Type, Any, get_type_hints
from functools import wraps
from inspect import getmembers, ismethod, signature

def protocol_decorator(protocol: Type[Protocol]) -> callable:
    def decorator(cls: Type[Any]) -> Type[Any]:
        @wraps(cls)
        def wrapper(*args, **kwargs):
            return cls(*args, **kwargs)
        
        # Get the protocol's required methods
        protocol_methods = {
            name: method for name, method in getmembers(protocol, ismethod)
            if not name.startswith('__')
        }
        
        # Get type hints and signatures for protocol methods
        protocol_hints = get_type_hints(protocol)
        
        # Check if the class implements all required methods
        for method_name in protocol_methods:
            if not hasattr(cls, method_name):
                raise TypeError(
                    f"Class {cls.__name__} does not implement the {protocol.__name__} "
                    f"(missing '{method_name}' method)"
                )
            
            # Check method signature
            cls_method = getattr(cls, method_name)
            protocol_method = protocol_methods[method_name]
            
            cls_sig = signature(cls_method)
            protocol_sig = signature(protocol_method)
            
            if cls_sig != protocol_sig:
                raise TypeError(
                    f"Class {cls.__name__}'s '{method_name}' method has incorrect signature. "
                    f"Expected {protocol_sig}, got {cls_sig}"
                )
            
            # Check return type
            if method_name in protocol_hints:
                cls_hints = get_type_hints(cls_method)
                expected_return = protocol_hints[method_name]
                actual_return = cls_hints.get('return', Any)
                
                if actual_return != expected_return:
                    raise TypeError(
                        f"Class {cls.__name__}'s '{method_name}' method has incorrect return type. "
                        f"Expected {expected_return}, got {actual_return}"
                    )
        
        # Add flag to indicate the class follows the protocol
        flag_name = f"is_{protocol.__name__.lower()}"
        setattr(cls, flag_name, True)
        
        return wrapper
    
    return decorator

# Example usage
class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

@protocol_decorator(GreeterProtocol)
class Student:
    def greet(self) -> str:
        return 42  # Error: Class Student's greet method has incorrect return type. Expected str, got int

@protocol_decorator(GreeterProtocol)
class Teacher:
    def greet(self, extra: int) -> str:
        return "Hello, I'm a teacher!"  # Error: Signature mismatch

@protocol_decorator(GreeterProtocol)
class Principal:
    def greet(self) -> str:
        return "Hello, I'm a principal!"  # Works fine

# Testing
principal = Principal()
print(principal.greet())  # Output: Hello, I'm a principal!
print(Principal.is_greeterprotocol)  # Output: True

This version checks:

1. If the method exists.
2. If the method’s signature (parameters) matches the protocol’s.
3. If the method’s return type matches the protocol’s.

Edge Case 2: Inherited Methods

What if a class inherits a method from a parent class? Our decorator should still recognize it. Let’s test this:

@protocol_decorator(GreeterProtocol)
class Person:
    def greet(self) -> str:
        return "Hello, I'm a person!""

class Student(Person):
    pass

print(Student.is_greeterprotocol)  # Output: True
print(Student().greet())  # Output: Hello, I'm a person!"

Good news: Our decorator already handles this because hasattr checks inherited attributes. So, no changes needed here!

Edge Case 3: Multiple Protocols

What if a class needs to follow multiple protocols? We can apply multiple decorators:

class NotifierProtocol(Protocol):
    def send_notification(self, message: str) -> None:
        ...

@protocol_decorator(GreeterProtocol)
@protocol_decorator(NotifierProtocol)
class Teacher:
    def greet(self) -> str:
        return "Hello, I'm a teacher!""
    def send_notification(self, message: str) -> None:
        print(f"Teacher sends: {message}")

print(Teacher.is_greeterprotocol)  # Output: True
print(Teacher.is_notifierprotocol)  # Output: True

This works because each decorator adds its own flag. But we could also create a decorator that checks multiple protocols at once, which we’ll explore in the next chapter.

---

Advanced Decorator: Checking Multiple Protocols

To make things even cooler, let’s create a decorator that can check multiple protocols at once. This way, we can say, “This class follows both GreeterProtocol and NotifierProtocol!” in one go.

from typing import Protocol, Type, List, Any, get_type_hints
from functools import wraps
from inspect import getmembers, ismethod, signature

def multi_protocol_decorator(protocols: List[Type[Protocol]]) -> callable:
    def decorator(cls: Type[Any]) -> Type[Any]:
        @wraps(cls)
        def wrapper(*args, **kwargs):
            return cls(*args, **kwargs)
        
        # Check each protocol
        for protocol in protocols:
            protocol_methods = {
                name: method for name, method in getmembers(protocol, ismethod)
                if not name.startswith('__')
            }
            protocol_hints = get_type_hints(protocol)
            
            # Check required methods
            for method_name in protocol_methods:
                if not hasattr(cls, method_name):
                    raise TypeError(
                        f"Class {cls.__name__} does not implement the {protocol.__name__} "
                        f"(missing '{method_name}' method)"
                    )
                
                # Check signature
                cls_method = getattr(cls, method_name)
                protocol_method = protocol_methods[method_name]
                
                cls_sig = signature(cls_method)
                protocol_sig = signature(protocol_method)
                
                if cls_sig != protocol_sig:
                    raise TypeError(
                        f"Class {cls.__name__}'s '{method_name}' method has incorrect signature. "
                        f"Expected {protocol_sig}, got {cls_sig}"
                    )
                
                # Check return type
                if method_name in protocol_hints:
                    cls_hints = get_type_hints(cls_method)
                    expected_return = protocol_hints[method_name]
                    actual_return = cls_hints.get('return', Any)
                    
                    if actual_return != expected_return:
                        raise TypeError(
                            f"Class {cls.__name__}'s '{method_name}' method has incorrect return type. "
                            f"Expected {expected_return}, got {actual_return}"
                        )
            
            # Add flag
            flag_name = f"is_{protocol.__name__.lower()}"
            setattr(cls, flag_name, True)
        
        return wrapper
    return decorator

# Define protocols
class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

class NotifierProtocol(Protocol):
    def send_notification(self, message: str) -> None:
        ...

# Apply decorator
@multi_protocol_decorator([GreeterProtocol, NotifierProtocol])
class Principal:
    def greet(self) -> str:
        return "Hello, I'm a principal!"
    def send_notification(self, message: str) -> None:
        print(f"Principal sends: {message}")

# Testing
principal = Principal()
principal.greet()  # Output: Hello, I'm a principal!
principal.send_notification("School closed!")  # Output: Principal sends: School closed!
print(Principal.is_greeterprotocol)  # Output: True
print(Principal.is_protocolprotocol)  # Output: True

This multi_protocol_decorator:

• Takes a list of protocols.
• Checks if the class implements each protocol’s methods.
• Adds a flag for each protocol the class passes.

This is super clean and lets you bundle multiple protocol checks into one decorator. Matlab, ekdum mast solution hai! (What a great solution!)

---

Performance Considerations

Before we wrap up, let’s talk about performance. Our decorator does a lot of checks at runtime, like inspecting methods, type hints, and signatures. For small projects, this is fine, but for large projects with many classes, it could slow things down. Here’s how to optimize:

1. Cache Results: Cache the protocol checks so they run only once per class.
2. Static Analysis: Use a static type checker like mypy to catch protocol errors at compile-time instead of runtime.
3. Simplify Checks: If you don’t need signature or type hint validation, simplify the decorator to just check only hasattr.

Here’s a lightweight version of the decorator for performance:

from typing import Protocol, Type, List, Any
from functools import wraps

def lightweight_protocol_decorator(protocol: Type[Protocol]) -> callable:
    def decorator(cls: Type[Any]) -> Type[Any]:
        @wraps(cls)
        def wrapper(*args, **kwargs):
            return cls(*args, **kwargs)
        
        # Check only method presence
        protocol_methods = [name for name in dir(protocol) if not name.startswith('__')]
        
        for method_name in protocol_methods:
            if not hasattr(cls, method_name):
                raise TypeError(
                    f"Class {cls.__name__} does not implement {protocol.__name__} "
                    f"(missing '{method_name}' method)"
                    )
        
        # Add flag
        flag_name = f"is_{protocol.__name__.lower()}"
        setattr(cls, flag_name, True)
        
        return wrapper
    return decorator

class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

@lightweight_protocol_decorator(GreeterProtocol)
class Teacher:
    def greet(self) -> str:
        return "Hello, I'm a teacher!"

print(Teacher.is_greeterprotocol)  # Output: True

This version skips signature and type hint checks, making it faster but less strict.

---

Wrapping Up and Final Code

Wow, we’ve covered a lot, na? We learned how to use decorators to flag classes as passing protocols, created reusable and robust code, handled edge cases, and even optimized for performance. Let’s put together the final version of our code for you to use:

from typing import Protocol, Type, List, Any, get_type_hints
from functools import wraps
from inspect import getmembers, ismethod, signature

def protocol_decorator(protocol: Type[Protocol]) -> callable:
    def decorator(cls: Type[Any]) -> Type[Any]:
        @wraps(cls)
        def wrapper(*args, **kwargs):
            return cls(*args, **kwargs)
        
        # Get protocol methods
        protocol_methods = {
            name: method for name, method in getmembers(protocol, ismethod)
            if not name.startswith('__')
        }
        protocol_hints = get_type_hints(protocol)
        
        # Check methods
        for method_name in protocol_methods:
            if not hasattr(cls, method_name):
                raise TypeError(
                    f"Class {cls.__name__} does not implement {protocol.__name__} "
                    f"(missing '{method_name}' method)"
                )
            
            cls_method = getattr(cls, method_name)
            protocol_method = protocol_methods[method_name]
            
            cls_sig = signature(cls_method)
            protocol_sig = signature(protocol_method)
            
            if cls_sig != protocol_sig:
                raise TypeError(
                    f"Class {cls.__name__}'s '{method_name}' method has incorrect signature. "
                    f"Expected {protocol_sig}, got {cls_sig}"
                )
            
            if method_name in protocol_hints:
                cls_hints = get_type_hints(cls_method)
                expected_return = protocol_hints[method_name]
                actual_return = cls_hints.get('return', Any)
                
                if actual_return != expected_return:
                    raise TypeError(
                        f"Class {cls.__name__}'s '{method_name}' method has incorrect return type. "
                        f"Expected {expected_return}, got {actual_return}"
                    )
        
        # Add flag
        flag_name = f"is_{protocol.__name__.lower()}"
        setattr(cls, flag_name, True)
        
        return wrapper
    return decorator

# Protocols
class GreeterProtocol(Protocol):
    def greet(self) -> str:
        ...

class NotifierProtocol(Protocol):
    def send_notification(self, message: str) -> None:
        ...

# Classes
@protocol_decorator(GreeterProtocol)
class Student:
    def greet(self) -> str:
        return "Hello, I'm a student!"

@protocol_decorator(NotifierProtocol)
class Teacher:
    def send_notification(self, message: str) -> None:
        print(f"Teacher sends: {message}")

@protocol_decorator(GreeterProtocol)
@protocol_decorator(NotifierProtocol)
class Principal:
    def greet(self) -> str:
        return "Hello, I'm a principal!"
    def send_notification(self, message: str) -> None:
        print(f"Principal sends: {message}")

# Testing
student = Student()
teacher = Teacher()
principal = Principal()

print(student.greet())  # Output: Hello, I'm a student!
teacher.send_notification("Exam tomorrow!")  # Output: Teacher sends: Exam tomorrow!
principal.greet()  # Output: Hello, I'm a principal!
principal.send_notification("School closed!")  # Output: Principal sends: School closed!
print(Student.is_greeterprotocol)  # Output: True
print(Teacher.is_notifierprotocol)  # Output: True
print(Principal.is_greeterprotocol)  # Output: True
print(Principal.is_notifierprotocol)  # Output: True

This code is robust, reusable, and handles all the cases we discussed. You can use it in your own projects to flag classes as passing protocols.

---

Conclusion

Phew, what a journey, bhai! We started with a simple idea—using a decorator to flag classes as passing a protocol—and ended up with a super powerful solution. You’ve learned:

• What protocols are and why they’re useful.
• How decorators work in Python.
• How to create a decorator to check if a class follows a protocol.
• How to make the decorator reusable for any protocol.
• How to handle edge cases and optimize performance.
• Real-world examples to make it all practical.

I hope you’re feeling like a Python ninja now! Decorators and protocols are like masala and tadka in cooking—they add flavor and make things awesome. Keep practicing, keep coding, and don’t forget to have fun along the way.

If you have any doubts, ping me in the comments, and I’ll explain in my desi style. Abhi ke liye, bas yahi kahunga: “Code karo, aur khush raho!” (Keep coding and stay happy!)

Happy Coding, Dosto! 😎