Class Decorators

Part of the Object Oriented Programming section of Coddy's Python journey — lesson 40 of 64.

Class decorators allow you to modify or enhance classes by wrapping them with another function. They work like function decorators but are applied to entire classes.

Here is a simple class without decoration:

class Person:
    def __init__(self, name):
        self.name = name
        
    def greet(self):
        return f"Hello, my name is {self.name}"

Create a class decorator that adds a new method:

def add_farewell(cls):
    def farewell(self):
        return f"Goodbye from {self.name}"
    
    cls.farewell = farewell  # Add the method to the class
    return cls

Apply the decorator to a class using @:

@add_farewell
class EnhancedPerson:
    def __init__(self, name):
        self.name = name
        
    def greet(self):
        return f"Hello, my name is {self.name}"

Now the class has both original and added methods:

person = EnhancedPerson("Alice")
print(person.greet())     # Hello, my name is Alice
print(person.farewell())  # Goodbye from Alice

You can also wrap an existing method — storing the original so you can still call it inside the wrapper. This is useful when you want to add behaviour (like tracking or logging) around a method the class already has:

def add_tracking(cls):
    original_greet = cls.greet  # Store the original method

    def tracked_greet(self):
        print(f"greet was called")   # Extra behaviour before
        return original_greet(self)  # Call the original method

    cls.greet = tracked_greet  # Replace the method on the class
    return cls

@add_tracking
class TrackedPerson:
    def __init__(self, name):
        self.name = name

    def greet(self):
        return f"Hello, my name is {self.name}"

When you call greet now, the wrapper runs first, then delegates to the original:

person = TrackedPerson("Alice")
print(person.greet())
# greet was called
# Hello, my name is Alice

The key steps for wrapping an existing method are:
1. Save the original method: original = cls.method
2. Define a new function that calls original(self) plus any extra logic.
3. Assign the new function back to the class: cls.method = new_function
4. Return the modified class.

Key Point: Class decorators take a class as input, modify or enhance it, and return the modified class. They can add methods, modify attributes, or wrap existing functionality. Use @decorator_name above the class definition to apply them. This provides a clean way to extend class functionality without inheritance.

Challenge

Medium

In this challenge, you'll implement a class decorator that adds method call tracking functionality.

decorator.py - Implement the add_counter class decorator
calculator.py - Contains the classes that will use your decorator

The driver.py file contains extensive test scenarios that will validate your implementation against various use cases, inheritance patterns, and edge cases. You don't need to modify this file, but examining it will help you understand the expected behavior.

Cheat sheet

Class decorators modify or enhance classes by wrapping them with another function. They take a class as input, modify it, and return the modified class.

Basic class decorator syntax:

def decorator_name(cls):
    # Modify the class
    return cls

@decorator_name
class MyClass:
    pass

Example - adding a method to a class:

def add_farewell(cls):
    def farewell(self):
        return f"Goodbye from {self.name}"
    
    cls.farewell = farewell  # Add the method to the class
    return cls

@add_farewell
class EnhancedPerson:
    def __init__(self, name):
        self.name = name
        
    def greet(self):
        return f"Hello, my name is {self.name}"

person = EnhancedPerson("Alice")
print(person.greet())     # Hello, my name is Alice
print(person.farewell())  # Goodbye from Alice

Example - wrapping an existing method (store original, then replace):

def loud_greet(cls):
    original_greet = cls.greet  # Store the original method

    def new_greet(self):
        result = original_greet(self)  # Call the original
        return result.upper()          # Enhance the result

    cls.greet = new_greet  # Replace with the wrapped version
    return cls

@loud_greet
class EnhancedPerson:
    def __init__(self, name):
        self.name = name

    def greet(self):
        return f"Hello, my name is {self.name}"

person = EnhancedPerson("Alice")
print(person.greet())  # HELLO, MY NAME IS ALICE

Key steps for wrapping an existing method:

original = cls.method_name — save the original method
def new_method(self): original(self) — call it inside the wrapper
cls.method_name = new_method — replace the method on the class

Class decorators can add methods, modify attributes, or wrap existing functionality without using inheritance.

Try it yourself

from calculator import Calculator

# Comprehensive test case handler
test_case = input()

if test_case == "basic_functionality":
    calc = Calculator()
    print(calc.add(5, 3))  # Should print 8
    print(calc.add(2, 7))  # Should print 9
    print(calc.subtract(10, 4))  # Should print 6
    print(calc.call_counts)  # Should print {'add': 2, 'subtract': 1}

elif test_case == "multiple_instances":
    calc1 = Calculator()
    calc2 = Calculator()
    
    print(calc1.add(5, 3))  # Should print 8
    print(calc2.add(2, 7))  # Should print 9
    print(calc1.subtract(10, 4))  # Should print 6
    
    # Both instances should share the same call_counts
    print(calc1.call_counts)  # Should print {'add': 2, 'subtract': 1}
    print(calc2.call_counts)  # Should print {'add': 2, 'subtract': 1}
    print(calc1.call_counts is calc2.call_counts)  # Should print True

elif test_case == "method_call_order":
    calc = Calculator()
    
    print(calc.add(1, 2))  # Should print 3
    print(calc.subtract(5, 2))  # Should print 3
    print(calc.add(10, 20))  # Should print 30
    print(calc.add(7, 3))  # Should print 10
    print(calc.call_counts)  # Should print {'add': 3, 'subtract': 1}

elif test_case == "large_values":
    calc = Calculator()
    
    print(calc.add(1000000, 2000000))  # Should print 3000000
    print(calc.subtract(5000000, 2000000))  # Should print 3000000
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 1}

elif test_case == "negative_values":
    calc = Calculator()
    
    print(calc.add(-5, -3))  # Should print -8
    print(calc.subtract(-10, -4))  # Should print -6
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 1}

elif test_case == "float_values":
    calc = Calculator()
    
    print(calc.add(5.5, 3.2))  # Should print 8.7
    print(calc.subtract(10.5, 4.2))  # Should print 6.3
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 1}

elif test_case == "zero_values":
    calc = Calculator()
    
    print(calc.add(0, 0))  # Should print 0
    print(calc.subtract(0, 0))  # Should print 0
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 1}

elif test_case == "counter_reset":
    calc = Calculator()
    
    print(calc.add(5, 3))  # Should print 8
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 0}
    
    # Reset the counter
    calc.call_counts = {"add": 0, "subtract": 0}
    
    print(calc.add(2, 7))  # Should print 9
    print(calc.call_counts)  # Should print {'add': 0, 'subtract': 0}

elif test_case == "counter_manipulation":
    calc = Calculator()
    
    print(calc.add(5, 3))  # Should print 8
    print(calc.call_counts)  # Should print {'add': 1, 'subtract': 0}
    
    # Manipulate the counter directly
    calc.call_counts["add"] = 100
    
    print(calc.add(2, 7))  # Should print 9
    print(calc.call_counts)  # Should print {'add': 101, 'subtract': 0}

elif test_case == "performance_test":
    calc = Calculator()
    
    # Perform 1000 add operations
    for i in range(1000):
        calc.add(i, i+1)
    
    # Perform 500 subtract operations
    for i in range(500):
        calc.subtract(i+10, i)
    
    print(calc.call_counts)  # Should print {'add': 1000, 'subtract': 500}