refactor: simplify warning suppression by relying on global filter

Author: shuoweil

bigframes/_magics.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import warnings
-
 from IPython.core import magic_arguments  # type: ignore
 from IPython.core.getipython import get_ipython
 from IPython.display import display
@@ -44,11 +42,9 @@ def _cell_magic(line, cell):
         print("Query is missing.")
         return
     pyformat_args = ipython.user_ns
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore", category=FutureWarning)
-        dataframe = bigframes.pandas._read_gbq_colab(
-            cell, pyformat_args=pyformat_args, dry_run=args.dry_run
-        )
+    dataframe = bigframes.pandas._read_gbq_colab(
+        cell, pyformat_args=pyformat_args, dry_run=args.dry_run
+    )
     if args.destination_var:
         ipython.push({args.destination_var: dataframe})
 

bigframes/pandas/io/api.py

@@ -328,50 +328,47 @@ def _read_gbq_colab(
     if pyformat_args is None:
         pyformat_args = {}
 
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore", category=FutureWarning)
-
-        # Only try to set the global location if it's not a dry run. We don't want
-        # to bind to a location too early. This is especially important if the query
-        # only refers to local data and not any BigQuery tables.
-        if dry_run:
-            result = _try_read_gbq_colab_sessionless_dry_run(
-                query_or_table, pyformat_args=pyformat_args
-            )
-
-            if result is not None:
-                return result
+    # Only try to set the global location if it's not a dry run. We don't want
+    # to bind to a location too early. This is especially important if the query
+    # only refers to local data and not any BigQuery tables.
+    if dry_run:
+        result = _try_read_gbq_colab_sessionless_dry_run(
+            query_or_table, pyformat_args=pyformat_args
+        )
 
-            # If we made it this far, we must have a session that has already
-            # started. That means we can safely call the "real" _read_gbq_colab,
-            # which generates slightly nicer SQL.
-        else:
-            # Delay formatting the query with the special "session-less" logic. This
-            # avoids doing unnecessary work if the session already has a location or has
-            # already started.
-            create_query = functools.partial(
-                bigframes.core.pyformat.pyformat,
-                query_or_table,
-                pyformat_args=pyformat_args,
-                dry_run=True,
-            )
-            _set_default_session_location_if_possible_deferred_query(create_query)
-            if not config.options.bigquery._session_started:
-                # Don't warning about Polars in SQL cell.
-                # Related to b/437090788.
-                try:
-                    bigframes._importing.import_polars()
-                    warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning)
-                    config.options.bigquery.enable_polars_execution = True
-                except ImportError:
-                    pass  # don't fail if polars isn't available
-
-        return global_session.with_default_session(
-            bigframes.session.Session._read_gbq_colab,
+        if result is not None:
+            return result
+
+        # If we made it this far, we must have a session that has already
+        # started. That means we can safely call the "real" _read_gbq_colab,
+        # which generates slightly nicer SQL.
+    else:
+        # Delay formatting the query with the special "session-less" logic. This
+        # avoids doing unnecessary work if the session already has a location or has
+        # already started.
+        create_query = functools.partial(
+            bigframes.core.pyformat.pyformat,
             query_or_table,
             pyformat_args=pyformat_args,
-            dry_run=dry_run,
+            dry_run=True,
         )
+        _set_default_session_location_if_possible_deferred_query(create_query)
+        if not config.options.bigquery._session_started:
+            # Don't warning about Polars in SQL cell.
+            # Related to b/437090788.
+            try:
+                bigframes._importing.import_polars()
+                warnings.simplefilter("ignore", bigframes.exceptions.PreviewWarning)
+                config.options.bigquery.enable_polars_execution = True
+            except ImportError:
+                pass  # don't fail if polars isn't available
+
+    return global_session.with_default_session(
+        bigframes.session.Session._read_gbq_colab,
+        query_or_table,
+        pyformat_args=pyformat_args,
+        dry_run=dry_run,
+    )
 
 
 def read_gbq_model(model_name: str):