[Python] TypeVar with UnionType to express lambda parameter type to Pyright

I have a collection of objects which all inherit from Component. I'm writing a function that allows one to select a subset of objects, based on type, and pass them to a lambda. The simple case shown below works, both executing correctly at runtime and also correctly inferring lambda types in VS Code/Pyright:

from collections.abc import Callable
from dataclasses import dataclass
from types import UnionType
from typing import Any, TypeVar

@dataclass
class Component:
c: int = 0

@dataclass
class Foo(Component):
f: int = 0

@dataclass
class Bar(Component):
f: int = 0

T = TypeVar("T", bound=Component)

def query(
collection: list[Component],
value: Callable[[T], Any],
kind: type[T] = Component
) -> list:
return [value(item) for item in collection if isinstance(item, kind)]

coll = [Foo(f=1), Bar(f=2), Component(c=3)]
query(coll, lambda x: x.f, kind=Foo)
#=> [1]


However, I want to be able to specify a UnionType for selecting more than one type. The following changes work at runtime, but make Pyright cranky:

def query(
…,
kind: type[T] | UnionType = Component
…,
# Adding `| UnionType` to the `kind` parameter causes a problem to be shown on the `item` in `value(item)`:
#
# Argument of type "Component" cannot be assigned to parameter of type "T@query"
# Type "Component" is incompatible with type "T@query"

query(coll, lambda x: x.f, kind=Foo | Bar)
#=> [1, 2]
# Adding `Foo | Bar` when calling the query() function causes Pyright to think that
# `x` is of type `Component`, and so a problem is shown for `.f`:
#
# Cannot access attribute "f" for class "Component*"
# Attribute "f" is unknown


How can I express this to make Pyright happy, so that I get clean static analysis and Intellisense for my code?

Currently targeting Python 3.11+, but happy to select a more modern version if it helps.


The code above is a simplified case for clarity: query() does far more than shown here, uses generators and not lists, the class hierarchy and properties are significantly more interesting than shown. This is hopefully obvious, and the simplification desirable. I mention it only to forestall comments or answers that focus on the trivial setup, instead of the meat of the question.

Continue reading...