From 65e0dde60a39197aa32475a4a01cdffdc0bb6350 Mon Sep 17 00:00:00 2001
From: Stefano Cretti <stefanocretti99@gmail.com>
Date: Thu, 25 Jun 2026 23:04:44 +0200
Subject: [PATCH] typing: DataSource.to_polars and .to_pandas

Added type hinting for the methods `to_pandas` and `to_polars` of DataSource in py-oxbow.
---
 py-oxbow/oxbow/_core/base.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/py-oxbow/oxbow/_core/base.py b/py-oxbow/oxbow/_core/base.py
index b550dd2..454e623 100644
--- a/py-oxbow/oxbow/_core/base.py
+++ b/py-oxbow/oxbow/_core/base.py
@@ -3,7 +3,7 @@
 import pathlib
 import warnings
 from abc import abstractmethod
-from typing import IO, Any, Callable, Generator, Literal
+from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Literal, overload
 
 try:
     from typing import Self
@@ -21,6 +21,10 @@
     BatchReaderFragment,
 )
 
+if TYPE_CHECKING:
+    import pandas as pd
+    import polars as pl
+
 
 class DataSource:
     """
@@ -226,7 +230,7 @@ def dataset(self) -> BatchReaderDataset:
         """
         return BatchReaderDataset(self.fragments())
 
-    def to_pandas(self):
+    def to_pandas(self) -> "pd.DataFrame":
         """
         Convert the dataset to a Pandas DataFrame.
 
@@ -239,7 +243,12 @@ def to_pandas(self):
 
     pd = to_pandas
 
-    def to_polars(self, lazy=False):
+    @overload
+    def to_polars(self, lazy: Literal[True]) -> "pl.LazyFrame": ...
+    @overload
+    def to_polars(self, lazy: Literal[False] = ...) -> "pl.DataFrame": ...
+
+    def to_polars(self, lazy: bool = False) -> "pl.DataFrame | pl.LazyFrame":
         """
         Convert the data source to a Polars DataFrame or LazyFrame.