diff --git a/google/cloud/bigquery/_tqdm_helpers.py b/google/cloud/bigquery/_tqdm_helpers.py index f2355ab3b..ae289d8a6 100644 --- a/google/cloud/bigquery/_tqdm_helpers.py +++ b/google/cloud/bigquery/_tqdm_helpers.py @@ -22,6 +22,7 @@ try: import tqdm # type: ignore + except ImportError: # pragma: NO COVER tqdm = None @@ -48,7 +49,7 @@ def get_progress_bar(progress_bar_type, description, total, unit): if progress_bar_type == "tqdm": return tqdm.tqdm(desc=description, total=total, unit=unit) elif progress_bar_type == "tqdm_notebook": - return tqdm.tqdm_notebook(desc=description, total=total, unit=unit) + return tqdm.notebook.tqdm(desc=description, total=total, unit=unit) elif progress_bar_type == "tqdm_gui": return tqdm.tqdm_gui(desc=description, total=total, unit=unit) except (KeyError, TypeError): diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index b5f2e58c6..fca43f1ee 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -45,7 +45,9 @@ geopandas = None try: - from tqdm import tqdm + import tqdm + from tqdm.std import TqdmDeprecationWarning + except (ImportError, AttributeError): # pragma: NO COVER tqdm = None @@ -2798,7 +2800,7 @@ def test_to_arrow_w_bqstorage_no_streams(self): @unittest.skipIf(tqdm is None, "Requires `tqdm`") @mock.patch("tqdm.tqdm_gui") - @mock.patch("tqdm.tqdm_notebook") + @mock.patch("tqdm.notebook.tqdm") @mock.patch("tqdm.tqdm") def test_to_arrow_progress_bar(self, tqdm_mock, tqdm_notebook_mock, tqdm_gui_mock): from google.cloud.bigquery.schema import SchemaField @@ -3146,7 +3148,7 @@ def test_to_dataframe_datetime_out_of_pyarrow_bounds(self): @unittest.skipIf(pandas is None, "Requires `pandas`") @unittest.skipIf(tqdm is None, "Requires `tqdm`") @mock.patch("tqdm.tqdm_gui") - @mock.patch("tqdm.tqdm_notebook") + @mock.patch("tqdm.notebook.tqdm") @mock.patch("tqdm.tqdm") def test_to_dataframe_progress_bar( self, tqdm_mock, tqdm_notebook_mock, tqdm_gui_mock @@ -3249,7 +3251,7 @@ def test_to_dataframe_no_tqdm(self): @unittest.skipIf(pandas is None, "Requires `pandas`") @unittest.skipIf(tqdm is None, "Requires `tqdm`") @mock.patch("tqdm.tqdm_gui", new=None) # will raise TypeError on call - @mock.patch("tqdm.tqdm_notebook", new=None) # will raise TypeError on call + @mock.patch("tqdm.notebook.tqdm", new=None) # will raise TypeError on call @mock.patch("tqdm.tqdm", new=None) # will raise TypeError on call def test_to_dataframe_tqdm_error(self): from google.cloud.bigquery.schema import SchemaField @@ -3281,7 +3283,10 @@ def test_to_dataframe_tqdm_error(self): # Warn that a progress bar was requested, but creating the tqdm # progress bar failed. for warning in warned: - self.assertIs(warning.category, UserWarning) + self.assertIn( + warning.category, + [UserWarning, DeprecationWarning, TqdmDeprecationWarning], + ) @unittest.skipIf(pandas is None, "Requires `pandas`") def test_to_dataframe_w_empty_results(self):