fix: to_pydantic typing (#2517)

currently, to_pydantic will always return LanceModel. If type checking
is enabled in my project. I have to use `cast(data,
List[RealModelType])` to solve type error. This PR uses generic to solve
this problem.
This commit is contained in:
Chen Chongchen
2025-07-25 03:30:15 +08:00
committed by GitHub
parent 12d4ce4cfe
commit 3cd7dd3375

View File

@@ -14,7 +14,7 @@ from typing import (
Literal,
Optional,
Tuple,
Type,
TypeVar,
Union,
Any,
)
@@ -58,6 +58,8 @@ if TYPE_CHECKING:
else:
from typing_extensions import Self
T = TypeVar("T", bound="LanceModel")
# Pydantic validation function for vector queries
def ensure_vector_query(
@@ -746,8 +748,8 @@ class LanceQueryBuilder(ABC):
return self.to_arrow(timeout=timeout).to_pylist()
def to_pydantic(
self, model: Type[LanceModel], *, timeout: Optional[timedelta] = None
) -> List[LanceModel]:
self, model: type[T], *, timeout: Optional[timedelta] = None
) -> list[T]:
"""Return the table as a list of pydantic models.
Parameters