howto: a decorator that sometimes takes a generator

[bukzor]

This is a boiled-down toy example from a larger set of requirements in my real project.

This is close as I got on my own:

def imap(transformer: Callable[[T], T], generator: Generator[T]) -> Generator[T]:
    for x in generator:
        yield transformer(x)


def transformed(transformer: Callable[[T], T]):
    def decorator(f: Callable[P, T]):
        @wraps(f)
        def decorated(*args: P.args, **kwargs: P.kwargs) -> T:
            result = f(*args, **kwargs)
            if isinstance(result, types.GeneratorType):

                # E Argument type is partially unknown: # Argument type is "GeneratorType[Unknown, Unknown, Unknown]*" Pyright (reportUnknownArgumentType) [21, 42]
                return imap(transformer, result)
            else:
                return transformer(result)
        return decorated
    return decorator

Desiderata:

• decorator should not change the type of the decorated function
• we have an ability to process generated items (if we have a generator)
• avoid poor typing: Any, cast, Unknown

If it's not possible to have all these, I'd like to understand why, at least.

[Glinte]

This can be explained in two step.

1. Your type hint is wrong. Your decorator function actually accepts Callable[P, T | Generator[T]],

playground

def double(x: int):
    return 2 * x

@transformed(double)
#^^^^^^^^^^^^^^^^^^^
#Argument of type "() -> Generator[int, None, None]" cannot be assigned to parameter "f" of type "(**P@decorator) -> int" in function "decorator"
#  Type "() -> Generator[int, None, None]" is not assignable to type "(**P@decorator) -> int"
#    Function return type "Generator[int, None, None]" is incompatible with type "int"
#      "Generator[int, None, None]" is not assignable to "int"  (reportArgumentType)
#  Argument type is "GeneratorType[Unknown, Unknown, Unknown]*"  (reportUnknownArgumentType)
def despite_what_you_think_this_wont_type_check() -> Generator[int]:
    for i in range(10):
        yield i

Also, the decorator function returns a different type of function depending on the output type, so we also need @overload.

playground

from collections.abc import Callable, Generator
from typing import TypeVar, ParamSpec, overload
import types
from functools import wraps

T = TypeVar("T")
P = ParamSpec("P")

def imap(transformer: Callable[[T], T], generator: Generator[T]) -> Generator[T]:
    for x in generator:
        yield transformer(x)


def transformed(transformer: Callable[[T], T]):
    @overload
    def decorator(f: Callable[P, Generator[T]]) -> Callable[P, Generator[T]]: ...
    
    @overload
    def decorator(f: Callable[P, T]) -> Callable[P, T]: ...

    def decorator(f: Callable[P, T | Generator[T]]):
        @wraps(f)
        def decorated(*args: P.args, **kwargs: P.kwargs) -> T | Generator[T]:
            result = f(*args, **kwargs)
            if isinstance(result, types.GeneratorType):
                # E Argument type is partially unknown: # Argument type is "GeneratorType[Unknown, Unknown, Unknown]*" Pyright (reportUnknownArgumentType) [21, 42]
                return imap(transformer, result)
#                                        ^^^^^^
#Argument type is partially unknown
#  Argument corresponds to parameter "generator" in function "imap"
#  Argument type is "GeneratorType[Unknown, Unknown, Unknown]* | GeneratorType[T@transformed, None, None]"  (reportUnknownArgumentType)
            else:
                return transformer(result)
        return decorated
    return decorator

2. This still doesn't work

Unfortunately, this increases the number of errors to 4. I included the most readable one above. What is it saying? Pyright thinks that result is still possible to be GeneratorType[Unknown, Unknown, Unknown], why? This is actually because T can be matched to Generator[Any, Any, Any] too. As far as pyright is concerned, you could have passed a transformer with type Callable[Generator, Generator].

Two remaining errors result from functools.wraps not behaving sanely when there are other errors, I don't understand those either. And one seem to be a pyright bug?

If you understand that T can be matched to Generator[Any, Any, Any], you would probably think about specifying T as "Anything but not a Generator. Unfortunately, you cannot express negation in python's type system (python/typing#801).

And here are some tangents to end it

1. you can use the following to get syntax highlighting

```python
code here
```​

2. it would be nice if you included all the imports, variables and mention the you used strict mode, which is not the default

[bukzor]

I believe the core issue here is union ambiguity; when R itself can be a
generator type, the union R | Generator[R] creates genuine ambiguity:

# If R = Generator[int], then R | Generator[R] becomes:
result: Generator[int] | Generator[Generator[int]]

# After isinstance(result, Generator), both cases match:
if isinstance(result, Generator):  # True for both cases!
    # Is this Generator[int] or Generator[Generator[int]]?

This explains the Generator[Unknown] types pyright reports - it correctly
identifies that the yield type cannot be determined from the isinstance check
alone.

Type System Constraints That Block Clean Solutions

1. No parameterized constraints: Cannot write
   F: Callable[P, R | Generator[R]] where P, R are type variables
2. No overload dispatch with bodies: @overload only allows stubs, not
   implementations
3. Union narrowing limitations: isinstance(result, Generator) doesn't
   narrow generic parameters

Workaround Solution: Separate Functions

The type-safe approach requires splitting the decorator:

def transformed_value[R, **P](transformer: Callable[[R], R]) -> Callable[
  [Callable[P, R]],
  Callable[P, R]
]: ...
def transformed_generator[R, **P](transformer: Callable[[R], R]) -> Callable[
  [Callable[P, Generator[R]]],
  Callable[P, Generator[R]]
]: ...

This eliminates all casts and runtime type ambiguity, at the cost of requiring
users to choose explicitly.

Sufficient Typing Extensions

The pattern would be naturally expressible with either:

1. Overloaded dispatch with function bodies
2. A better way to type-narrow a union of generics