From 71f87862730bef472b87e39818b3400d7ee1f055 Mon Sep 17 00:00:00 2001 From: Jonathan Romano Date: Wed, 11 Dec 2024 00:07:25 -0500 Subject: [PATCH] fix: More accurate type hints for Python @planning_entity --- .../src/main/python/domain/_annotations.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/python/python-core/src/main/python/domain/_annotations.py b/python/python-core/src/main/python/domain/_annotations.py index 6a2a5b3b88..dfd5ba77bb 100644 --- a/python/python-core/src/main/python/domain/_annotations.py +++ b/python/python-core/src/main/python/domain/_annotations.py @@ -1,12 +1,13 @@ from _jpyinterpreter import JavaAnnotation from enum import Enum from jpype import JImplements, JOverride, JClass -from typing import Union, List, Callable, Type, TypeVar +from typing import Union, List, Callable, Type, TypeVar, overload from ._variable_listener import VariableListener from .._timefold_java_interop import ensure_init Solution_ = TypeVar('Solution_') +Entity_ = TypeVar('Entity_') class PlanningId(JavaAnnotation): @@ -741,9 +742,13 @@ def __init__(self, delegate): def accept(self, solution, entity): return self.delegate(solution, entity) - -def planning_entity(entity_class: Type = None, /, *, pinning_filter: Callable = None) -> Union[Type, - Callable[[Type], Type]]: +@overload +def planning_entity(entity_class: Type[Entity_]) -> Type[Entity_]: + ... +@overload +def planning_entity(pinning_filter: Callable = None) -> Callable[[Type[Entity_]], Type[Entity_]]: + ... +def planning_entity(entity_class = None, /, *, pinning_filter = None): """ Specifies that the class is a planning entity. There are two types of entities: