diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 81359ebb36..664313ce72 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -24,7 +24,6 @@ import typing from typing import ( Any, - Callable, Dict, IO, Iterable, @@ -504,6 +503,9 @@ def _read_gbq_colab( *, pyformat_args: Optional[Dict[str, Any]] = None, dry_run: bool = False, + callback: abc.Callable[ + [Any], None + ] = lambda _: None, # TODO: use an Event class not Any. ) -> Union[dataframe.DataFrame, pandas.Series]: """A version of read_gbq that has the necessary default values for use in colab integrations. @@ -535,6 +537,7 @@ def _read_gbq_colab( index_col=bigframes.enums.DefaultIndexKind.NULL, force_total_order=False, dry_run=typing.cast(Union[Literal[False], Literal[True]], dry_run), + callback=callback, ) @overload @@ -992,7 +995,7 @@ def read_csv( pandas.Series, pandas.Index, np.ndarray[Any, Any], - Callable[[Any], bool], + abc.Callable[[Any], bool], ] ] = None, dtype: Optional[Dict] = None, diff --git a/bigframes/session/loader.py b/bigframes/session/loader.py index 1e32f3d860..04a4a75ee1 100644 --- a/bigframes/session/loader.py +++ b/bigframes/session/loader.py @@ -14,6 +14,7 @@ from __future__ import annotations +from collections import abc import copy import dataclasses import datetime @@ -22,6 +23,7 @@ import os import typing from typing import ( + Any, Dict, Generator, Hashable, @@ -743,6 +745,9 @@ def read_gbq_query( filters: third_party_pandas_gbq.FiltersType = (), dry_run: bool = False, force_total_order: Optional[bool] = None, + callback: abc.Callable[ + [Any], None + ] = lambda _: None, # TODO: use an Event class not Any. ) -> dataframe.DataFrame | pandas.Series: import bigframes.dataframe as dataframe