Custom Decorators

James wants every Note method call logged during debugging. He starts adding print statements:

class Note:
    def add_tag(self, tag: str) -> None:
        print(f"Calling add_tag with {tag}")  # ← debugging
        if tag not in self.tags:
            self.tags.append(tag)

    def remove_tag(self, tag: str) -> None:
        print(f"Calling remove_tag with {tag}")  # ← debugging
        if tag not in self.tags:
            raise ValueError(f"Tag '{tag}' not found")
        self.tags.remove(tag)

Two methods, two identical print lines. If he has ten methods, he writes ten identical lines. When debugging is done, he deletes ten lines. When debugging starts again, he adds them back.

"That is cross-cutting logic," Emma says. "The logging behavior has nothing to do with tags. It applies to every method the same way. Write it once, apply it everywhere. That is a decorator."

---

If you're new to programming

A decorator is a function that takes another function as input and returns a new function that adds behavior. The @decorator syntax is a shortcut: instead of calling add_tag = log_calls(add_tag) manually, you write @log_calls above the function definition.

If you've coded before

You have already used decorators: @dataclass transforms a class, @property transforms a method into a descriptor, @pytest.fixture registers a function as a test fixture. Writing a custom decorator means creating a function that accepts a callable, wraps it in a new callable, and returns the wrapper. functools.wraps copies __name__, __doc__, and other attributes from the original to the wrapper.

---

Decorators You Already Know

You have used decorators throughout this book without writing your own:

Decorator	Chapter	What It Does
@dataclass	51	Auto-generates __init__, __repr__, __eq__ for a class
@pytest.fixture	52	Registers a function as a test fixture
@property	61 L1	Turns a method into a computed attribute
@staticmethod	61 L2	Makes a method callable without self or cls
@classmethod	61 L2	Passes cls as the first argument

Every one of these follows the same pattern: take a function (or class), modify its behavior, return the result. Now you will write your own.

---

Writing @log_calls

A decorator is a function that wraps another function:

import functools

def log_calls(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}")
        result = func(*args, **kwargs)
        print(f"Finished {func.__name__}")
        return result
    return wrapper

Apply it:

@log_calls
def greet(name: str) -> str:
    return f"Hello, {name}"

print(greet("James"))

Output:

Calling greet
Finished greet
Hello, James

What Just Happened

The @log_calls line is equivalent to writing:

def greet(name: str) -> str:
    return f"Hello, {name}"

greet = log_calls(greet)  # ← this is what @log_calls does

1. Python passes greet to log_calls
2. log_calls creates wrapper, which calls the original greet inside
3. log_calls returns wrapper
4. greet now points to wrapper
5. Calling greet("James") runs wrapper("James"), which prints, calls the original, prints again, and returns the result

James's analogy: a quality inspection stamp. Every product coming off the line gets the same stamp process, regardless of what the product is. The stamp (decorator) is applied once to the production line (function definition), and every unit (function call) gets stamped automatically.

---

Why functools.wraps Matters

Without @functools.wraps(func), the wrapper replaces the original function's identity:

