chore: format warning message (#1281)

* chore: format warning message

* fix unit tests

* fix mypy

* use regional endpoints warning

Author: chelsea-lin

bigframes/_config/bigquery_options.py

@@ -25,7 +25,7 @@
 
 import bigframes.constants
 import bigframes.enums
-import bigframes.exceptions
+import bigframes.exceptions as bfe
 
 SESSION_STARTED_MESSAGE = (
     "Cannot change '{attribute}' once a session has started. "
@@ -55,15 +55,12 @@ def _get_validated_location(value: Optional[str]) -> Optional[str]:
         bigframes.constants.ALL_BIGQUERY_LOCATIONS,
         key=lambda item: jellyfish.levenshtein_distance(location, item),
     )
-    warnings.warn(
-        UNKNOWN_LOCATION_MESSAGE.format(location=location, possibility=possibility),
-        # There are many layers before we get to (possibly) the user's code:
-        # -> bpd.options.bigquery.location = "us-central-1"
-        # -> location.setter
-        # -> _get_validated_location
-        stacklevel=3,
-        category=bigframes.exceptions.UnknownLocationWarning,
-    )
+    # There are many layers before we get to (possibly) the user's code:
+    # -> bpd.options.bigquery.location = "us-central-1"
+    # -> location.setter
+    # -> _get_validated_location
+    msg = UNKNOWN_LOCATION_MESSAGE.format(location=location, possibility=possibility)
+    warnings.warn(msg, stacklevel=3, category=bfe.UnknownLocationWarning)
 
     return value
 
@@ -275,10 +272,11 @@ def use_regional_endpoints(self, value: bool):
             )
 
         if value:
-            warnings.warn(
+            msg = (
                 "Use of regional endpoints is a feature in preview and "
                 "available only in selected regions and projects. "
             )
+            warnings.warn(msg, category=bfe.PreviewWarning, stacklevel=2)
 
         self._use_regional_endpoints = value
 
@@ -334,9 +332,12 @@ def client_endpoints_override(self) -> dict:
 
     @client_endpoints_override.setter
     def client_endpoints_override(self, value: dict):
-        warnings.warn(
-            "This is an advanced configuration option for directly setting endpoints. Incorrect use may lead to unexpected behavior or system instability. Proceed only if you fully understand its implications."
+        msg = (
+            "This is an advanced configuration option for directly setting endpoints. "
+            "Incorrect use may lead to unexpected behavior or system instability. "
+            "Proceed only if you fully understand its implications."
         )
+        warnings.warn(msg)
 
         if self._session_started and self._client_endpoints_override != value:
             raise ValueError(

bigframes/_config/experiment_options.py

@@ -14,6 +14,8 @@
 
 import warnings
 
+import bigframes.exceptions as bfe
+
 
 class ExperimentOptions:
     """
@@ -31,9 +33,11 @@ def semantic_operators(self) -> bool:
     @semantic_operators.setter
     def semantic_operators(self, value: bool):
         if value is True:
-            warnings.warn(
-                "Semantic operators are still under experiments, and are subject to change in the future."
+            msg = (
+                "Semantic operators are still under experiments, and are subject "
+                "to change in the future."
             )
+            warnings.warn(msg, category=bfe.PreviewWarning)
         self._semantic_operators = value
 
     @property
@@ -43,7 +47,9 @@ def blob(self) -> bool:
     @blob.setter
     def blob(self, value: bool):
         if value is True:
-            warnings.warn(
-                "BigFrames Blob is still under experiments. It may not work and subject to change in the future."
+            msg = (
+                "BigFrames Blob is still under experiments. It may not work and "
+                "subject to change in the future."
             )
+            warnings.warn(msg, category=bfe.PreviewWarning)
         self._blob = value

bigframes/core/__init__.py

@@ -39,6 +39,7 @@
 import bigframes.core.utils
 from bigframes.core.window_spec import WindowSpec
 import bigframes.dtypes
+import bigframes.exceptions as bfe
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
 
@@ -106,10 +107,11 @@ def from_table(
         if offsets_col and primary_key:
             raise ValueError("must set at most one of 'offests', 'primary_key'")
         if any(i.field_type == "JSON" for i in table.schema if i.name in schema.names):
-            warnings.warn(
-                "Interpreting JSON column(s) as StringDtype and pyarrow.large_string. This behavior may change in future versions.",
-                bigframes.exceptions.PreviewWarning,
+            msg = (
+                "Interpreting JSON column(s) as pyarrow.large_string. "
+                "This behavior may change in future versions."
             )
+            warnings.warn(msg, bfe.PreviewWarning)
         # define data source only for needed columns, this makes row-hashing cheaper
         table_def = nodes.GbqTable.from_table(table, columns=schema.names)
 
@@ -228,10 +230,8 @@ def slice(
         self, start: Optional[int], stop: Optional[int], step: Optional[int]
     ) -> ArrayValue:
         if self.node.order_ambiguous and not (self.session._strictly_ordered):
-            warnings.warn(
-                "Window ordering may be ambiguous, this can cause unstable results.",
-                bigframes.exceptions.AmbiguousWindowWarning,
-            )
+            msg = "Window ordering may be ambiguous, this can cause unstable results."
+            warnings.warn(msg, bfe.AmbiguousWindowWarning)
         return ArrayValue(
             nodes.SliceNode(
                 self.node,
@@ -252,10 +252,10 @@ def promote_offsets(self) -> Tuple[ArrayValue, str]:
                     "Generating offsets not supported in partial ordering mode"
                 )
             else:
-                warnings.warn(
-                    "Window ordering may be ambiguous, this can cause unstable results.",
-                    bigframes.exceptions.AmbiguousWindowWarning,
+                msg = (
+                    "Window ordering may be ambiguous, this can cause unstable results."
                 )
+                warnings.warn(msg, category=bfe.AmbiguousWindowWarning)
 
         return (
             ArrayValue(
@@ -391,10 +391,8 @@ def project_window_op(
                         "Generating offsets not supported in partial ordering mode"
                     )
                 else:
-                    warnings.warn(
-                        "Window ordering may be ambiguous, this can cause unstable results.",
-                        bigframes.exceptions.AmbiguousWindowWarning,
-                    )
+                    msg = "Window ordering may be ambiguous, this can cause unstable results."
+                    warnings.warn(msg, category=bfe.AmbiguousWindowWarning)
 
         output_name = self._gen_namespaced_uid()
         return (

bigframes/core/blocks.py

@@ -63,7 +63,7 @@
 import bigframes.core.utils as utils
 import bigframes.core.window_spec as windows
 import bigframes.dtypes
-import bigframes.exceptions
+import bigframes.exceptions as bfe
 import bigframes.features
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
@@ -137,10 +137,8 @@ def __init__(
                 )
 
         if len(index_columns) == 0:
-            warnings.warn(
-                "Creating object with Null Index. Null Index is a preview feature.",
-                category=bigframes.exceptions.NullIndexPreviewWarning,
-            )
+            msg = "Creating object with Null Index. Null Index is a preview feature."
+            warnings.warn(msg, category=bfe.NullIndexPreviewWarning)
         self._index_columns = tuple(index_columns)
         # Index labels don't need complicated hierarchical access so can store as tuple
         self._index_labels = (
@@ -616,13 +614,13 @@ def _materialize_local(
                     " # Setting it to None will download all the data\n"
                     f"{constants.FEEDBACK_LINK}"
                 )
-
-            warnings.warn(
+            msg = (
                 f"The data size ({table_mb:.2f} MB) exceeds the maximum download limit of"
-                f"({max_download_size} MB). It will be downsampled to {max_download_size} MB for download."
-                "\nPlease refer to the documentation for configuring the downloading limit.",
-                UserWarning,
+                f"({max_download_size} MB). It will be downsampled to {max_download_size} "
+                "MB for download.\nPlease refer to the documentation for configuring "
+                "the downloading limit."
             )
+            warnings.warn(msg, category=UserWarning)
             total_rows = execute_result.total_rows
             # Remove downsampling config from subsequent invocations, as otherwise could result in many
             # iterations if downsampling undershoots

bigframes/core/compile/ibis_types.py

@@ -32,6 +32,7 @@
 import pyarrow as pa
 
 import bigframes.dtypes
+import bigframes.exceptions as bfe
 
 # Type hints for Ibis data types supported by BigQuery DataFrame
 IbisDtype = Union[
@@ -305,10 +306,11 @@ def ibis_dtype_to_bigframes_dtype(
 
     # Temporary: Will eventually support an explicit json type instead of casting to string.
     if isinstance(ibis_dtype, ibis_dtypes.JSON):
-        warnings.warn(
-            "Interpreting JSON as string. This behavior may change in future versions.",
-            bigframes.exceptions.PreviewWarning,
+        msg = (
+            "Interpreting JSON column(s) as pyarrow.large_string. This behavior may change "
+            "in future versions."
         )
+        warnings.warn(msg, category=bfe.PreviewWarning)
         return bigframes.dtypes.JSON_DTYPE
 
     if ibis_dtype in IBIS_TO_BIGFRAMES:

bigframes/core/compile/polars/__init__.py

@@ -22,4 +22,5 @@
 
     __all__ = ["PolarsCompiler"]
 except Exception:
-    warnings.warn("Polars compiler not available as polars is not installed.")
+    msg = "Polars compiler not available as polars is not installed."
+    warnings.warn(msg)

bigframes/core/global_session.py

@@ -22,6 +22,7 @@
 import google.auth.exceptions
 
 import bigframes._config
+import bigframes.exceptions as bfe
 import bigframes.session
 
 _global_session: Optional[bigframes.session.Session] = None
@@ -38,11 +39,11 @@ def _try_close_session(session: bigframes.session.Session):
         session_id = session.session_id
         location = session._location
         project_id = session._project
-        warnings.warn(
+        msg = (
             f"Session cleanup failed for session with id: {session_id}, "
-            f"location: {location}, project: {project_id}",
-            category=bigframes.exceptions.CleanupFailedWarning,
+            f"location: {location}, project: {project_id}"
         )
+        warnings.warn(msg, category=bfe.CleanupFailedWarning)
         traceback.print_tb(e.__traceback__)
 
 

bigframes/core/indexers.py

@@ -29,7 +29,7 @@
 import bigframes.core.scalar
 import bigframes.dataframe
 import bigframes.dtypes
-import bigframes.exceptions
+import bigframes.exceptions as bfe
 import bigframes.operations as ops
 import bigframes.series
 
@@ -407,11 +407,12 @@ def _struct_accessor_check_and_warn(
         return
 
     if not bigframes.dtypes.is_string_like(series.index.dtype):
-        warnings.warn(
-            "Are you trying to access struct fields? If so, please use Series.struct.field(...) method instead.",
-            category=bigframes.exceptions.BadIndexerKeyWarning,
-            stacklevel=7,  # Stack depth from series.__getitem__ to here
+        msg = (
+            "Are you trying to access struct fields? If so, please use Series.struct.field(...) "
+            "method instead."
         )
+        # Stack depth from series.__getitem__ to here
+        warnings.warn(msg, stacklevel=7, category=bfe.BadIndexerKeyWarning)
 
 
 @typing.overload

bigframes/core/utils.py

@@ -21,7 +21,7 @@
 import pandas as pd
 import typing_extensions
 
-import bigframes.exceptions as exc
+import bigframes.exceptions as bfe
 
 UNNAMED_COLUMN_ID = "bigframes_unnamed_column"
 UNNAMED_INDEX_ID = "bigframes_unnamed_index"
@@ -170,11 +170,6 @@ def merge_column_labels(
     return pd.Index(result_labels)
 
 
-def warn_preview(msg=""):
-    """Warn a preview API."""
-    warnings.warn(msg, exc.PreviewWarning)
-
-
 def preview(*, name: str):
     """Decorate to warn of a preview API."""
 
@@ -183,7 +178,7 @@ def decorator(func):
 
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
-            warn_preview(msg=msg)
+            warnings.warn(msg, category=bfe.PreviewWarning)
             return func(*args, **kwargs)
 
         return wrapper

bigframes/dataframe.py

@@ -68,7 +68,7 @@
 import bigframes.core.window
 import bigframes.core.window_spec as windows
 import bigframes.dtypes
-import bigframes.exceptions
+import bigframes.exceptions as bfe
 import bigframes.formatting_helpers as formatter
 import bigframes.operations as ops
 import bigframes.operations.aggregations
@@ -1481,10 +1481,8 @@ def to_arrow(
         Returns:
             pyarrow.Table: A pyarrow Table with all rows and columns of this DataFrame.
         """
-        warnings.warn(
-            "to_arrow is in preview. Types and unnamed / duplicate name columns may change in future.",
-            category=bigframes.exceptions.PreviewWarning,
-        )
+        msg = "to_arrow is in preview. Types and unnamed / duplicate name columns may change in future."
+        warnings.warn(msg, category=bfe.PreviewWarning)
 
         pa_table, query_job = self._block.to_arrow(ordered=ordered)
         self._set_internal_query_job(query_job)
@@ -3920,10 +3918,8 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame:
 
     def apply(self, func, *, axis=0, args: typing.Tuple = (), **kwargs):
         if utils.get_axis_number(axis) == 1:
-            warnings.warn(
-                "axis=1 scenario is in preview.",
-                category=bigframes.exceptions.PreviewWarning,
-            )
+            msg = "axis=1 scenario is in preview."
+            warnings.warn(msg, category=bfe.PreviewWarning)
 
             # Check if the function is a remote function
             if not hasattr(func, "bigframes_remote_function"):