chore: add overload to to_datetime and cleanup type errors in tests/system/small/test_pandas.py (#766)

* chore: add overload to to_datetime and cleanup type errors in tests/unit/test_pandas.py

* update import

Author: Genesis929

bigframes/pandas/__init__.py

@@ -699,6 +699,32 @@ def read_gbq_function(function_name: str):
 read_gbq_function.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_function)
 
 
+@typing.overload
+def to_datetime(
+    arg: vendored_pandas_datetimes.local_scalars,
+    *,
+    utc: bool = False,
+    format: Optional[str] = None,
+    unit: Optional[str] = None,
+) -> Union[pandas.Timestamp, datetime]:
+    ...
+
+
+@typing.overload
+def to_datetime(
+    arg: Union[
+        vendored_pandas_datetimes.local_iterables,
+        bigframes.series.Series,
+        bigframes.dataframe.DataFrame,
+    ],
+    *,
+    utc: bool = False,
+    format: Optional[str] = None,
+    unit: Optional[str] = None,
+) -> bigframes.series.Series:
+    ...
+
+
 def to_datetime(
     arg: Union[
         vendored_pandas_datetimes.local_scalars,

tests/system/small/test_pandas.py

@@ -394,12 +394,8 @@ def test_cut(scalars_dfs):
 
     # make sure the result is a supported dtype
     assert bf_result.dtype == bpd.Int64Dtype()
-
-    # TODO(b/340884971): fix type error
-    bf_result = bf_result.to_pandas()  # type: ignore
     pd_result = pd_result.astype("Int64")
-    # TODO(b/340884971): fix type error
-    pd.testing.assert_series_equal(bf_result, pd_result)  # type: ignore
+    pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)
 
 
 def test_cut_default_labels(scalars_dfs):
@@ -529,13 +525,9 @@ def test_qcut(scalars_dfs, q):
         scalars_pandas_df["float64_col"], q, labels=False, duplicates="drop"
     )
     bf_result = bpd.qcut(scalars_df["float64_col"], q, labels=False, duplicates="drop")
-
-    # TODO(b/340884971): fix type error
-    bf_result = bf_result.to_pandas()  # type: ignore
     pd_result = pd_result.astype("Int64")
 
-    # TODO(b/340884971): fix type error
-    pd.testing.assert_series_equal(bf_result, pd_result)  # type: ignore
+    pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)
 
 
 @pytest.mark.parametrize(
@@ -572,9 +564,8 @@ def test_to_datetime_scalar(arg, utc, unit, format):
     ],
 )
 def test_to_datetime_iterable(arg, utc, unit, format):
-    # TODO(b/340884971): fix type error
     bf_result = (
-        bpd.to_datetime(arg, utc=utc, unit=unit, format=format)  # type: ignore
+        bpd.to_datetime(arg, utc=utc, unit=unit, format=format)
         .to_pandas()
         .astype("datetime64[ns, UTC]" if utc else "datetime64[ns]")
     )
@@ -589,9 +580,8 @@ def test_to_datetime_iterable(arg, utc, unit, format):
 def test_to_datetime_series(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     col = "int64_too"
-    # TODO(b/340884971): fix type error
     bf_result = (
-        bpd.to_datetime(scalars_df[col], unit="s").to_pandas().astype("datetime64[s]")  # type: ignore
+        bpd.to_datetime(scalars_df[col], unit="s").to_pandas().astype("datetime64[s]")
     )
     pd_result = pd.Series(pd.to_datetime(scalars_pandas_df[col], unit="s"))
     pd.testing.assert_series_equal(
@@ -614,8 +604,7 @@ def test_to_datetime_series(scalars_dfs):
     ],
 )
 def test_to_datetime_unit_param(arg, unit):
-    # TODO(b/340884971): fix type error
-    bf_result = bpd.to_datetime(arg, unit=unit).to_pandas().astype("datetime64[ns]")  # type: ignore
+    bf_result = bpd.to_datetime(arg, unit=unit).to_pandas().astype("datetime64[ns]")
     pd_result = pd.Series(pd.to_datetime(arg, unit=unit)).dt.floor("us")
     pd.testing.assert_series_equal(
         bf_result, pd_result, check_index_type=False, check_names=False
@@ -632,9 +621,8 @@ def test_to_datetime_unit_param(arg, unit):
     ],
 )
 def test_to_datetime_format_param(arg, utc, format):
-    # TODO(b/340884971): fix type error
     bf_result = (
-        bpd.to_datetime(arg, utc=utc, format=format)  # type: ignore
+        bpd.to_datetime(arg, utc=utc, format=format)
         .to_pandas()
         .astype("datetime64[ns, UTC]" if utc else "datetime64[ns]")
     )
@@ -686,9 +674,8 @@ def test_to_datetime_format_param(arg, utc, format):
     ],
 )
 def test_to_datetime_string_inputs(arg, utc, output_in_utc, format):
-    # TODO(b/340884971): fix type error
     bf_result = (
-        bpd.to_datetime(arg, utc=utc, format=format)  # type: ignore
+        bpd.to_datetime(arg, utc=utc, format=format)
         .to_pandas()
         .astype("datetime64[ns, UTC]" if output_in_utc else "datetime64[ns]")
     )
@@ -730,9 +717,8 @@ def test_to_datetime_string_inputs(arg, utc, output_in_utc, format):
     ],
 )
 def test_to_datetime_timestamp_inputs(arg, utc, output_in_utc):
-    # TODO(b/340884971): fix type error
     bf_result = (
-        bpd.to_datetime(arg, utc=utc)  # type: ignore
+        bpd.to_datetime(arg, utc=utc)
         .to_pandas()
         .astype("datetime64[ns, UTC]" if output_in_utc else "datetime64[ns]")
     )

third_party/bigframes_vendored/pandas/core/tools/datetimes.py

@@ -1,14 +1,14 @@
 # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/tools/datetimes.py
 
 from datetime import datetime
-from typing import Iterable, Mapping, Union
+from typing import List, Mapping, Tuple, Union
 
 import pandas as pd
 
 from bigframes import constants, series
 
 local_scalars = Union[int, float, str, datetime]
-local_iterables = Union[Iterable, pd.Series, pd.DataFrame, Mapping]
+local_iterables = Union[List, Tuple, pd.Series, pd.DataFrame, Mapping]
 
 
 def to_datetime(