def log_calls_bad(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

@log_calls_bad
def greet(name: str) -> str:
    """Say hello."""
    return f"Hello, {name}"

print(greet.__name__)
print(greet.__doc__)

Output:

wrapper
None

The function's name and docstring are gone. functools.wraps copies them from the original:

@log_calls  # uses @functools.wraps(func)
def greet(name: str) -> str:
    """Say hello."""
    return f"Hello, {name}"

print(greet.__name__)
print(greet.__doc__)

Output:

greet
Say hello.

Rule: always use @functools.wraps(func) in your wrapper. It costs nothing and prevents debugging confusion.

---

Writing @validate_input for SmartNotes

Now write a decorator that validates string arguments are non-empty:

import functools

def validate_input(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        for arg in args[1:]:  # skip self (first argument)
            if isinstance(arg, str) and not arg.strip():
                raise ValueError(
                    f"{func.__name__}: empty string argument"
                )
        return func(*args, **kwargs)
    return wrapper

Apply it to Note methods:

class Note:
    def __init__(self, title: str, body: str, author: str = "Anonymous") -> None:
        self.title = title
        self.body = body
        self.author = author
        self.tags: list[str] = []

    @validate_input
    def add_tag(self, tag: str) -> None:
        if tag not in self.tags:
            self.tags.append(tag)

    @validate_input
    def remove_tag(self, tag: str) -> None:
        if tag not in self.tags:
            raise ValueError(f"Tag '{tag}' not found")
        self.tags.remove(tag)

note = Note("Test", "Body")
note.add_tag("python")
print(note.tags)
note.add_tag("")  # raises ValueError

Output:

['python']
ValueError: add_tag: empty string argument

One decorator, applied to every method that takes string input. Write once, use everywhere.

---

PRIMM-AI+ Practice: Decorator Tracing

Predict [AI-FREE]

Press Shift+Tab to enter Plan Mode before predicting.

Read this code and predict the exact output. Write your prediction and a confidence score from 1 to 5.

import functools

def count_calls(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        wrapper.call_count += 1
        print(f"{func.__name__} called {wrapper.call_count} time(s)")
        return func(*args, **kwargs)
    wrapper.call_count = 0
    return wrapper

@count_calls
def add(a: int, b: int) -> int:
    return a + b

result1 = add(2, 3)
result2 = add(10, 20)
print(f"Results: {result1}, {result2}")
print(f"Total calls: {add.call_count}")

Check your prediction

add called 1 time(s)
add called 2 time(s)
Results: 5, 30
Total calls: 2

The decorator adds a call_count attribute to the wrapper function. Each call increments it. Because functools.wraps preserves the name, the print shows add (not wrapper). The count persists across calls because wrapper is a single function object.

Run

Press Shift+Tab to exit Plan Mode.

Create a file called decorator_practice.py. Write the count_calls decorator and add function. Then apply @count_calls to a second function multiply(a, b). Verify that each function tracks its own count independently.

Investigate

In Claude Code, type:

/investigate @decorator_practice.py

Ask: "Why does each decorated function have its own call_count? Where is the count stored? What would happen if I stored the count as a global variable instead of on wrapper?"

---

Try With AI

Prompt 1: Understand the Wrapping Mechanism

I know that @log_calls above a function is the same
as func = log_calls(func). Walk me through what happens
step by step when Python processes this code:

@log_calls
def add_tag(self, tag: str) -> None:
    self.tags.append(tag)

What does log_calls receive? What does it return?
What does add_tag point to after the decoration?

Read the AI's explanation. Verify it matches what you learned: log_calls receives the original add_tag, creates wrapper, and returns wrapper. After decoration, add_tag points to wrapper.

What you're learning: You are building the mechanical understanding of decoration. The @ syntax is convenience; the real operation is function replacement through assignment.

Prompt 2: Decorator Use Cases

Give me five real-world use cases for custom decorators
in Python applications. For each, show a minimal
implementation (under 10 lines) and explain what
cross-cutting concern it addresses.

What you're learning: Decorators solve cross-cutting concerns: logging, timing, caching, authentication, validation. Seeing multiple use cases makes the pattern feel general, not specific to SmartNotes. Emma learns from this too: "I had not thought about using decorators for rate limiting. That is clever."

Prompt 3: Decorator for Your Domain

I work in [describe your professional domain]. I have
functions that need [describe a repeated behavior:
audit logging, input validation, retry on failure,
permission checking, etc.].

Write a decorator that handles this cross-cutting concern.
Show it applied to a realistic function from my domain.

What you're learning: You are transferring the decorator pattern to your own context. The AI generates the decorator, but you evaluate whether it correctly wraps the function, uses functools.wraps, and handles the cross-cutting concern appropriately.

---

James looks at the @validate_input decorator on add_tag and remove_tag. "I wrote the validation once, and it applies to both methods. If I add a third method that takes strings, I add one line: @validate_input."

Emma nods. "Think of it like your quality inspection stamp. Every product gets inspected the same way. You do not write a new inspection process for each product type. You write the process once and stamp it on every line."

"That is exactly what I was thinking," James says.

Emma pauses. "Honestly, I got the stamp analogy from you. I usually explain decorators as wrappers or envelopes, but the stamp is better. It captures the 'applied once, runs every time' behavior. The wrapper analogy makes people think the original function is gone. It is not gone. It is inside the stamp."

"Decorators modify behavior. Properties compute values. Static and class methods organize code. But there is still a problem. I want to test SmartNotes with a fake database instead of a real one. How do I swap the storage without changing the class?"

"That is a Protocol. Next lesson."