Add explicit polars overloads

Author: J-Meyers

_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: ...