diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index 8040f3c7..a4965632 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_duckdb-stubs/__init__.pyi @@ -318,7 +318,18 @@ class DuckDBPyConnection: def list_type(self, type: sqltypes.DuckDBPyType) -> sqltypes.DuckDBPyType: ... def load_extension(self, extension: str) -> None: ... def map_type(self, key: sqltypes.DuckDBPyType, value: sqltypes.DuckDBPyType) -> sqltypes.DuckDBPyType: ... - def pl(self, rows_per_batch: pytyping.SupportsInt = 1000000, *, lazy: bool = False) -> polars.DataFrame: ... + @pytyping.overload + def pl( + self, rows_per_batch: pytyping.SupportsInt = 1000000, *, lazy: pytyping.Literal[False] = ... + ) -> polars.DataFrame: ... + @pytyping.overload + def pl( + self, rows_per_batch: pytyping.SupportsInt = 1000000, *, lazy: pytyping.Literal[True] + ) -> polars.LazyFrame: ... + @pytyping.overload + def pl( + self, rows_per_batch: pytyping.SupportsInt = 1000000, *, lazy: bool = False + ) -> pytyping.Union[polars.DataFrame, polars.LazyFrame]: ... def query(self, query: str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ... def query_progress(self) -> float: ... def read_csv( @@ -596,7 +607,16 @@ class DuckDBPyRelation: ) -> DuckDBPyRelation: ... def order(self, order_expr: str) -> DuckDBPyRelation: ... def percent_rank(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ... - def pl(self, batch_size: pytyping.SupportsInt = 1000000, *, lazy: bool = False) -> polars.DataFrame: ... + @pytyping.overload + def pl( + self, batch_size: pytyping.SupportsInt = 1000000, *, lazy: pytyping.Literal[False] = ... + ) -> polars.DataFrame: ... + @pytyping.overload + def pl(self, batch_size: pytyping.SupportsInt = 1000000, *, lazy: pytyping.Literal[True]) -> polars.LazyFrame: ... + @pytyping.overload + def pl( + self, batch_size: pytyping.SupportsInt = 1000000, *, lazy: bool = False + ) -> pytyping.Union[polars.DataFrame, polars.LazyFrame]: ... def product( self, column: str, groups: str = "", window_spec: str = "", projected_columns: str = "" ) -> DuckDBPyRelation: ... @@ -1241,12 +1261,27 @@ def map_type( def order( df: pandas.DataFrame, order_expr: str, *, connection: DuckDBPyConnection | None = None ) -> DuckDBPyRelation: ... +@pytyping.overload def pl( rows_per_batch: pytyping.SupportsInt = 1000000, *, - lazy: bool = False, + lazy: pytyping.Literal[False] = ..., connection: DuckDBPyConnection | None = None, ) -> polars.DataFrame: ... +@pytyping.overload +def pl( + rows_per_batch: pytyping.SupportsInt = 1000000, + *, + lazy: pytyping.Literal[True], + connection: DuckDBPyConnection | None = None, +) -> polars.LazyFrame: ... +@pytyping.overload +def pl( + rows_per_batch: pytyping.SupportsInt = 1000000, + *, + lazy: bool = False, + connection: DuckDBPyConnection | None = None, +) -> pytyping.Union[polars.DataFrame, polars.LazyFrame]: ... def project( df: pandas.DataFrame, *args: str | Expression, groups: str = "", connection: DuckDBPyConnection | None = None ) -> DuckDBPyRelation: ...