diff --git a/doc/source/whatsnew/v1.0.0.rst b/doc/source/whatsnew/v1.0.0.rst index db23bfdc8a5bd..6e1cf9ab66201 100644 --- a/doc/source/whatsnew/v1.0.0.rst +++ b/doc/source/whatsnew/v1.0.0.rst @@ -600,6 +600,7 @@ I/O - Bug in :meth:`DataFrame.to_clipboard` which did not work reliably in ipython (:issue:`22707`) - Bug in :func:`read_json` where default encoding was not set to ``utf-8`` (:issue:`29565`) - Bug in :class:`PythonParser` where str and bytes were being mixed when dealing with the decimal field (:issue:`29650`) +- :meth:`read_gbq` now accepts ``progress_bar_type`` to display progress bar while the data downloads. (:issue:`29857`) - Plotting diff --git a/pandas/io/gbq.py b/pandas/io/gbq.py index b120de1b3011a..9fac7a52b370d 100644 --- a/pandas/io/gbq.py +++ b/pandas/io/gbq.py @@ -27,6 +27,7 @@ def read_gbq( use_bqstorage_api=None, private_key=None, verbose=None, + progress_bar_type=None, ): """ Load data from Google BigQuery. @@ -134,6 +135,30 @@ def read_gbq( Deprecated in pandas-gbq version 0.4.0. Use the `logging module to adjust verbosity instead `__. + progress_bar_type : Optional, str + If set, use the `tqdm `__ library to + display a progress bar while the data downloads. Install the + ``tqdm`` package to use this feature. + + Possible values of ``progress_bar_type`` include: + + ``None`` + No progress bar. + ``'tqdm'`` + Use the :func:`tqdm.tqdm` function to print a progress bar + to :data:`sys.stderr`. + ``'tqdm_notebook'`` + Use the :func:`tqdm.tqdm_notebook` function to display a + progress bar as a Jupyter notebook widget. + ``'tqdm_gui'`` + Use the :func:`tqdm.tqdm_gui` function to display a + progress bar as a graphical dialog box. + + Note that his feature requires version 0.12.0 or later of the + ``pandas-gbq`` package. And it requires the ``tqdm`` package. Slightly + different than ``pandas-gbq``, here the default is ``None``. + + .. versionadded:: 1.0.0 Returns ------- @@ -152,6 +177,9 @@ def read_gbq( # START: new kwargs. Don't populate unless explicitly set. if use_bqstorage_api is not None: kwargs["use_bqstorage_api"] = use_bqstorage_api + + if progress_bar_type is not None: + kwargs["progress_bar_type"] = progress_bar_type # END: new kwargs # START: deprecated kwargs. Don't populate unless explicitly set. diff --git a/pandas/tests/io/test_gbq.py b/pandas/tests/io/test_gbq.py index 52147f4e1afc7..75c80bb0b8025 100644 --- a/pandas/tests/io/test_gbq.py +++ b/pandas/tests/io/test_gbq.py @@ -144,6 +144,24 @@ def mock_read_gbq(sql, **kwargs): assert "use_bqstorage_api" not in captured_kwargs +@pytest.mark.parametrize("progress_bar", [None, "foo"]) +def test_read_gbq_progress_bar_type_kwarg(monkeypatch, progress_bar): + # GH 29857 + captured_kwargs = {} + + def mock_read_gbq(sql, **kwargs): + captured_kwargs.update(kwargs) + return DataFrame([[1.0]]) + + monkeypatch.setattr("pandas_gbq.read_gbq", mock_read_gbq) + pd.read_gbq("SELECT 1", progress_bar_type=progress_bar) + + if progress_bar: + assert "progress_bar_type" in captured_kwargs + else: + assert "progress_bar_type" not in captured_kwargs + + @pytest.mark.single class TestToGBQIntegrationWithServiceAccountKeyPath: @classmethod