Adding Type Hints to Decorators and Generators

• Decorators and generators are advanced constructs that require specialized type hints to make their transformations and data flows explicit.
• Properly typed decorators allow MyPy to understand how they preserve or change function signatures.
• Typed generators clarify the types of values yielded, values accepted via .send(), and final return values.

Typing Decorators

• Decorators take a function (Callable) and return a new function; using Callable[..., Any] types them broadly but loses specific signature information.
• To preserve the original function’s signature, define a TypeVar bound to Callable[..., Any] and use it for both the decorator’s input and output types.
• Inside the decorator, the wrapper can use *args: Any, **kwargs: Any -> Any, while TypeVar ensures the decorated function’s overall type remains correct.

Typing Generators

• Use Generator[YieldType, SendType, ReturnType] to specify a generator’s yield type, the type accepted by .send(), and its return type on completion.
• If a generator does not use send(), set SendType to None; if it has no explicit return, set ReturnType to None.
• The count_up generator is typed as Generator[int, None, str], yielding integers and returning a string message.
• The accumulate_and_send generator is typed as Generator[float, float, None], yielding a running total, accepting floats via send(), and returning nothing.

Iterable & Iterator

• For functions that consume sequences of items, use Iterable[T] to accept any iterable of T (lists, tuples, generators).
• Use Iterator[T] when a function specifically expects an iterator object supporting __next__().

from typing import (
    Callable,
    Any,
    TypeVar,
    ParamSpec,
    Generator,
    Iterable,
)
import functools

# Section: Typing Decorators (simple_logging_decorator)

def simple_logging_decorator(
    func: Callable[..., Any],
) -> Callable[..., Any]:
    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        print(f"LOG: Calling {func.__name__}")
        result = func(*args, **kwargs)
        print(f"LOG: {func.__name__} returned {result}")

        return result

    return wrapper

@simple_logging_decorator
def add(x: int, y: int) -> int:
    return x + y

result_add = add(3, 5)

# Section: Typing Decorators (better_logging_decorator with TypeVar)

P = ParamSpec("P")
R = TypeVar("R")

def better_logging_decorator(
    func: Callable[P, R],
) -> Callable[P, R]:
    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
        print(f"LOG: Calling {func.__name__}")
        result = func(*args, **kwargs)
        print(f"LOG: {func.__name__} returned {result}")

        return result

    return wrapper

@better_logging_decorator
def subtract(x: int, y: int) -> int:
    return x - y

result_subtract = subtract(3, 5)

# Section: Typing Generators

def count_up_to(limit: int) -> Generator[int, None, str]:
    for i in range(limit):
        yield i

    return "Counting complete!"

def accumulate_and_send() -> (
    Generator[float, float | None, None]
):
    total = 0.0

    try:
        while True:
            sent = yield total

            if sent:
                total += sent
    except GeneratorExit:
        pass

test_accumulate = accumulate_and_send()
next(test_accumulate)
print(test_accumulate.send(1.0))
print(next(test_accumulate))
print(test_accumulate.send(2.0))
print(test_accumulate.send(3.0))
print(next(test_accumulate))

# Section: Iterable & Iterator

def process_items(items: Iterable[str]) -> list[str]:
    return [item.upper() for item in items]

print(process_items(["a", "b"]))
print(process_items(("a", "b")))
print(process_items({"a", "b"}))
print(process_items({"a": "b", "hello": "world"}))