Python type hints partially broken for @planning_entity classes

[luxaritas]

Describe the bug
When using planning_entity annotated classes, their types do not properly pass through in generics, such as when using ConstraintFactory#for_each.

Expected behavior
In the below code, the type of stream would be UniConstraintStream[Job] and intellisense would show x as a valid property of job

Actual behavior
the type of stream would be UniConstraintStream and intellisense shows job as Any

To Reproduce

@planning_entity
@dataclass
class Job:
    x: int

def constraint(factory: ConstraintFactory):
    stream = factory.for_each(Job)
    stream.filter(lambda job: job.x == 5)

Environment

Timefold Solver Version or Git ref:
1.16.0b0

Output of java -version:

openjdk version "23" 2024-09-17
OpenJDK Runtime Environment (build 23)
OpenJDK 64-Bit Server VM (build 23, mixed mode, sharing)

Output of uname -a or ver:

Linux 6.6.63-1-MANJARO #1 SMP PREEMPT_DYNAMIC Sat Nov 23 02:15:34 UTC 2024 x86_64 GNU/Linux

Additional information

@planning_entity is annotated as follows:

def planning_entity(entity_class: Type = None, /, *, pinning_filter: Callable = None) -> Union[Type, Callable[[Type], Type]]:

Shouldn't entity_class be TypeVared, and the return just be that typevar (and not a union with callable)?

NB: I'm specifically using Pylance in VS Code