From 481b58c9a273b8e8ae5530b210b049075ddc52e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 20:33:35 +0000 Subject: [PATCH 01/11] feat: create `Series.bigquery.function_name` accessors for array and AEAD functions --- packages/bigframes/bigframes/__init__.py | 1 + .../extensions/bigframes/__init__.py | 9 +- .../extensions/bigframes/series_accessor.py | 59 ++ .../extensions/core/series_accessor.py | 704 ++++++++++++++++++ .../bigframes/extensions/pandas/__init__.py | 8 +- .../extensions/pandas/series_accessor.py | 68 ++ .../googlesql/global_namespace/array.py | 44 +- packages/bigframes/bigframes/series.py | 31 +- packages/bigframes/docs/reference/index.rst | 4 +- .../scripts/data/sql-functions/aead.yaml | 3 + .../global_namespace/aead_encryption.yaml | 3 + .../sql-functions/global_namespace/array.yaml | 57 +- .../scripts/generate_bigframes_bigquery.py | 155 +++- .../templates/bigframes_series_accessor.py.j2 | 40 + .../templates/core_series_accessor.py.j2 | 83 +++ .../templates/pandas_series_accessor.py.j2 | 49 ++ .../unit/extensions/bigframes/__init__.py | 13 + .../bigframes/test_series_accessor.py | 78 ++ .../extensions/pandas/test_series_accessor.py | 137 ++++ 19 files changed, 1536 insertions(+), 10 deletions(-) create mode 100644 packages/bigframes/bigframes/extensions/bigframes/series_accessor.py create mode 100644 packages/bigframes/bigframes/extensions/core/series_accessor.py create mode 100644 packages/bigframes/bigframes/extensions/pandas/series_accessor.py create mode 100644 packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 create mode 100644 packages/bigframes/scripts/templates/core_series_accessor.py.j2 create mode 100644 packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 create mode 100644 packages/bigframes/tests/unit/extensions/bigframes/__init__.py create mode 100644 packages/bigframes/tests/unit/extensions/bigframes/test_series_accessor.py create mode 100644 packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py diff --git a/packages/bigframes/bigframes/__init__.py b/packages/bigframes/bigframes/__init__.py index 7061300b5cc5..533726343a59 100644 --- a/packages/bigframes/bigframes/__init__.py +++ b/packages/bigframes/bigframes/__init__.py @@ -42,6 +42,7 @@ # Register pandas extensions import bigframes.extensions.pandas.dataframe_accessor # noqa: F401, E402 +import bigframes.extensions.pandas.series_accessor # noqa: F401, E402 from bigframes._config.bigquery_options import BigQueryOptions # noqa: E402 from bigframes.core.global_session import ( # noqa: E402 close_session, diff --git a/packages/bigframes/bigframes/extensions/bigframes/__init__.py b/packages/bigframes/bigframes/extensions/bigframes/__init__.py index 859b51d71ca8..439a8189dedf 100644 --- a/packages/bigframes/bigframes/extensions/bigframes/__init__.py +++ b/packages/bigframes/bigframes/extensions/bigframes/__init__.py @@ -16,5 +16,12 @@ BigframesAIAccessor, BigframesBigQueryDataFrameAccessor, ) +from bigframes.extensions.bigframes.series_accessor import ( + BigframesBigQuerySeriesAccessor, +) -__all__ = ["BigframesAIAccessor", "BigframesBigQueryDataFrameAccessor"] +__all__ = [ + "BigframesAIAccessor", + "BigframesBigQueryDataFrameAccessor", + "BigframesBigQuerySeriesAccessor", +] diff --git a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py new file mode 100644 index 000000000000..45e17a970674 --- /dev/null +++ b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py @@ -0,0 +1,59 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# + +from __future__ import annotations + +from typing import Optional, TypeVar, cast + +import bigframes.extensions.core.series_accessor as core_accessor +import bigframes.series +from bigframes.core.logging import log_adapter + +S = TypeVar("S", bound="bigframes.series.Series") + + +@log_adapter.class_logger +class BigframesBigQuerySeriesAccessor(core_accessor.BigQuerySeriesAccessor[S]): + def __init__(self, bf_obj: S): + super().__init__(bf_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + return self._obj + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series) + + @property + def aead(self) -> BigframesAeadSeriesAccessor[S]: + return BigframesAeadSeriesAccessor(self._obj) + + +@log_adapter.class_logger +class BigframesAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[S]): + def __init__(self, bf_obj: S): + super().__init__(bf_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + return self._obj + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series) diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py new file mode 100644 index 000000000000..580fc7cc3e46 --- /dev/null +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -0,0 +1,704 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# + +from __future__ import annotations + +import abc +from typing import Any, Generic, Literal, Optional, TypeVar, Union, cast + +import bigframes.core.col +import bigframes.core.sentinels as sentinels +import bigframes.series as series +from bigframes import dtypes + +S = TypeVar("S") + + +class AbstractBigQuerySeriesAccessor(abc.ABC, Generic[S]): + def __init__(self, obj: S): + self._obj = obj + + @abc.abstractmethod + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> series.Series: + """Convert the accessor's object to a BigFrames Series.""" + + @abc.abstractmethod + def _to_series(self, bf_series: series.Series) -> S: + """Convert a BigFrames Series to the accessor's object type.""" + + +class BigQuerySeriesAccessor(AbstractBigQuerySeriesAccessor[S]): + """Series accessor for BigQuery functions.""" + + @property + @abc.abstractmethod + def aead(self) -> AeadSeriesAccessor[S]: + """Accessor for BigQuery aead functions.""" + + def deterministic_decrypt_bytes( + self, + ciphertext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Uses the matching key from `keyset` to decrypt `ciphertext` and verifies the integrity of the data using `additional_data`. Returns an error if decryption fails.""" + from bigframes.operations.googlesql.global_namespace.aead_encryption import ( + deterministic_decrypt_bytes as deterministic_decrypt_bytes_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + ciphertext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = deterministic_decrypt_bytes_impl( + bf_series, + ciphertext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) + + def deterministic_decrypt_string( + self, + ciphertext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING.""" + from bigframes.operations.googlesql.global_namespace.aead_encryption import ( + deterministic_decrypt_string as deterministic_decrypt_string_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + ciphertext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = deterministic_decrypt_string_impl( + bf_series, + ciphertext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) + + def deterministic_encrypt( + self, + plaintext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Encrypts `plaintext` using the primary cryptographic key in `keyset` using deterministic AEAD. The algorithm of the primary key must be `DETERMINISTIC_AEAD_AES_SIV_CMAC_256`. Binds the ciphertext to the context defined by `additional_data`. Returns `NULL` if any input is `NULL`.""" + from bigframes.operations.googlesql.global_namespace.aead_encryption import ( + deterministic_encrypt as deterministic_encrypt_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + plaintext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = deterministic_encrypt_impl( + bf_series, + plaintext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) + + def array_concat( + self, + array_expression_2: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Any, Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Concatenates one or more arrays with the same element type into a single array.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_concat as array_concat_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + array_expression_2, + ) + + bf_series = self._bf_from_series(session) + result = array_concat_impl( + bf_series, + array_expression_2, + ) + return self._to_series(cast(series.Series, result)) + + def array_first( + self, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array and returns the first element in the array.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_first as array_first_impl, + ) + + bf_series = self._bf_from_series(session) + result = array_first_impl( + bf_series, + ) + return self._to_series(cast(series.Series, result)) + + def array_first_n( + self, + n: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], int], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Returns a prefix of `input_array` consisting of the first `n` elements.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_first_n as array_first_n_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + n, + ) + + bf_series = self._bf_from_series(session) + result = array_first_n_impl( + bf_series, + n, + ) + return self._to_series(cast(series.Series, result)) + + def array_includes( + self, + search_value: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Any, Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array and returns `TRUE` if there is an element in the array that is equal to the search_value.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_includes as array_includes_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + search_value, + ) + + bf_series = self._bf_from_series(session) + result = array_includes_impl( + bf_series, + search_value, + ) + return self._to_series(cast(series.Series, result)) + + def array_includes_all( + self, + search_values: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Any, Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array to search and an array of search values. Returns `TRUE` if all search values are in the array to search, otherwise returns `FALSE`.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_includes_all as array_includes_all_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + search_values, + ) + + bf_series = self._bf_from_series(session) + result = array_includes_all_impl( + bf_series, + search_values, + ) + return self._to_series(cast(series.Series, result)) + + def array_includes_any( + self, + search_values: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Any, Literal[sentinels.Sentinel.ARGUMENT_DEFAULT]], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array to search and an array of search values. Returns `TRUE` if any search values are in the array to search, otherwise returns `FALSE`.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_includes_any as array_includes_any_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + search_values, + ) + + bf_series = self._bf_from_series(session) + result = array_includes_any_impl( + bf_series, + search_values, + ) + return self._to_series(cast(series.Series, result)) + + def array_is_distinct( + self, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Returns `TRUE` if the array contains no repeated elements, using the same equality comparison logic as `SELECT DISTINCT`.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_is_distinct as array_is_distinct_impl, + ) + + bf_series = self._bf_from_series(session) + result = array_is_distinct_impl( + bf_series, + ) + return self._to_series(cast(series.Series, result)) + + def array_last( + self, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array and returns the last element in the array.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_last as array_last_impl, + ) + + bf_series = self._bf_from_series(session) + result = array_last_impl( + bf_series, + ) + return self._to_series(cast(series.Series, result)) + + def array_length( + self, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Compute the length of each array element in the Series. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import bigframes.bigquery as bbq + + >>> s = bpd.Series([[1, 2, 8, 3], [], [3, 4]]) + >>> bbq.array_length(s) + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([[1, 2, 8, 3], [], [3, 4]]) + >>> ps.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also apply this function directly to Series using `apply`. + + >>> s.apply(bbq.array_length, by_row=False) + 0 4 + 1 0 + 2 2 + dtype: Int64 + + Args: + series (bigframes.series.Series): A Series with array columns. + + Returns: + bigframes.series.Series: A Series of integer values indicating + the length of each element in the Series. + """ + from bigframes.operations.googlesql.global_namespace.array import ( + array_length as array_length_impl, + ) + + bf_series = self._bf_from_series(session) + result = array_length_impl( + bf_series, + ) + return self._to_series(cast(series.Series, result)) + + def array_reverse( + self, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Returns the input `ARRAY` with elements in reverse order.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_reverse as array_reverse_impl, + ) + + bf_series = self._bf_from_series(session) + result = array_reverse_impl( + bf_series, + ) + return self._to_series(cast(series.Series, result)) + + def array_slice( + self, + start_offset: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], int], + ], + end_offset: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], int], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Returns an array containing zero or more consecutive elements from the input array.""" + from bigframes.operations.googlesql.global_namespace.array import ( + array_slice as array_slice_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + start_offset, + end_offset, + ) + + bf_series = self._bf_from_series(session) + result = array_slice_impl( + bf_series, + start_offset, + end_offset, + ) + return self._to_series(cast(series.Series, result)) + + def array_to_string( + self, + delimiter: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + null_text: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ] = sentinels.Sentinel.ARGUMENT_DEFAULT, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Converts array elements within a Series into delimited strings. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import bigframes.bigquery as bbq + + >>> s = bpd.Series([["H", "i", "!"], ["Hello", "World"], np.nan, [], ["Hi"]]) + >>> bbq.array_to_string(s, delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([["H", "i", "!"], ["Hello", "World"], None, [], ["Hi"]]) + >>> ps.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + Args: + series (bigframes.series.Series): A Series containing arrays. + delimiter (str): The string used to separate array elements. + null_text (str, optional): The string to replace any NULL values in the array with. + + Returns: + bigframes.series.Series: A Series containing delimited strings. + """ + from bigframes.operations.googlesql.global_namespace.array import ( + array_to_string as array_to_string_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + delimiter, + null_text, + ) + + bf_series = self._bf_from_series(session) + result = array_to_string_impl( + bf_series, + delimiter, + null_text, + ) + return self._to_series(cast(series.Series, result)) + + def flatten( + self, + depth: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], int], + ] = sentinels.Sentinel.ARGUMENT_DEFAULT, + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Takes an array of nested data and flattens a specific part of it into a single, flat array with the [array elements field access operator][array-el-field-operator]. Returns `NULL` if the input value is `NULL`.""" + from bigframes.operations.googlesql.global_namespace.array import ( + flatten as flatten_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + depth, + ) + + bf_series = self._bf_from_series(session) + result = flatten_impl( + bf_series, + depth, + ) + return self._to_series(cast(series.Series, result)) + + +class AeadSeriesAccessor(AbstractBigQuerySeriesAccessor[S]): + """Series accessor for BigQuery aead functions.""" + + def decrypt_bytes( + self, + ciphertext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Uses the matching key from keyset to decrypt ciphertext and verifies the integrity of the data using additional_data. Returns an error if decryption or verification fails.""" + from bigframes.operations.googlesql.aead import ( + decrypt_bytes as decrypt_bytes_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + ciphertext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = decrypt_bytes_impl( + bf_series, + ciphertext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) + + def decrypt_string( + self, + ciphertext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Like AEAD.DECRYPT_BYTES, but where additional_data is of type STRING.""" + from bigframes.operations.googlesql.aead import ( + decrypt_string as decrypt_string_impl, + ) + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + ciphertext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = decrypt_string_impl( + bf_series, + ciphertext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) + + def encrypt( + self, + plaintext: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + additional_data: Union[ + series.Series, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """Encrypts plaintext using the primary cryptographic key in keyset. The algorithm of the primary key must be AEAD_AES_GCM_256. Binds the ciphertext to the context defined by additional_data. Returns NULL if any input is NULL.""" + from bigframes.operations.googlesql.aead import encrypt as encrypt_impl + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + + session = googlesql._find_session( + plaintext, + additional_data, + ) + + bf_series = self._bf_from_series(session) + result = encrypt_impl( + bf_series, + plaintext, + additional_data, + ) + return self._to_series(cast(series.Series, result)) diff --git a/packages/bigframes/bigframes/extensions/pandas/__init__.py b/packages/bigframes/bigframes/extensions/pandas/__init__.py index d47acd3b05e0..6af1f769b5ba 100644 --- a/packages/bigframes/bigframes/extensions/pandas/__init__.py +++ b/packages/bigframes/bigframes/extensions/pandas/__init__.py @@ -21,5 +21,11 @@ from bigframes.extensions.pandas.dataframe_accessor import ( PandasBigQueryDataFrameAccessor, ) +from bigframes.extensions.pandas.series_accessor import ( + PandasBigQuerySeriesAccessor, +) -__all__ = ["PandasBigQueryDataFrameAccessor"] +__all__ = [ + "PandasBigQueryDataFrameAccessor", + "PandasBigQuerySeriesAccessor", +] diff --git a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py new file mode 100644 index 000000000000..09cb8be8dcaa --- /dev/null +++ b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py @@ -0,0 +1,68 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: scripts/generate_bigframes_bigquery.py +# + +from __future__ import annotations + +from typing import Optional, TypeVar, cast + +import pandas +import pandas.api.extensions + +import bigframes.core.global_session as bf_session +import bigframes.extensions.core.series_accessor as core_accessor +import bigframes.series +from bigframes.core.logging import log_adapter + +S = TypeVar("S", bound="pandas.Series") + + +@pandas.api.extensions.register_series_accessor("bigquery") +@log_adapter.class_logger +class PandasBigQuerySeriesAccessor(core_accessor.BigQuerySeriesAccessor[S]): + def __init__(self, pandas_obj: S): + super().__init__(pandas_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + if session is None: + session = bf_session.get_global_session() + return cast(bigframes.series.Series, session.read_pandas(self._obj)) + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series.to_pandas(ordered=True)) + + @property + def aead(self) -> PandasAeadSeriesAccessor[S]: + return PandasAeadSeriesAccessor(self._obj) + + +@log_adapter.class_logger +class PandasAeadSeriesAccessor(core_accessor.AeadSeriesAccessor[S]): + def __init__(self, pandas_obj: S): + super().__init__(pandas_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + if session is None: + session = bf_session.get_global_session() + return cast(bigframes.series.Series, session.read_pandas(self._obj)) + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series.to_pandas(ordered=True)) diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index 3b0f93656442..4140e3d6620e 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -696,7 +696,26 @@ def array_length( 2 2 dtype: Int64 - You can also apply this function directly to Series. + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([[1, 2, 8, 3], [], [3, 4]]) + >>> ps.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also apply this function directly to Series using `apply`. >>> s.apply(bbq.array_length, by_row=False) 0 4 @@ -790,6 +809,29 @@ def array_to_string( 4 Hi dtype: string + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([["H", "i", "!"], ["Hello", "World"], None, [], ["Hi"]]) + >>> ps.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + Args: series (bigframes.series.Series): A Series containing arrays. delimiter (str): The string used to separate array elements. diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index a28c2f14cc9d..4bdd9092b451 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -42,6 +42,7 @@ import bigframes_vendored.constants as constants import bigframes_vendored.pandas.core.series as vendored_pandas_series import google.cloud.bigquery as bigquery +import google.cloud.bigquery.job import numpy import pandas import pyarrow as pa @@ -80,6 +81,7 @@ from bigframes.core.window import rolling if typing.TYPE_CHECKING: + import bigframes.extensions.bigframes.series_accessor as series_bigquery_accessor import bigframes.geopandas.geoseries import bigframes.operations.datetimes as datetimes import bigframes.operations.strings as strings @@ -118,7 +120,7 @@ def __init__( *, session: Optional[bigframes.session.Session] = None, ): - self._query_job: Optional[bigquery.QueryJob] = None + self._query_job: Optional[google.cloud.bigquery.job.QueryJob] = None import bigframes.pandas # Ignore object dtype if provided, as it provides no additional @@ -301,7 +303,26 @@ def keys(self) -> indexes.Index: return self.index @property - def query_job(self) -> Optional[bigquery.QueryJob]: + def bigquery( + self, + ) -> series_bigquery_accessor.BigframesBigQuerySeriesAccessor: + """ + Accessor for BigQuery functionality. + + Returns: + bigframes.extensions.core.series_accessor.BigQuerySeriesAccessor: + Accessor that exposes BigQuery functionality on a Series, + with method names closer to SQL. + """ + # Import the accessor here to avoid circular imports. + import bigframes.extensions.bigframes.series_accessor + + return bigframes.extensions.bigframes.series_accessor.BigframesBigQuerySeriesAccessor( + self + ) + + @property + def query_job(self) -> Optional[google.cloud.bigquery.job.QueryJob]: """BigQuery job metadata for the most recent query. Returns: @@ -355,7 +376,9 @@ def sql(self) -> str: def transpose(self) -> Series: return self - def _set_internal_query_job(self, query_job: Optional[bigquery.QueryJob]): + def _set_internal_query_job( + self, query_job: Optional[google.cloud.bigquery.job.QueryJob] + ): self._query_job = query_job def __len__(self): @@ -815,7 +838,7 @@ def to_pandas_batches( ) return map(lambda df: cast(pandas.Series, df.squeeze(1)), batches) - def _compute_dry_run(self) -> bigquery.QueryJob: + def _compute_dry_run(self) -> google.cloud.bigquery.job.QueryJob: _, query_job = self._block._compute_dry_run((self._value_column,)) return query_job diff --git a/packages/bigframes/docs/reference/index.rst b/packages/bigframes/docs/reference/index.rst index 60934582e969..99228010b249 100644 --- a/packages/bigframes/docs/reference/index.rst +++ b/packages/bigframes/docs/reference/index.rst @@ -23,13 +23,15 @@ packages. Pandas Extensions ~~~~~~~~~~~~~~~~~ -BigQuery DataFrames provides extensions to pandas DataFrame objects. +BigQuery DataFrames provides extensions to pandas DataFrame and Series objects. .. autosummary:: :toctree: api bigframes.extensions.core.dataframe_accessor.BigQueryDataFrameAccessor bigframes.extensions.core.dataframe_accessor.AIAccessor + bigframes.extensions.core.series_accessor.BigQuerySeriesAccessor + bigframes.extensions.core.series_accessor.AeadSeriesAccessor ML APIs ~~~~~~~ diff --git a/packages/bigframes/scripts/data/sql-functions/aead.yaml b/packages/bigframes/scripts/data/sql-functions/aead.yaml index 6c289a96e886..198248782d7d 100644 --- a/packages/bigframes/scripts/data/sql-functions/aead.yaml +++ b/packages/bigframes/scripts/data/sql-functions/aead.yaml @@ -2,6 +2,7 @@ urn: extension:google:bq_scalar_functions scalar_functions: - name: "aead.decrypt_bytes" description: "Uses the matching key from keyset to decrypt ciphertext and verifies the integrity of the data using additional_data. Returns an error if decryption or verification fails." + series_accessor_arg: keyset impls: # Signature: aead.decrypt_bytes:vbin_vbin_vbin - args: @@ -35,6 +36,7 @@ scalar_functions: return: binary - name: "aead.decrypt_string" description: "Like AEAD.DECRYPT_BYTES, but where additional_data is of type STRING." + series_accessor_arg: keyset impls: # Signature: aead.decrypt_string:vbin_vbin_str - args: @@ -68,6 +70,7 @@ scalar_functions: return: string - name: "aead.encrypt" description: "Encrypts plaintext using the primary cryptographic key in keyset. The algorithm of the primary key must be AEAD_AES_GCM_256. Binds the ciphertext to the context defined by additional_data. Returns NULL if any input is NULL." + series_accessor_arg: keyset impls: # Signature: aead.encrypt:vbin_str_str - args: diff --git a/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml b/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml index ffd26e5e0e7b..1e62de0f2a65 100644 --- a/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml +++ b/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml @@ -2,6 +2,7 @@ urn: extension:google:bq_scalar_functions scalar_functions: - name: "deterministic_decrypt_bytes" description: "Uses the matching key from `keyset` to decrypt `ciphertext` and verifies the integrity of the data using `additional_data`. Returns an error if decryption fails." + series_accessor_arg: keyset impls: # Signature: deterministic_decrypt_bytes:vbin_vbin_vbin - args: @@ -35,6 +36,7 @@ scalar_functions: return: binary - name: "deterministic_decrypt_string" description: "Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING." + series_accessor_arg: keyset impls: # Signature: deterministic_decrypt_string:vbin_vbin_str - args: @@ -68,6 +70,7 @@ scalar_functions: return: string - name: "deterministic_encrypt" description: "Encrypts `plaintext` using the primary cryptographic key in `keyset` using deterministic AEAD. The algorithm of the primary key must be `DETERMINISTIC_AEAD_AES_SIV_CMAC_256`. Binds the ciphertext to the context defined by `additional_data`. Returns `NULL` if any input is `NULL`." + series_accessor_arg: keyset impls: # Signature: deterministic_encrypt:vbin_str_str - args: diff --git a/packages/bigframes/scripts/data/sql-functions/global_namespace/array.yaml b/packages/bigframes/scripts/data/sql-functions/global_namespace/array.yaml index a7d01a9143ce..aa9230c251be 100644 --- a/packages/bigframes/scripts/data/sql-functions/global_namespace/array.yaml +++ b/packages/bigframes/scripts/data/sql-functions/global_namespace/array.yaml @@ -2,6 +2,7 @@ urn: extension:google:bq_scalar_functions scalar_functions: - name: "array_concat" description: "Concatenates one or more arrays with the same element type into a single array." + series_accessor_arg: array_expression_1 impls: # Signature: array_concat:list_list - args: @@ -16,6 +17,7 @@ scalar_functions: return: list - name: "array_first" description: "Takes an array and returns the first element in the array." + series_accessor_arg: array_expression impls: # Signature: array_first:list - args: @@ -26,6 +28,7 @@ scalar_functions: return: any1 - name: "array_first_n" description: "Returns a prefix of `input_array` consisting of the first `n` elements." + series_accessor_arg: input_array impls: # Signature: array_first_n:list_i64 - args: @@ -40,6 +43,7 @@ scalar_functions: return: list - name: "array_includes" description: "Takes an array and returns `TRUE` if there is an element in the array that is equal to the search_value." + series_accessor_arg: array_to_search impls: # Signature: array_includes:list_any - args: @@ -54,6 +58,7 @@ scalar_functions: return: boolean - name: "array_includes_all" description: "Takes an array to search and an array of search values. Returns `TRUE` if all search values are in the array to search, otherwise returns `FALSE`." + series_accessor_arg: array_to_search impls: # Signature: array_includes_all:list_list - args: @@ -68,6 +73,7 @@ scalar_functions: return: boolean - name: "array_includes_any" description: "Takes an array to search and an array of search values. Returns `TRUE` if any search values are in the array to search, otherwise returns `FALSE`." + series_accessor_arg: array_to_search impls: # Signature: array_includes_any:list_list - args: @@ -82,6 +88,7 @@ scalar_functions: return: boolean - name: "array_is_distinct" description: "Returns `TRUE` if the array contains no repeated elements, using the same equality comparison logic as `SELECT DISTINCT`." + series_accessor_arg: array_expression impls: # Signature: array_is_distinct:list - args: @@ -92,6 +99,7 @@ scalar_functions: return: boolean - name: "array_last" description: "Takes an array and returns the last element in the array." + series_accessor_arg: array_expression impls: # Signature: array_last:list - args: @@ -116,7 +124,26 @@ scalar_functions: 2 2 dtype: Int64 - You can also apply this function directly to Series. + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([[1, 2, 8, 3], [], [3, 4]]) + >>> ps.bigquery.array_length() + 0 4 + 1 0 + 2 2 + dtype: Int64 + + You can also apply this function directly to Series using `apply`. >>> s.apply(bbq.array_length, by_row=False) 0 4 @@ -130,6 +157,7 @@ scalar_functions: Returns: bigframes.series.Series: A Series of integer values indicating the length of each element in the Series. + series_accessor_arg: series impls: # Signature: array_length:list - args: @@ -140,6 +168,7 @@ scalar_functions: return: i64 - name: "array_reverse" description: "Returns the input `ARRAY` with elements in reverse order." + series_accessor_arg: value impls: # Signature: array_reverse:list - args: @@ -150,6 +179,7 @@ scalar_functions: return: list - name: "array_slice" description: "Returns an array containing zero or more consecutive elements from the input array." + series_accessor_arg: array_to_slice impls: # Signature: array_slice:list_i64_i64 - args: @@ -184,6 +214,29 @@ scalar_functions: 4 Hi dtype: string + You can call this function using the Series `bigquery` accessor. + + >>> s.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + You can also use this accessor on a pandas Series after importing bigframes. + + >>> import bigframes + >>> import pandas as pd + >>> ps = pd.Series([["H", "i", "!"], ["Hello", "World"], None, [], ["Hi"]]) + >>> ps.bigquery.array_to_string(delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + Args: series (bigframes.series.Series): A Series containing arrays. delimiter (str): The string used to separate array elements. @@ -191,6 +244,7 @@ scalar_functions: Returns: bigframes.series.Series: A Series containing delimited strings. + series_accessor_arg: series impls: # Signature: array_to_string:list_str_str - args: @@ -224,6 +278,7 @@ scalar_functions: return: binary - name: "flatten" description: "Takes an array of nested data and flattens a specific part of it into a single, flat array with the [array elements field access operator][array-el-field-operator]. Returns `NULL` if the input value is `NULL`." + series_accessor_arg: array_to_flatten impls: # Signature: flatten:list_i64 - args: diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index afdda2a5f98b..e3053ab60453 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -131,6 +131,11 @@ def load_templates(): "test_operation": env.get_template("test_operation.py.j2"), "license": env.get_template("license.py.j2"), "signature_def": env.get_template("signature_def.py.j2"), + "core_series_accessor": env.get_template("core_series_accessor.py.j2"), + "bigframes_series_accessor": env.get_template( + "bigframes_series_accessor.py.j2" + ), + "pandas_series_accessor": env.get_template("pandas_series_accessor.py.j2"), } @@ -326,6 +331,9 @@ def parse_scalar_functions(data, module_name, signature_def_template, is_global= # Test args test_args = _get_test_args(args_by_name, arg_order) + # Read series_accessor_arg + series_accessor_arg = func_data.get("series_accessor_arg") + functions_list.append( { "name": python_name, @@ -333,6 +341,7 @@ def parse_scalar_functions(data, module_name, signature_def_template, is_global= "description": func_data["description"], "args": func_args, "test_args": test_args, + "series_accessor_arg": series_accessor_arg, } ) @@ -386,6 +395,7 @@ def process_yaml_file(yaml_file, templates): output_file = OUTPUT_DIR.joinpath(module_path).with_suffix(".py") is_global = "global_namespace" in module_path.parts + namespace = get_namespace(yaml_file) ops_list, functions_list = parse_scalar_functions( data, module_name, @@ -432,12 +442,155 @@ def process_yaml_file(yaml_file, templates): run_ruff(test_output_file) print(f" Generated {test_output_file}") + # Collect functions for Series accessor + accessor_functions = [] + for func in functions_list: + if func.get("series_accessor_arg"): + import_module = ( + f"bigframes.operations.googlesql.{'.'.join(module_path.parts)}" + ) + accessor_functions.append( + { + "name": func["name"], + "import_module": import_module, + "namespace": namespace, + "description": func["description"], + "args": func["args"], + "series_accessor_arg": func["series_accessor_arg"], + } + ) + + return accessor_functions + + +def get_namespace(yaml_file: pathlib.Path) -> tuple[str, ...] | None: + rel_path = yaml_file.relative_to(DATA_DIR) + parts = rel_path.with_suffix("").parts + if "global_namespace" in parts: + return None + return parts + + +def get_class_name(ns_tuple: tuple[str, ...], prefix: str = "") -> str: + if not ns_tuple: + return f"{prefix}BigQuerySeriesAccessor" + camel_parts = [part.capitalize() for part in ns_tuple] + return f"{prefix}{''.join(camel_parts)}SeriesAccessor" + + +def generate_series_accessors(functions: list[dict], templates: dict): + print("Generating Series accessors...") + # Find all active namespaces + active_namespaces = set() + for func in functions: + ns = func["namespace"] or () + for i in range(len(ns) + 1): + active_namespaces.add(ns[:i]) + + # Sort namespaces by depth so parents come first + sorted_namespaces = sorted(list(active_namespaces), key=len) + + # Build namespace definitions + ns_defs = [] + ns_by_tuple = {} + for ns in sorted_namespaces: + class_name = get_class_name(ns) + bf_class_name = get_class_name(ns, prefix="Bigframes") + pd_class_name = get_class_name(ns, prefix="Pandas") + + ns_def = { + "ns_tuple": ns, + "class_name": class_name, + "bigframes_class_name": bf_class_name, + "pandas_class_name": pd_class_name, + "is_root": len(ns) == 0, + "description": ( + f"Series accessor for BigQuery {'.'.join(ns)} functions." + if ns + else "Series accessor for BigQuery functions." + ), + "children": [], + "functions": [], + } + ns_defs.append(ns_def) + ns_by_tuple[ns] = ns_def + + # Populate functions + for func in functions: + ns = func["namespace"] or () + ns_by_tuple[ns]["functions"].append(func) + + # Populate children properties + for ns in sorted_namespaces: + if len(ns) > 0: + parent_ns = ns[:-1] + parent_def = ns_by_tuple[parent_ns] + child_def = ns_by_tuple[ns] + parent_def["children"].append( + { + "prop_name": ns[-1], + "class_name": child_def["class_name"], + "bigframes_class_name": child_def["bigframes_class_name"], + "pandas_class_name": child_def["pandas_class_name"], + } + ) + + # Render and write core + core_output_file = pathlib.Path("bigframes/extensions/core/series_accessor.py") + core_output_file.parent.mkdir(parents=True, exist_ok=True) + ensure_init_py( + core_output_file.parent, pathlib.Path("bigframes"), templates["license"] + ) + core_content = templates["core_series_accessor"].render( + script_path="scripts/generate_bigframes_bigquery.py", + namespaces=ns_defs, + ) + with open(core_output_file, "w") as f: + f.write(core_content) + run_ruff(core_output_file) + print(f" Generated {core_output_file}") + + # Render and write bigframes + bf_output_file = pathlib.Path("bigframes/extensions/bigframes/series_accessor.py") + bf_output_file.parent.mkdir(parents=True, exist_ok=True) + ensure_init_py( + bf_output_file.parent, pathlib.Path("bigframes"), templates["license"] + ) + bf_content = templates["bigframes_series_accessor"].render( + script_path="scripts/generate_bigframes_bigquery.py", + namespaces=ns_defs, + ) + with open(bf_output_file, "w") as f: + f.write(bf_content) + run_ruff(bf_output_file) + print(f" Generated {bf_output_file}") + + # Render and write pandas + pd_output_file = pathlib.Path("bigframes/extensions/pandas/series_accessor.py") + pd_output_file.parent.mkdir(parents=True, exist_ok=True) + ensure_init_py( + pd_output_file.parent, pathlib.Path("bigframes"), templates["license"] + ) + pd_content = templates["pandas_series_accessor"].render( + script_path="scripts/generate_bigframes_bigquery.py", + namespaces=ns_defs, + ) + with open(pd_output_file, "w") as f: + f.write(pd_content) + run_ruff(pd_output_file) + print(f" Generated {pd_output_file}") + def main(): templates = load_templates() + all_accessor_functions = [] for yaml_file in sorted(DATA_DIR.glob("**/*.yaml")): - process_yaml_file(yaml_file, templates) + accessor_funcs = process_yaml_file(yaml_file, templates) + all_accessor_functions.extend(accessor_funcs) + + if all_accessor_functions: + generate_series_accessors(all_accessor_functions, templates) if __name__ == "__main__": diff --git a/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 b/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 new file mode 100644 index 000000000000..f9307d1401e1 --- /dev/null +++ b/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 @@ -0,0 +1,40 @@ +{% include 'license.py.j2' %} + +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: {{ script_path }} +# + +from __future__ import annotations + +from typing import cast, Optional, TypeVar + +from bigframes.core.logging import log_adapter +import bigframes.extensions.core.series_accessor as core_accessor +import bigframes.series + +S = TypeVar("S", bound="bigframes.series.Series") + + +{% for ns in namespaces %} +@log_adapter.class_logger +class {{ ns.bigframes_class_name }}(core_accessor.{{ ns.class_name }}[S]): + def __init__(self, bf_obj: S): + super().__init__(bf_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + return self._obj + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series) + + {% for child in ns.children %} + @property + def {{ child.prop_name }}(self) -> {{ child.bigframes_class_name }}[S]: + return {{ child.bigframes_class_name }}(self._obj) + + {% endfor %} + +{% endfor %} diff --git a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 new file mode 100644 index 000000000000..4a3a04efa815 --- /dev/null +++ b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 @@ -0,0 +1,83 @@ +{% include 'license.py.j2' %} + +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: {{ script_path }} +# + +from __future__ import annotations + +import abc +from typing import Any, cast, Generic, Literal, Optional, TypeVar, Union + +from bigframes import dtypes +import bigframes.core.col +import bigframes.core.sentinels as sentinels +import bigframes.series as series + +S = TypeVar("S") + +class AbstractBigQuerySeriesAccessor(abc.ABC, Generic[S]): + def __init__(self, obj: S): + self._obj = obj + + @abc.abstractmethod + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> series.Series: + """Convert the accessor's object to a BigFrames Series.""" + + @abc.abstractmethod + def _to_series(self, bf_series: series.Series) -> S: + """Convert a BigFrames Series to the accessor's object type.""" + + +{% for ns in namespaces %} +class {{ ns.class_name }}(AbstractBigQuerySeriesAccessor[S]): + """{{ ns.description }}""" + + {% for child in ns.children %} + @property + @abc.abstractmethod + def {{ child.prop_name }}(self) -> {{ child.class_name }}[S]: + """Accessor for BigQuery {{ child.prop_name }} functions.""" + + {% endfor %} + {% for func in ns.functions %} + def {{ func.name }}( + self, + {% for arg in func.args if arg.name != func.series_accessor_arg %} + {{ arg.name }}: Union[series.Series, bigframes.core.col.Expression, {{ arg.type_hint }}]{% if arg.default %} = {{ arg.default }}{% endif %}, + {% endfor %} + *, + session: Optional[bigframes.session.Session] = None, + ) -> S: + """{{ func.description | indent(8) }}""" + from {{ func.import_module }} import {{ func.name }} as {{ func.name }}_impl + {% if func.args | length > 1 %} + + # Resolve session from other arguments if not passed + if session is None: + import bigframes.core.googlesql as googlesql + session = googlesql._find_session( + {% for arg in func.args if arg.name != func.series_accessor_arg %} + {{ arg.name }}, + {% endfor %} + ) + {% endif %} + + bf_series = self._bf_from_series(session) + result = {{ func.name }}_impl( + {% for arg in func.args %} + {% if arg.name == func.series_accessor_arg %} + bf_series, + {% else %} + {{ arg.name }}, + {% endif %} + {% endfor %} + ) + return self._to_series(cast(series.Series, result)) + + {% endfor %} + +{% endfor %} diff --git a/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 b/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 new file mode 100644 index 000000000000..a26fb328549e --- /dev/null +++ b/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 @@ -0,0 +1,49 @@ +{% include 'license.py.j2' %} + +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated by the script: {{ script_path }} +# + +from __future__ import annotations + +from typing import cast, Optional, TypeVar + +import pandas +import pandas.api.extensions + +import bigframes.core.global_session as bf_session +from bigframes.core.logging import log_adapter +import bigframes.extensions.core.series_accessor as core_accessor +import bigframes.series + +S = TypeVar("S", bound="pandas.Series") + + +{% for ns in namespaces %} +{% if ns.is_root %} +@pandas.api.extensions.register_series_accessor("bigquery") +{% endif %} +@log_adapter.class_logger +class {{ ns.pandas_class_name }}(core_accessor.{{ ns.class_name }}[S]): + def __init__(self, pandas_obj: S): + super().__init__(pandas_obj) + + def _bf_from_series( + self, session: Optional[bigframes.session.Session] = None + ) -> bigframes.series.Series: + if session is None: + session = bf_session.get_global_session() + return cast(bigframes.series.Series, session.read_pandas(self._obj)) + + def _to_series(self, bf_series: bigframes.series.Series) -> S: + return cast(S, bf_series.to_pandas(ordered=True)) + + {% for child in ns.children %} + @property + def {{ child.prop_name }}(self) -> {{ child.pandas_class_name }}[S]: + return {{ child.pandas_class_name }}(self._obj) + + {% endfor %} + +{% endfor %} diff --git a/packages/bigframes/tests/unit/extensions/bigframes/__init__.py b/packages/bigframes/tests/unit/extensions/bigframes/__init__.py new file mode 100644 index 000000000000..58d482ea3866 --- /dev/null +++ b/packages/bigframes/tests/unit/extensions/bigframes/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/packages/bigframes/tests/unit/extensions/bigframes/test_series_accessor.py b/packages/bigframes/tests/unit/extensions/bigframes/test_series_accessor.py new file mode 100644 index 000000000000..4c74b60a1a03 --- /dev/null +++ b/packages/bigframes/tests/unit/extensions/bigframes/test_series_accessor.py @@ -0,0 +1,78 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import cast +from unittest.mock import MagicMock, patch + +import pytest + +import bigframes.series as series +from bigframes.testing import mocks + + +def test_bigframes_series_has_accessor(monkeypatch: pytest.MonkeyPatch): + # Arrange + from bigframes.extensions.bigframes.series_accessor import ( + BigframesBigQuerySeriesAccessor, + ) + + bf_df = mocks.create_dataframe(monkeypatch, data={"col": [1, 2]}) + bf_series = cast(series.Series, bf_df["col"]) + + # Act + has_bq = hasattr(bf_series, "bigquery") + bq_obj = bf_series.bigquery + + # Assert + assert has_bq + assert isinstance(bq_obj, BigframesBigQuerySeriesAccessor) + + +@patch("bigframes.operations.googlesql.global_namespace.array.array_length") +def test_bigframes_series_accessor_global_routing( + mock_array_length, monkeypatch: pytest.MonkeyPatch +): + # Arrange + bf_df = mocks.create_dataframe(monkeypatch, data={"col": [[1, 2], [3, 4, 5]]}) + bf_series = cast(series.Series, bf_df["col"]) + mock_result_series = MagicMock() + mock_array_length.return_value = mock_result_series + + # Act + result = bf_series.bigquery.array_length() + + # Assert + mock_array_length.assert_called_once_with(bf_series) + assert result is mock_result_series + + +@patch("bigframes.operations.googlesql.aead.encrypt") +def test_bigframes_series_accessor_namespaced_routing( + mock_encrypt, monkeypatch: pytest.MonkeyPatch +): + # Arrange + bf_df = mocks.create_dataframe(monkeypatch, data={"keyset": [b"key1", b"key2"]}) + keyset_series = cast(series.Series, bf_df["keyset"]) + mock_result_series = MagicMock() + mock_encrypt.return_value = mock_result_series + + plaintext = "my secret" + additional_data = "context" + + # Act + result = keyset_series.bigquery.aead.encrypt(plaintext, additional_data) + + # Assert + mock_encrypt.assert_called_once_with(keyset_series, plaintext, additional_data) + assert result is mock_result_series diff --git a/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py b/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py new file mode 100644 index 000000000000..b459cad3a6c9 --- /dev/null +++ b/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py @@ -0,0 +1,137 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +import bigframes # noqa: F401 registers pandas extensions +import bigframes.series as series + + +def test_pandas_series_registers_accessor(): + # Arrange + from bigframes.extensions.pandas.series_accessor import ( + PandasBigQuerySeriesAccessor, + ) + + s = pd.Series([1, 2]) + + # Act + has_bq = hasattr(s, "bigquery") + bq_obj = s.bigquery + + # Assert + assert has_bq + assert isinstance(bq_obj, PandasBigQuerySeriesAccessor) + + +@patch("bigframes.operations.googlesql.global_namespace.array.array_length") +def test_pandas_series_accessor_global_routing(mock_array_length): + # Arrange + mock_bf_series = MagicMock() + mock_bf_series.to_pandas.return_value = pd.Series([2, 3]) + mock_array_length.return_value = mock_bf_series + mock_session = MagicMock() + mock_bf_self = MagicMock() + mock_session.read_pandas.return_value = mock_bf_self + + s = pd.Series([[1, 2], [3, 4, 5]]) + + # Act + result = s.bigquery.array_length(session=mock_session) + + # Assert + mock_session.read_pandas.assert_called_once_with(s) + mock_array_length.assert_called_once_with(mock_bf_self) + mock_bf_series.to_pandas.assert_called_once_with(ordered=True) + pd.testing.assert_series_equal(result, pd.Series([2, 3])) + + +@patch("bigframes.operations.googlesql.aead.encrypt") +def test_pandas_series_accessor_namespaced_routing(mock_encrypt): + # Arrange + mock_bf_series = MagicMock() + mock_bf_series.to_pandas.return_value = pd.Series([b"encrypted1", b"encrypted2"]) + mock_encrypt.return_value = mock_bf_series + mock_session = MagicMock() + mock_bf_self = MagicMock() + mock_session.read_pandas.return_value = mock_bf_self + + keyset_series = pd.Series([b"key1", b"key2"]) + plaintext = "my secret" + additional_data = "context" + + # Act + result = keyset_series.bigquery.aead.encrypt( # type: ignore + plaintext, additional_data, session=mock_session + ) + + # Assert + mock_session.read_pandas.assert_called_once_with(keyset_series) + mock_encrypt.assert_called_once_with(mock_bf_self, plaintext, additional_data) + mock_bf_series.to_pandas.assert_called_once_with(ordered=True) + pd.testing.assert_series_equal(result, pd.Series([b"encrypted1", b"encrypted2"])) + + +@patch("bigframes.operations.googlesql.global_namespace.array.array_concat") +def test_pandas_series_accessor_global_routing_uses_series_session(mock_array_concat): + # Arrange + mock_bf_series = MagicMock() + mock_bf_series.to_pandas.return_value = pd.Series([[1, 2, 3, 4]]) + mock_array_concat.return_value = mock_bf_series + mock_session = MagicMock() + mock_bf_other = MagicMock(spec=series.Series) + mock_bf_other._session = mock_session + mock_bf_self = MagicMock() + mock_session.read_pandas.return_value = mock_bf_self + s = pd.Series([[1, 2]]) + + # Act + result = s.bigquery.array_concat(mock_bf_other) + + # Assert + assert result is not None + mock_session.read_pandas.assert_called_once_with(s) + mock_array_concat.assert_called_once_with(mock_bf_self, mock_bf_other) + + +@patch("bigframes.operations.googlesql.aead.encrypt") +def test_pandas_series_accessor_namespaced_routing_uses_series_session( + mock_encrypt, +): + # Arrange + mock_bf_series = MagicMock() + mock_bf_series.to_pandas.return_value = pd.Series([b"encrypted1", b"encrypted2"]) + mock_encrypt.return_value = mock_bf_series + mock_session = MagicMock() + mock_bf_plaintext = MagicMock(spec=series.Series) + mock_bf_plaintext._session = mock_session + mock_bf_self = MagicMock() + mock_session.read_pandas.return_value = mock_bf_self + keyset_series = pd.Series([b"key1", b"key2"]) + additional_data = "context" + + # Act + result = keyset_series.bigquery.aead.encrypt( # type: ignore + mock_bf_plaintext, additional_data + ) + + # Assert + assert result is not None + mock_session.read_pandas.assert_called_once_with(keyset_series) + mock_encrypt.assert_called_once_with( + mock_bf_self, mock_bf_plaintext, additional_data + ) From 6e1604474a11f1df868613857b5f10c9064491a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 20:59:11 +0000 Subject: [PATCH 02/11] ruff fix --- packages/bigframes/bigframes/_config/auth.py | 1 - packages/bigframes/bigframes/bigquery/_operations/array.py | 1 - .../core/compile/ibis_compiler/scalar_op_compiler.py | 2 +- .../core/compile/sqlglot/expressions/generic_ops.py | 2 +- packages/bigframes/bigframes/core/nodes.py | 1 - .../bigframes/bigframes/extensions/core/series_accessor.py | 1 - packages/bigframes/bigframes/operations/__init__.py | 1 + .../bigframes/bigframes/operations/googlesql/__init__.py | 3 --- packages/bigframes/bigframes/operations/googlesql/aead.py | 6 +----- .../googlesql/global_namespace/aead_encryption.py | 6 +----- .../operations/googlesql/global_namespace/array.py | 5 +---- packages/bigframes/bigframes/series.py | 1 - .../bigframes/bigframes/session/bq_caching_executor.py | 5 ++--- packages/bigframes/bigframes/session/clients.py | 2 -- .../bigframes/bigframes/session/direct_gbq_execution.py | 7 ++----- packages/bigframes/bigframes/session/execution_spec.py | 2 +- packages/bigframes/bigframes/session/proxy_executor.py | 4 +--- .../bigframes/tests/system/small/bigquery/test_json.py | 1 - .../tests/system/small/engines/test_aggregation.py | 3 +-- .../tests/system/small/engines/test_googlesql_ops.py | 4 +--- .../bigframes/tests/system/small/engines/test_windowing.py | 4 +--- packages/bigframes/tests/unit/core/sql/test_ml.py | 4 ---- .../tests/unit/extensions/pandas/test_series_accessor.py | 1 - .../bigframes/tests/unit/session/test_proxy_executor.py | 3 --- packages/bigframes/tests/unit/test_pandas.py | 1 - 25 files changed, 15 insertions(+), 56 deletions(-) diff --git a/packages/bigframes/bigframes/_config/auth.py b/packages/bigframes/bigframes/_config/auth.py index 3b05c878db95..f1c069b5310a 100644 --- a/packages/bigframes/bigframes/_config/auth.py +++ b/packages/bigframes/bigframes/_config/auth.py @@ -23,7 +23,6 @@ import pydata_google_auth import bigframes._config.bigquery_options as bigquery_options -from bigframes._config import options _SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] diff --git a/packages/bigframes/bigframes/bigquery/_operations/array.py b/packages/bigframes/bigframes/bigquery/_operations/array.py index 7cc414879a87..0a3c5d66217a 100644 --- a/packages/bigframes/bigframes/bigquery/_operations/array.py +++ b/packages/bigframes/bigframes/bigquery/_operations/array.py @@ -24,7 +24,6 @@ import bigframes_vendored.constants as constants import bigframes.core.groupby as groupby -import bigframes.operations as ops import bigframes.operations.aggregations as agg_ops import bigframes.series as series diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py index 4108675bc11f..31a8459923c4 100644 --- a/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py +++ b/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_compiler.py @@ -308,7 +308,7 @@ def googlesql_scalar_op_impl( ) arg_spec = op.args[-1] if isinstance(operand.op(), ibis_generic.OmittedArg): - assert arg_spec.optional, f"Argument omitted, but not optional" + assert arg_spec.optional, "Argument omitted, but not optional" continue target_idx = len(final_operands) diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/generic_ops.py b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/generic_ops.py index 644e7bc365f7..a6cfb52685a0 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/expressions/generic_ops.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/expressions/generic_ops.py @@ -94,7 +94,7 @@ def _(*operands: TypedExpr, op: ops.GoogleSqlScalarOp) -> sge.Expression: ) arg_spec = op.args[-1] if operand.is_omitted: - assert arg_spec.optional, f"Argument omitted, but not optional" + assert arg_spec.optional, "Argument omitted, but not optional" continue elif arg_spec.arg_name: args.append(sge.Kwarg(this=arg_spec.arg_name, expression=operand.expr)) diff --git a/packages/bigframes/bigframes/core/nodes.py b/packages/bigframes/bigframes/core/nodes.py index e51ae61c725e..e88a78fae5cd 100644 --- a/packages/bigframes/bigframes/core/nodes.py +++ b/packages/bigframes/bigframes/core/nodes.py @@ -40,7 +40,6 @@ from bigframes.core.ordering import OrderingExpression, RowOrdering if typing.TYPE_CHECKING: - import bigframes.core.ordering as orderings import bigframes.session diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index 580fc7cc3e46..ecfb37fc4814 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -24,7 +24,6 @@ import bigframes.core.col import bigframes.core.sentinels as sentinels import bigframes.series as series -from bigframes import dtypes S = TypeVar("S") diff --git a/packages/bigframes/bigframes/operations/__init__.py b/packages/bigframes/bigframes/operations/__init__.py index dd036bec5a26..473204d0c672 100644 --- a/packages/bigframes/bigframes/operations/__init__.py +++ b/packages/bigframes/bigframes/operations/__init__.py @@ -426,6 +426,7 @@ "GeoStDistanceOp", "GeoStLengthOp", "GeoStRegionStatsOp", + "GeoStSimplifyOp", # AI ops "AIClassify", "AIGenerate", diff --git a/packages/bigframes/bigframes/operations/googlesql/__init__.py b/packages/bigframes/bigframes/operations/googlesql/__init__.py index 0100784bda1d..edec5b84f8ec 100644 --- a/packages/bigframes/bigframes/operations/googlesql/__init__.py +++ b/packages/bigframes/bigframes/operations/googlesql/__init__.py @@ -17,11 +17,8 @@ import dataclasses import typing -from enum import Enum, auto -from typing import Callable, Iterable import bigframes.operations as ops -import bigframes.operations.type as op_typing from bigframes import dtypes diff --git a/packages/bigframes/bigframes/operations/googlesql/aead.py b/packages/bigframes/bigframes/operations/googlesql/aead.py index 8c461b3ff04b..f719d7d69893 100644 --- a/packages/bigframes/bigframes/operations/googlesql/aead.py +++ b/packages/bigframes/bigframes/operations/googlesql/aead.py @@ -18,15 +18,11 @@ from __future__ import annotations -import datetime -import decimal -from typing import Any, Literal, Optional, TypeVar, Union +from typing import Literal, Union import bigframes.core.col -import bigframes.core.expression as ex import bigframes.core.googlesql import bigframes.core.sentinels as sentinels -import bigframes.operations as ops import bigframes.series as series from bigframes import dtypes from bigframes.operations import googlesql diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py index dd67027287a8..4613ddd7e6df 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/aead_encryption.py @@ -18,15 +18,11 @@ from __future__ import annotations -import datetime -import decimal -from typing import Any, Literal, Optional, TypeVar, Union +from typing import Literal, Union import bigframes.core.col -import bigframes.core.expression as ex import bigframes.core.googlesql import bigframes.core.sentinels as sentinels -import bigframes.operations as ops import bigframes.series as series from bigframes import dtypes from bigframes.operations import googlesql diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index 4140e3d6620e..aed69c0f8c8a 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -18,15 +18,12 @@ from __future__ import annotations -import datetime import decimal -from typing import Any, Literal, Optional, TypeVar, Union +from typing import Any, Literal, Union import bigframes.core.col -import bigframes.core.expression as ex import bigframes.core.googlesql import bigframes.core.sentinels as sentinels -import bigframes.operations as ops import bigframes.series as series from bigframes import dtypes from bigframes.operations import googlesql diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index 4bdd9092b451..72486c1611c6 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -41,7 +41,6 @@ import bigframes_vendored.constants as constants import bigframes_vendored.pandas.core.series as vendored_pandas_series -import google.cloud.bigquery as bigquery import google.cloud.bigquery.job import numpy import pandas diff --git a/packages/bigframes/bigframes/session/bq_caching_executor.py b/packages/bigframes/bigframes/session/bq_caching_executor.py index 9948480d5cac..19f26bd9c322 100644 --- a/packages/bigframes/bigframes/session/bq_caching_executor.py +++ b/packages/bigframes/bigframes/session/bq_caching_executor.py @@ -19,7 +19,7 @@ import dataclasses import math import threading -from typing import Literal, Mapping, Optional, Sequence, Tuple +from typing import Literal, Optional, Sequence, Tuple import google.api_core.exceptions import google.cloud.bigquery_storage_v1 @@ -41,8 +41,7 @@ import bigframes.session.metrics import bigframes.session.planner import bigframes.session.temporary_storage -from bigframes._config import ComputeOptions -from bigframes.core import bq_data, compile, guid, identifiers, local_data, rewrite +from bigframes.core import compile, guid, identifiers, local_data, rewrite from bigframes.core.compile.sqlglot import sql as sg_sql from bigframes.core.compile.sqlglot import sqlglot_ir from bigframes.session import ( diff --git a/packages/bigframes/bigframes/session/clients.py b/packages/bigframes/bigframes/session/clients.py index f3046c7cbe0f..49822bac16b6 100644 --- a/packages/bigframes/bigframes/session/clients.py +++ b/packages/bigframes/bigframes/session/clients.py @@ -14,9 +14,7 @@ """Clients manages the connection to Google APIs.""" -import os import threading -import typing from typing import Optional, Sequence, Tuple import google.api_core.client_info diff --git a/packages/bigframes/bigframes/session/direct_gbq_execution.py b/packages/bigframes/bigframes/session/direct_gbq_execution.py index 24f48fe7d54b..6b1ac76d28a2 100644 --- a/packages/bigframes/bigframes/session/direct_gbq_execution.py +++ b/packages/bigframes/bigframes/session/direct_gbq_execution.py @@ -14,7 +14,7 @@ from __future__ import annotations import asyncio -from typing import Callable, Literal, Mapping, Optional, Tuple +from typing import Literal, Mapping, Optional, Tuple import google.api_core.exceptions import google.cloud.bigquery.job as bq_job @@ -24,15 +24,12 @@ import bigframes import bigframes.core.compile -import bigframes.core.compile.ibis_compiler.ibis_compiler as ibis_compiler -import bigframes.core.compile.sqlglot.compiler as sqlglot_compiler import bigframes.core.events -import bigframes.core.schema as schemata import bigframes.session._io.bigquery as bq_io import bigframes.session.metrics from bigframes import exceptions as bfe from bigframes.core import bq_data, compile, nodes -from bigframes.core.compile.configs import CompileRequest, CompileResult +from bigframes.core.compile.configs import CompileRequest from bigframes.session import execution_spec, executor, semi_executor _WRITE_DISPOSITIONS = { diff --git a/packages/bigframes/bigframes/session/execution_spec.py b/packages/bigframes/bigframes/session/execution_spec.py index fc5c11c1c403..9a095b23a8d7 100644 --- a/packages/bigframes/bigframes/session/execution_spec.py +++ b/packages/bigframes/bigframes/session/execution_spec.py @@ -15,7 +15,7 @@ from __future__ import annotations import dataclasses -from typing import TYPE_CHECKING, Literal, Mapping, Optional, Union +from typing import Literal, Mapping, Optional, Union from google.cloud import bigquery diff --git a/packages/bigframes/bigframes/session/proxy_executor.py b/packages/bigframes/bigframes/session/proxy_executor.py index c4fe6584bd27..8f673a8bd1f4 100644 --- a/packages/bigframes/bigframes/session/proxy_executor.py +++ b/packages/bigframes/bigframes/session/proxy_executor.py @@ -14,10 +14,9 @@ from __future__ import annotations -import typing import uuid import warnings -from typing import Mapping, Optional +from typing import Optional import google.cloud.bigquery as bigquery import google.cloud.exceptions @@ -30,7 +29,6 @@ execution_spec, executor, loader, - metrics, temporary_storage, ) diff --git a/packages/bigframes/tests/system/small/bigquery/test_json.py b/packages/bigframes/tests/system/small/bigquery/test_json.py index c83d38a49067..4fc4d2283ece 100644 --- a/packages/bigframes/tests/system/small/bigquery/test_json.py +++ b/packages/bigframes/tests/system/small/bigquery/test_json.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import geopandas as gpd # type: ignore import pandas as pd diff --git a/packages/bigframes/tests/system/small/engines/test_aggregation.py b/packages/bigframes/tests/system/small/engines/test_aggregation.py index e6e4ac571578..669eae9ebf75 100644 --- a/packages/bigframes/tests/system/small/engines/test_aggregation.py +++ b/packages/bigframes/tests/system/small/engines/test_aggregation.py @@ -19,12 +19,11 @@ from bigframes.core import ( agg_expressions, array_value, - events, expression, identifiers, nodes, ) -from bigframes.session import direct_gbq_execution, polars_executor +from bigframes.session import polars_executor from bigframes.testing.engine_utils import assert_equivalence_execution pytest.importorskip("polars") diff --git a/packages/bigframes/tests/system/small/engines/test_googlesql_ops.py b/packages/bigframes/tests/system/small/engines/test_googlesql_ops.py index 2ea7070bbd30..e47308fa3557 100644 --- a/packages/bigframes/tests/system/small/engines/test_googlesql_ops.py +++ b/packages/bigframes/tests/system/small/engines/test_googlesql_ops.py @@ -12,13 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re import pytest -import bigframes.dtypes import bigframes.operations.googlesql as gsql_ops -from bigframes.core import array_value, expression +from bigframes.core import array_value from bigframes.session import polars_executor from bigframes.testing.engine_utils import assert_equivalence_execution diff --git a/packages/bigframes/tests/system/small/engines/test_windowing.py b/packages/bigframes/tests/system/small/engines/test_windowing.py index fcc1dad928f8..8235fe0ef6bf 100644 --- a/packages/bigframes/tests/system/small/engines/test_windowing.py +++ b/packages/bigframes/tests/system/small/engines/test_windowing.py @@ -13,19 +13,17 @@ # limitations under the License. import pytest -from google.cloud import bigquery import bigframes.operations.aggregations as agg_ops from bigframes.core import ( agg_expressions, array_value, - events, expression, identifiers, nodes, window_spec, ) -from bigframes.session import direct_gbq_execution, polars_executor +from bigframes.session import polars_executor from bigframes.testing.engine_utils import assert_equivalence_execution pytest.importorskip("polars") diff --git a/packages/bigframes/tests/unit/core/sql/test_ml.py b/packages/bigframes/tests/unit/core/sql/test_ml.py index d2f789fc6309..a03d8cd805ab 100644 --- a/packages/bigframes/tests/unit/core/sql/test_ml.py +++ b/packages/bigframes/tests/unit/core/sql/test_ml.py @@ -15,11 +15,7 @@ import pytest import bigframes.bigquery as bbq -import bigframes.core.col as col -import bigframes.core.expression as ex import bigframes.core.sql.ml -import bigframes.dtypes as dtypes -import bigframes.operations.numeric_ops as numeric_ops pytest.importorskip("pytest_snapshot") diff --git a/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py b/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py index b459cad3a6c9..bfb68323f6da 100644 --- a/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py +++ b/packages/bigframes/tests/unit/extensions/pandas/test_series_accessor.py @@ -15,7 +15,6 @@ from unittest.mock import MagicMock, patch import pandas as pd -import pytest import bigframes # noqa: F401 registers pandas extensions import bigframes.series as series diff --git a/packages/bigframes/tests/unit/session/test_proxy_executor.py b/packages/bigframes/tests/unit/session/test_proxy_executor.py index a1a8f168995b..8e3760cdb035 100644 --- a/packages/bigframes/tests/unit/session/test_proxy_executor.py +++ b/packages/bigframes/tests/unit/session/test_proxy_executor.py @@ -16,12 +16,9 @@ import google.cloud.bigquery as bigquery import google.cloud.exceptions -import pyarrow as pa import pytest import bigframes -import bigframes.core.nodes as nodes -import bigframes.core.schema as schemata from bigframes.session.proxy_executor import DualCompilerProxyExecutor diff --git a/packages/bigframes/tests/unit/test_pandas.py b/packages/bigframes/tests/unit/test_pandas.py index 7afb0b5dfdb3..c85d92e024da 100644 --- a/packages/bigframes/tests/unit/test_pandas.py +++ b/packages/bigframes/tests/unit/test_pandas.py @@ -14,7 +14,6 @@ import inspect import re -import sys import unittest.mock as mock import pandas as pd From 916542c884075f823d46fcd94b016c0be269ab4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:14:41 +0000 Subject: [PATCH 03/11] format --- packages/bigframes/tests/system/small/engines/conftest.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/bigframes/tests/system/small/engines/conftest.py b/packages/bigframes/tests/system/small/engines/conftest.py index c599b256b7fc..823ba9806d58 100644 --- a/packages/bigframes/tests/system/small/engines/conftest.py +++ b/packages/bigframes/tests/system/small/engines/conftest.py @@ -60,7 +60,6 @@ def bq_engine( bigquery_client: bigquery.Client, bigquery_storage_read_client: google.cloud.bigquery_storage_v1.BigQueryReadClient, ): - return direct_gbq_execution.DirectGbqExecutor( bigquery_client, bqstoragereadclient=bigquery_storage_read_client, From facc3a3d36864cf33236560f1afbc86a56668d3f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:22:27 +0000 Subject: [PATCH 04/11] fix imports --- .../bigframes/extensions/core/series_accessor.py | 14 +++++++++++++- .../operations/googlesql/global_namespace/array.py | 5 +++++ packages/bigframes/bigframes/series.py | 10 +++++----- .../scripts/templates/core_series_accessor.py.j2 | 14 +++++++++++++- 4 files changed, 36 insertions(+), 7 deletions(-) diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index ecfb37fc4814..30fab9919501 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -19,12 +19,24 @@ from __future__ import annotations import abc -from typing import Any, Generic, Literal, Optional, TypeVar, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Literal, + Optional, + TypeVar, + Union, + cast, +) import bigframes.core.col import bigframes.core.sentinels as sentinels import bigframes.series as series +if TYPE_CHECKING: + import bigframes.session + S = TypeVar("S") diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index 94adbad1839d..aed69c0f8c8a 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -298,6 +298,7 @@ def _ARRAY_TO_STRING_SIG(*args): # Pad args with None to match max expected args args = args + (None,) * (3 - len(args)) # Try matching impl 0 + any1_val = None match_ok = True if match_ok and args[0] is not None: if not dtypes.is_array_like(args[0]): @@ -334,6 +335,7 @@ def _ARRAY_TO_STRING_SIG(*args): return dtypes.STRING_DTYPE # Try matching impl 1 + any1_val = None match_ok = True if match_ok and args[0] is not None: if not dtypes.is_array_like(args[0]): @@ -427,6 +429,7 @@ def _GENERATE_ARRAY_SIG(*args): # Pad args with None to match max expected args args = args + (None,) * (3 - len(args)) # Try matching impl 0 + any1_val = None match_ok = True if match_ok and args[0] is not None: try: @@ -450,6 +453,7 @@ def _GENERATE_ARRAY_SIG(*args): return dtypes.list_type(dtypes.INT_DTYPE) # Try matching impl 1 + any1_val = None match_ok = True if match_ok and args[0] is not None: try: @@ -482,6 +486,7 @@ def _GENERATE_ARRAY_SIG(*args): return dtypes.list_type(dtypes.NUMERIC_DTYPE) # Try matching impl 2 + any1_val = None match_ok = True if match_ok and args[0] is not None: try: diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index 79977b90416b..e3a8b7e49539 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -41,7 +41,7 @@ import bigframes_vendored.constants as constants import bigframes_vendored.pandas.core.series as vendored_pandas_series -import google.cloud.bigquery as bigquery +import google.cloud.bigquery.job import numpy import pandas import pyarrow as pa @@ -119,7 +119,7 @@ def __init__( *, session: Optional[bigframes.session.Session] = None, ): - self._query_job: Optional[bigquery.QueryJob] = None + self._query_job: Optional[google.cloud.bigquery.job.QueryJob] = None import bigframes.pandas # Ignore object dtype if provided, as it provides no additional @@ -321,7 +321,7 @@ def bigquery( ) @property - def query_job(self) -> Optional[bigquery.QueryJob]: + def query_job(self) -> Optional[google.cloud.bigquery.job.QueryJob]: """BigQuery job metadata for the most recent query. Returns: @@ -375,7 +375,7 @@ def sql(self) -> str: def transpose(self) -> Series: return self - def _set_internal_query_job(self, query_job: Optional[bigquery.QueryJob]): + def _set_internal_query_job(self, query_job: Optional[google.cloud.bigquery.job.QueryJob]): self._query_job = query_job def __len__(self): @@ -835,7 +835,7 @@ def to_pandas_batches( ) return map(lambda df: cast(pandas.Series, df.squeeze(1)), batches) - def _compute_dry_run(self) -> bigquery.QueryJob: + def _compute_dry_run(self) -> google.cloud.bigquery.job.QueryJob: _, query_job = self._block._compute_dry_run((self._value_column,)) return query_job diff --git a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 index 4a3a04efa815..ee64518b21e4 100644 --- a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 +++ b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 @@ -8,13 +8,25 @@ from __future__ import annotations import abc -from typing import Any, cast, Generic, Literal, Optional, TypeVar, Union +from typing import ( + TYPE_CHECKING, + Any, + cast, + Generic, + Literal, + Optional, + TypeVar, + Union, +) from bigframes import dtypes import bigframes.core.col import bigframes.core.sentinels as sentinels import bigframes.series as series +if TYPE_CHECKING: + import bigframes.session + S = TypeVar("S") class AbstractBigQuerySeriesAccessor(abc.ABC, Generic[S]): From 23a09a4031d481380837344c0052c09169be3887 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:28:21 +0000 Subject: [PATCH 05/11] remove unused --- .../operations/googlesql/global_namespace/array.py | 5 ----- .../bigframes/scripts/generate_bigframes_bigquery.py | 11 ++++++++++- .../bigframes/scripts/templates/signature_def.py.j2 | 2 ++ 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py index aed69c0f8c8a..94adbad1839d 100644 --- a/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py +++ b/packages/bigframes/bigframes/operations/googlesql/global_namespace/array.py @@ -298,7 +298,6 @@ def _ARRAY_TO_STRING_SIG(*args): # Pad args with None to match max expected args args = args + (None,) * (3 - len(args)) # Try matching impl 0 - any1_val = None match_ok = True if match_ok and args[0] is not None: if not dtypes.is_array_like(args[0]): @@ -335,7 +334,6 @@ def _ARRAY_TO_STRING_SIG(*args): return dtypes.STRING_DTYPE # Try matching impl 1 - any1_val = None match_ok = True if match_ok and args[0] is not None: if not dtypes.is_array_like(args[0]): @@ -429,7 +427,6 @@ def _GENERATE_ARRAY_SIG(*args): # Pad args with None to match max expected args args = args + (None,) * (3 - len(args)) # Try matching impl 0 - any1_val = None match_ok = True if match_ok and args[0] is not None: try: @@ -453,7 +450,6 @@ def _GENERATE_ARRAY_SIG(*args): return dtypes.list_type(dtypes.INT_DTYPE) # Try matching impl 1 - any1_val = None match_ok = True if match_ok and args[0] is not None: try: @@ -486,7 +482,6 @@ def _GENERATE_ARRAY_SIG(*args): return dtypes.list_type(dtypes.NUMERIC_DTYPE) # Try matching impl 2 - any1_val = None match_ok = True if match_ok and args[0] is not None: try: diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index e3053ab60453..2a1699cdd471 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -46,7 +46,7 @@ "ruff", "check", "--select", - "I", + "I,F", "--fix", ] + RUFF_COMMON_ARGS RUFF_FORMAT_ARGS = [ @@ -217,6 +217,15 @@ def _validate_types(impls): def _generate_signature_def(python_name, impls, sql_name, template): + for impl in impls: + uses_any1 = False + if "any1" in str(impl["return"]): + uses_any1 = True + for arg in impl["args"]: + if "any1" in str(arg["value"]): + uses_any1 = True + impl["uses_any1"] = uses_any1 + return_types = {impl["return"] for impl in impls} # Optimization: if all impls return the same concrete type, diff --git a/packages/bigframes/scripts/templates/signature_def.py.j2 b/packages/bigframes/scripts/templates/signature_def.py.j2 index ad1871f7df6d..341b889df4b8 100644 --- a/packages/bigframes/scripts/templates/signature_def.py.j2 +++ b/packages/bigframes/scripts/templates/signature_def.py.j2 @@ -3,7 +3,9 @@ def {{ func_name }}(*args): args = args + (None,) * ({{ max_args }} - len(args)) {% for impl in impls %} # Try matching impl {{ loop.index0 }} + {% if impl.uses_any1 %} any1_val = None + {% endif %} match_ok = True {% for arg in impl.args %} {% set idx = loop.index0 %} From 00b68982175b3caacf8e391fe1effdf5d981dfd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:34:16 +0000 Subject: [PATCH 06/11] import feedback --- .../bigframes/extensions/bigframes/series_accessor.py | 1 + .../bigframes/bigframes/extensions/core/series_accessor.py | 5 +---- .../bigframes/bigframes/extensions/pandas/series_accessor.py | 1 + .../scripts/templates/bigframes_series_accessor.py.j2 | 1 + .../bigframes/scripts/templates/core_series_accessor.py.j2 | 5 +---- .../bigframes/scripts/templates/pandas_series_accessor.py.j2 | 1 + 6 files changed, 6 insertions(+), 8 deletions(-) diff --git a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py index 45e17a970674..b67d007b88e7 100644 --- a/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/bigframes/series_accessor.py @@ -22,6 +22,7 @@ import bigframes.extensions.core.series_accessor as core_accessor import bigframes.series +import bigframes.session from bigframes.core.logging import log_adapter S = TypeVar("S", bound="bigframes.series.Series") diff --git a/packages/bigframes/bigframes/extensions/core/series_accessor.py b/packages/bigframes/bigframes/extensions/core/series_accessor.py index 30fab9919501..4c0f261b83cd 100644 --- a/packages/bigframes/bigframes/extensions/core/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/core/series_accessor.py @@ -20,7 +20,6 @@ import abc from typing import ( - TYPE_CHECKING, Any, Generic, Literal, @@ -33,9 +32,7 @@ import bigframes.core.col import bigframes.core.sentinels as sentinels import bigframes.series as series - -if TYPE_CHECKING: - import bigframes.session +import bigframes.session S = TypeVar("S") diff --git a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py index 09cb8be8dcaa..837664c6e1f5 100644 --- a/packages/bigframes/bigframes/extensions/pandas/series_accessor.py +++ b/packages/bigframes/bigframes/extensions/pandas/series_accessor.py @@ -26,6 +26,7 @@ import bigframes.core.global_session as bf_session import bigframes.extensions.core.series_accessor as core_accessor import bigframes.series +import bigframes.session from bigframes.core.logging import log_adapter S = TypeVar("S", bound="pandas.Series") diff --git a/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 b/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 index f9307d1401e1..f5ed3d045485 100644 --- a/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 +++ b/packages/bigframes/scripts/templates/bigframes_series_accessor.py.j2 @@ -12,6 +12,7 @@ from typing import cast, Optional, TypeVar from bigframes.core.logging import log_adapter import bigframes.extensions.core.series_accessor as core_accessor import bigframes.series +import bigframes.session S = TypeVar("S", bound="bigframes.series.Series") diff --git a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 index ee64518b21e4..5881fe6963b9 100644 --- a/packages/bigframes/scripts/templates/core_series_accessor.py.j2 +++ b/packages/bigframes/scripts/templates/core_series_accessor.py.j2 @@ -9,7 +9,6 @@ from __future__ import annotations import abc from typing import ( - TYPE_CHECKING, Any, cast, Generic, @@ -23,9 +22,7 @@ from bigframes import dtypes import bigframes.core.col import bigframes.core.sentinels as sentinels import bigframes.series as series - -if TYPE_CHECKING: - import bigframes.session +import bigframes.session S = TypeVar("S") diff --git a/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 b/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 index a26fb328549e..76f3d4797531 100644 --- a/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 +++ b/packages/bigframes/scripts/templates/pandas_series_accessor.py.j2 @@ -16,6 +16,7 @@ import bigframes.core.global_session as bf_session from bigframes.core.logging import log_adapter import bigframes.extensions.core.series_accessor as core_accessor import bigframes.series +import bigframes.session S = TypeVar("S", bound="pandas.Series") From 6e9ef19a30755e1ee0c632590dd7ce5e4f896a0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:40:42 +0000 Subject: [PATCH 07/11] revert gapic gen changes --- .../gapic/templates/noxfile.py.j2 | 10 +- .../gapic/%name_%version/%sub/test_macros.j2 | 10 +- .../integration/goldens/asset/noxfile.py | 3 +- .../unit/gapic/asset_v1/test_asset_service.py | 64 +++++------ .../goldens/credentials/noxfile.py | 3 +- .../credentials_v1/test_iam_credentials.py | 47 ++++---- .../integration/goldens/eventarc/noxfile.py | 3 +- .../unit/gapic/eventarc_v1/test_eventarc.py | 101 +++++++++--------- .../integration/goldens/logging/noxfile.py | 3 +- .../logging_v2/test_config_service_v2.py | 54 +++++----- .../logging_v2/test_logging_service_v2.py | 51 +++++---- .../logging_v2/test_metrics_service_v2.py | 50 +++++---- .../goldens/logging_internal/noxfile.py | 3 +- .../logging_v2/test_config_service_v2.py | 54 +++++----- .../logging_v2/test_logging_service_v2.py | 51 +++++---- .../logging_v2/test_metrics_service_v2.py | 50 +++++---- .../integration/goldens/redis/noxfile.py | 3 +- .../unit/gapic/redis_v1/test_cloud_redis.py | 69 ++++++------ .../goldens/redis_selective/noxfile.py | 3 +- .../unit/gapic/redis_v1/test_cloud_redis.py | 69 ++++++------ .../goldens/storagebatchoperations/noxfile.py | 3 +- .../test_storage_batch_operations.py | 65 ++++++----- 22 files changed, 381 insertions(+), 388 deletions(-) diff --git a/packages/gapic-generator/gapic/templates/noxfile.py.j2 b/packages/gapic-generator/gapic/templates/noxfile.py.j2 index f6c17446780d..c240871b994e 100644 --- a/packages/gapic-generator/gapic/templates/noxfile.py.j2 +++ b/packages/gapic-generator/gapic/templates/noxfile.py.j2 @@ -158,7 +158,7 @@ def lint(session): "ruff", "format", "--check", f"--target-version=py{ALL_PYTHON[0].replace('.', '')}", - "--line-length=88", + "--line-length=88", *LINT_PATHS, ) @@ -174,7 +174,7 @@ def lint(session): def blacken(session): """(Deprecated) Legacy session. Please use 'nox -s format'.""" session.log("WARNING: The 'blacken' session is deprecated and will be removed in a future release. Please use 'nox -s format' in the future.") - + # Just run the ruff formatter (keeping legacy behavior of only formatting, not sorting imports) session.install(RUFF_VERSION) session.run( @@ -513,14 +513,14 @@ def prerelease_deps(session, protobuf_implementation): # Extract the base package name, safely ignoring version bounds and spaces # (e.g., "grpcio>=1.75.1" becomes "grpcio") parsed_deps = { - dep: re.match(r"^([a-zA-Z0-9_-]+)", dep).group(1) + dep: re.match(r"^([a-zA-Z0-9_-]+)", dep).group(1) for dep in prerel_deps } # Dynamically sort local packages vs PyPI dependencies local_paths = [] pypi_deps = [] - + for dep, pkg_name in parsed_deps.items(): if (deps_dir / pkg_name).exists(): local_paths.append(str(deps_dir / pkg_name)) @@ -625,7 +625,7 @@ def core_deps_from_source(session, protobuf_implementation): # Batch the pip installation to avoid sequential overhead dep_paths = [str(deps_dir / dep) for dep in core_dependencies_from_source] - + session.install(*dep_paths, "--no-deps", "--ignore-installed") print(f"Installed {', '.join(core_dependencies_from_source)} locally from {deps_dir}") diff --git a/packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2 b/packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2 index 7ed31142cc9f..a612b9397f65 100644 --- a/packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2 +++ b/packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_macros.j2 @@ -1222,7 +1222,7 @@ def test_{{ method_name }}_rest_required_fields(request_type={{ method.input.ide if key == "{{ auto_populated_field|camel_case }}": assert _UUID4_RE.match(value) break - + # Include {{ auto_populated_field|camel_case }} within expected_params with value mock.ANY expected_params = [p for p in expected_params if p[0] != "{{ auto_populated_field|camel_case }}"] expected_params.append( @@ -1723,7 +1723,7 @@ def test_unsupported_parameter_rest_asyncio(): # TODO(https://github.com/googleapis/gapic-generator-python/issues/2142): Continue migrating the test cases # in macro::run_transport_tests_for_config into here, and then delete that macro in favor of this one. # TODO(https://github.com/googleapis/gapic-generator-python/issues/2153): As a follow up, migrate gRPC test cases -# into `run_transport_tests_for_config` and make any of the rest specific specific macros which are called within more generic. +# into `run_transport_tests_for_config` and make any of the rest specific specific macros which are called within more generic. #} {% macro run_transport_tests_for_config(service, api, transport, is_async) %} {% for method in service.methods.values() %} @@ -1779,7 +1779,7 @@ def test_unsupported_parameter_rest_asyncio(): {% endmacro %} {# initialize_client_with_transport_test adds coverage for transport clients. - # Note: This test case is needed because we aren't unconditionally + # Note: This test case is needed because we aren't unconditionally # generating the not implemented coverage test for every client. #} {% macro initialize_client_with_transport_test(service, transport, is_async) %} @@ -2082,7 +2082,7 @@ def test_initialize_client_w_{{transport_name}}(): assert response.raw_page is response {% endif %} - + {% if method.server_streaming %} {% if is_async %} assert isinstance(response, AsyncIterable) @@ -2092,7 +2092,7 @@ def test_initialize_client_w_{{transport_name}}(): response = next(response) {% endif %} {% endif %} - + # Establish that the response is the type that we expect. {% if method.void %} assert response is None diff --git a/packages/gapic-generator/tests/integration/goldens/asset/noxfile.py b/packages/gapic-generator/tests/integration/goldens/asset/noxfile.py index 931c366b9253..93e185b59d11 100755 --- a/packages/gapic-generator/tests/integration/goldens/asset/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/asset/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py b/packages/gapic-generator/tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py index 9fe39fcc151f..b6bbf1cdc830 100755 --- a/packages/gapic-generator/tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py +++ b/packages/gapic-generator/tests/integration/goldens/asset/tests/unit/gapic/asset_v1/test_asset_service.py @@ -13,23 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule -from requests import PreparedRequest, Request, Response +from proto.marshal.rules import wrappers +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -37,6 +40,26 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.asset_v1.services.asset_service import AssetServiceAsyncClient +from google.cloud.asset_v1.services.asset_service import AssetServiceClient +from google.cloud.asset_v1.services.asset_service import pagers +from google.cloud.asset_v1.services.asset_service import transports +from google.cloud.asset_v1.types import asset_service +from google.cloud.asset_v1.types import assets +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api_core.operation_async as operation_async # type: ignore import google.auth import google.protobuf.duration_pb2 as duration_pb2 # type: ignore @@ -44,29 +67,8 @@ import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore import google.rpc.status_pb2 as status_pb2 # type: ignore import google.type.expr_pb2 as expr_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.asset_v1.services.asset_service import ( - AssetServiceAsyncClient, - AssetServiceClient, - pagers, - transports, -) -from google.cloud.asset_v1.types import asset_service, assets -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/credentials/noxfile.py b/packages/gapic-generator/tests/integration/goldens/credentials/noxfile.py index 635c5d011d33..c991842b24ca 100755 --- a/packages/gapic-generator/tests/integration/goldens/credentials/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/credentials/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/credentials/tests/unit/gapic/credentials_v1/test_iam_credentials.py b/packages/gapic-generator/tests/integration/goldens/credentials/tests/unit/gapic/credentials_v1/test_iam_credentials.py index 8409ff7e1f19..14a6074a40f9 100755 --- a/packages/gapic-generator/tests/integration/goldens/credentials/tests/unit/gapic/credentials_v1/test_iam_credentials.py +++ b/packages/gapic-generator/tests/integration/goldens/credentials/tests/unit/gapic/credentials_v1/test_iam_credentials.py @@ -13,23 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule -from requests import PreparedRequest, Request, Response +from proto.marshal.rules import wrappers +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -37,27 +40,25 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False -import google.auth -import google.protobuf.duration_pb2 as duration_pb2 # type: ignore -import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - path_template, -) +from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials from google.auth.exceptions import MutualTLSChannelError -from google.iam.credentials_v1.services.iam_credentials import ( - IAMCredentialsAsyncClient, - IAMCredentialsClient, - transports, -) +from google.iam.credentials_v1.services.iam_credentials import IAMCredentialsAsyncClient +from google.iam.credentials_v1.services.iam_credentials import IAMCredentialsClient +from google.iam.credentials_v1.services.iam_credentials import transports from google.iam.credentials_v1.types import common from google.oauth2 import service_account +import google.auth +import google.protobuf.duration_pb2 as duration_pb2 # type: ignore +import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/eventarc/noxfile.py b/packages/gapic-generator/tests/integration/goldens/eventarc/noxfile.py index 1f74b0f3c52f..1ec5368a9dd4 100755 --- a/packages/gapic-generator/tests/integration/goldens/eventarc/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/eventarc/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/eventarc/tests/unit/gapic/eventarc_v1/test_eventarc.py b/packages/gapic-generator/tests/integration/goldens/eventarc/tests/unit/gapic/eventarc_v1/test_eventarc.py index 1c0738c28abb..538dd2b2bac4 100755 --- a/packages/gapic-generator/tests/integration/goldens/eventarc/tests/unit/gapic/eventarc_v1/test_eventarc.py +++ b/packages/gapic-generator/tests/integration/goldens/eventarc/tests/unit/gapic/eventarc_v1/test_eventarc.py @@ -13,23 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule -from requests import PreparedRequest, Request, Response +from proto.marshal.rules import wrappers +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -37,64 +40,56 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False -import google.api_core.operation_async as operation_async # type: ignore -import google.auth -import google.protobuf.duration_pb2 as duration_pb2 # type: ignore -import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore -import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -import google.rpc.code_pb2 as code_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) +from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials from google.auth.exceptions import MutualTLSChannelError -from google.cloud.eventarc_v1.services.eventarc import ( - EventarcAsyncClient, - EventarcClient, - pagers, - transports, -) -from google.cloud.eventarc_v1.types import ( - channel, - channel_connection, - discovery, - enrollment, - eventarc, - google_api_source, - google_channel_config, - logging_config, - message_bus, - network_config, - pipeline, - trigger, -) +from google.cloud.eventarc_v1.services.eventarc import EventarcAsyncClient +from google.cloud.eventarc_v1.services.eventarc import EventarcClient +from google.cloud.eventarc_v1.services.eventarc import pagers +from google.cloud.eventarc_v1.services.eventarc import transports +from google.cloud.eventarc_v1.types import channel from google.cloud.eventarc_v1.types import channel as gce_channel +from google.cloud.eventarc_v1.types import channel_connection from google.cloud.eventarc_v1.types import channel_connection as gce_channel_connection +from google.cloud.eventarc_v1.types import discovery +from google.cloud.eventarc_v1.types import enrollment from google.cloud.eventarc_v1.types import enrollment as gce_enrollment +from google.cloud.eventarc_v1.types import eventarc +from google.cloud.eventarc_v1.types import google_api_source from google.cloud.eventarc_v1.types import google_api_source as gce_google_api_source -from google.cloud.eventarc_v1.types import ( - google_channel_config as gce_google_channel_config, -) +from google.cloud.eventarc_v1.types import google_channel_config +from google.cloud.eventarc_v1.types import google_channel_config as gce_google_channel_config +from google.cloud.eventarc_v1.types import logging_config +from google.cloud.eventarc_v1.types import message_bus from google.cloud.eventarc_v1.types import message_bus as gce_message_bus +from google.cloud.eventarc_v1.types import network_config +from google.cloud.eventarc_v1.types import pipeline from google.cloud.eventarc_v1.types import pipeline as gce_pipeline +from google.cloud.eventarc_v1.types import trigger from google.cloud.eventarc_v1.types import trigger as gce_trigger from google.cloud.location import locations_pb2 -from google.iam.v1 import ( - iam_policy_pb2, # type: ignore - options_pb2, # type: ignore - policy_pb2, # type: ignore -) -from google.longrunning import operations_pb2 # type: ignore +from google.iam.v1 import iam_policy_pb2 # type: ignore +from google.iam.v1 import options_pb2 # type: ignore +from google.iam.v1 import policy_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account +import google.api_core.operation_async as operation_async # type: ignore +import google.auth +import google.protobuf.duration_pb2 as duration_pb2 # type: ignore +import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore +import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore +import google.rpc.code_pb2 as code_pb2 # type: ignore + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging/noxfile.py b/packages/gapic-generator/tests/integration/goldens/logging/noxfile.py index ea66871f9e45..448aec3ef2b0 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/logging/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_config_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_config_service_v2.py index f63bd5482b8a..57522704960d 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_config_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_config_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,34 +34,32 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False -import google.api_core.operation_async as operation_async # type: ignore -import google.auth -import google.protobuf.empty_pb2 as empty_pb2 # type: ignore -import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore -import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) +from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.config_service_v2 import ( - ConfigServiceV2AsyncClient, - ConfigServiceV2Client, - pagers, - transports, -) +from google.cloud.logging_v2.services.config_service_v2 import ConfigServiceV2AsyncClient +from google.cloud.logging_v2.services.config_service_v2 import ConfigServiceV2Client +from google.cloud.logging_v2.services.config_service_v2 import pagers +from google.cloud.logging_v2.services.config_service_v2 import transports from google.cloud.logging_v2.types import logging_config -from google.longrunning import operations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account +import google.api_core.operation_async as operation_async # type: ignore +import google.auth +import google.protobuf.empty_pb2 as empty_pb2 # type: ignore +import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore +import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_logging_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_logging_service_v2.py index a2763274363c..544b2fd557db 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_logging_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_logging_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,6 +34,23 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2AsyncClient +from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client +from google.cloud.logging_v2.services.logging_service_v2 import pagers +from google.cloud.logging_v2.services.logging_service_v2 import transports +from google.cloud.logging_v2.types import log_entry +from google.cloud.logging_v2.types import logging +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api.monitored_resource_pb2 as monitored_resource_pb2 # type: ignore import google.auth import google.logging.type.http_request_pb2 as http_request_pb2 # type: ignore @@ -42,26 +59,8 @@ import google.protobuf.duration_pb2 as duration_pb2 # type: ignore import google.protobuf.struct_pb2 as struct_pb2 # type: ignore import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.logging_service_v2 import ( - LoggingServiceV2AsyncClient, - LoggingServiceV2Client, - pagers, - transports, -) -from google.cloud.logging_v2.types import log_entry, logging -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_metrics_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_metrics_service_v2.py index 51d143b43335..9f598dbcf95e 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_metrics_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging/tests/unit/gapic/logging_v2/test_metrics_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,6 +34,22 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.logging_v2.services.metrics_service_v2 import MetricsServiceV2AsyncClient +from google.cloud.logging_v2.services.metrics_service_v2 import MetricsServiceV2Client +from google.cloud.logging_v2.services.metrics_service_v2 import pagers +from google.cloud.logging_v2.services.metrics_service_v2 import transports +from google.cloud.logging_v2.types import logging_metrics +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api.distribution_pb2 as distribution_pb2 # type: ignore import google.api.label_pb2 as label_pb2 # type: ignore import google.api.launch_stage_pb2 as launch_stage_pb2 # type: ignore @@ -41,26 +57,8 @@ import google.auth import google.protobuf.duration_pb2 as duration_pb2 # type: ignore import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.metrics_service_v2 import ( - MetricsServiceV2AsyncClient, - MetricsServiceV2Client, - pagers, - transports, -) -from google.cloud.logging_v2.types import logging_metrics -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging_internal/noxfile.py b/packages/gapic-generator/tests/integration/goldens/logging_internal/noxfile.py index ea66871f9e45..448aec3ef2b0 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging_internal/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/logging_internal/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_config_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_config_service_v2.py index 96a54034b854..9152aa72c6a1 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_config_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_config_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,34 +34,32 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False -import google.api_core.operation_async as operation_async # type: ignore -import google.auth -import google.protobuf.empty_pb2 as empty_pb2 # type: ignore -import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore -import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) +from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.config_service_v2 import ( - BaseConfigServiceV2AsyncClient, - BaseConfigServiceV2Client, - pagers, - transports, -) +from google.cloud.logging_v2.services.config_service_v2 import BaseConfigServiceV2AsyncClient +from google.cloud.logging_v2.services.config_service_v2 import BaseConfigServiceV2Client +from google.cloud.logging_v2.services.config_service_v2 import pagers +from google.cloud.logging_v2.services.config_service_v2 import transports from google.cloud.logging_v2.types import logging_config -from google.longrunning import operations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account +import google.api_core.operation_async as operation_async # type: ignore +import google.auth +import google.protobuf.empty_pb2 as empty_pb2 # type: ignore +import google.protobuf.field_mask_pb2 as field_mask_pb2 # type: ignore +import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_logging_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_logging_service_v2.py index a2763274363c..544b2fd557db 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_logging_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_logging_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,6 +34,23 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2AsyncClient +from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client +from google.cloud.logging_v2.services.logging_service_v2 import pagers +from google.cloud.logging_v2.services.logging_service_v2 import transports +from google.cloud.logging_v2.types import log_entry +from google.cloud.logging_v2.types import logging +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api.monitored_resource_pb2 as monitored_resource_pb2 # type: ignore import google.auth import google.logging.type.http_request_pb2 as http_request_pb2 # type: ignore @@ -42,26 +59,8 @@ import google.protobuf.duration_pb2 as duration_pb2 # type: ignore import google.protobuf.struct_pb2 as struct_pb2 # type: ignore import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.logging_service_v2 import ( - LoggingServiceV2AsyncClient, - LoggingServiceV2Client, - pagers, - transports, -) -from google.cloud.logging_v2.types import log_entry, logging -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_metrics_service_v2.py b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_metrics_service_v2.py index 36b2af7dc0dc..0fc62ce795c0 100755 --- a/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_metrics_service_v2.py +++ b/packages/gapic-generator/tests/integration/goldens/logging_internal/tests/unit/gapic/logging_v2/test_metrics_service_v2.py @@ -13,20 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule +from proto.marshal.rules import wrappers try: from google.auth.aio import credentials as ga_credentials_async @@ -34,6 +34,22 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.logging_v2.services.metrics_service_v2 import BaseMetricsServiceV2AsyncClient +from google.cloud.logging_v2.services.metrics_service_v2 import BaseMetricsServiceV2Client +from google.cloud.logging_v2.services.metrics_service_v2 import pagers +from google.cloud.logging_v2.services.metrics_service_v2 import transports +from google.cloud.logging_v2.types import logging_metrics +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api.distribution_pb2 as distribution_pb2 # type: ignore import google.api.label_pb2 as label_pb2 # type: ignore import google.api.launch_stage_pb2 as launch_stage_pb2 # type: ignore @@ -41,26 +57,8 @@ import google.auth import google.protobuf.duration_pb2 as duration_pb2 # type: ignore import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -from google.api_core import ( - client_options, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.logging_v2.services.metrics_service_v2 import ( - BaseMetricsServiceV2AsyncClient, - BaseMetricsServiceV2Client, - pagers, - transports, -) -from google.cloud.logging_v2.types import logging_metrics -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/redis/noxfile.py b/packages/gapic-generator/tests/integration/goldens/redis/noxfile.py index ab84e007153d..d860093c9653 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/redis/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/redis/tests/unit/gapic/redis_v1/test_cloud_redis.py b/packages/gapic-generator/tests/integration/goldens/redis/tests/unit/gapic/redis_v1/test_cloud_redis.py index af594270faba..7d14d15849c9 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis/tests/unit/gapic/redis_v1/test_cloud_redis.py +++ b/packages/gapic-generator/tests/integration/goldens/redis/tests/unit/gapic/redis_v1/test_cloud_redis.py @@ -13,32 +13,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule - +from proto.marshal.rules import wrappers try: import aiohttp # type: ignore - from google.api_core.operations_v1 import AsyncOperationsRestClient from google.auth.aio.transport.sessions import AsyncAuthorizedSession + from google.api_core.operations_v1 import AsyncOperationsRestClient HAS_ASYNC_REST_EXTRA = True except ImportError: # pragma: NO COVER HAS_ASYNC_REST_EXTRA = False -from google.protobuf import json_format -from requests import PreparedRequest, Request, Response +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -46,6 +47,26 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.location import locations_pb2 +from google.cloud.redis_v1.services.cloud_redis import CloudRedisAsyncClient +from google.cloud.redis_v1.services.cloud_redis import CloudRedisClient +from google.cloud.redis_v1.services.cloud_redis import pagers +from google.cloud.redis_v1.services.cloud_redis import transports +from google.cloud.redis_v1.types import cloud_redis +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api_core.operation_async as operation_async # type: ignore import google.auth import google.protobuf.duration_pb2 as duration_pb2 # type: ignore @@ -54,30 +75,8 @@ import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore import google.type.dayofweek_pb2 as dayofweek_pb2 # type: ignore import google.type.timeofday_pb2 as timeofday_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.location import locations_pb2 -from google.cloud.redis_v1.services.cloud_redis import ( - CloudRedisAsyncClient, - CloudRedisClient, - pagers, - transports, -) -from google.cloud.redis_v1.types import cloud_redis -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/noxfile.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/noxfile.py index ab84e007153d..d860093c9653 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/tests/unit/gapic/redis_v1/test_cloud_redis.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/tests/unit/gapic/redis_v1/test_cloud_redis.py index a92110a6d9f3..1076ace45ba6 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/tests/unit/gapic/redis_v1/test_cloud_redis.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/tests/unit/gapic/redis_v1/test_cloud_redis.py @@ -13,32 +13,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence +import asyncio from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule - +from proto.marshal.rules import wrappers try: import aiohttp # type: ignore - from google.api_core.operations_v1 import AsyncOperationsRestClient from google.auth.aio.transport.sessions import AsyncAuthorizedSession + from google.api_core.operations_v1 import AsyncOperationsRestClient HAS_ASYNC_REST_EXTRA = True except ImportError: # pragma: NO COVER HAS_ASYNC_REST_EXTRA = False -from google.protobuf import json_format -from requests import PreparedRequest, Request, Response +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -46,6 +47,26 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False +from google.api_core import client_options +from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.location import locations_pb2 +from google.cloud.redis_v1.services.cloud_redis import CloudRedisAsyncClient +from google.cloud.redis_v1.services.cloud_redis import CloudRedisClient +from google.cloud.redis_v1.services.cloud_redis import pagers +from google.cloud.redis_v1.services.cloud_redis import transports +from google.cloud.redis_v1.types import cloud_redis +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account import google.api_core.operation_async as operation_async # type: ignore import google.auth import google.protobuf.duration_pb2 as duration_pb2 # type: ignore @@ -54,30 +75,8 @@ import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore import google.type.dayofweek_pb2 as dayofweek_pb2 # type: ignore import google.type.timeofday_pb2 as timeofday_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) -from google.api_core import exceptions as core_exceptions -from google.api_core import retry as retries -from google.auth import credentials as ga_credentials -from google.auth.exceptions import MutualTLSChannelError -from google.cloud.location import locations_pb2 -from google.cloud.redis_v1.services.cloud_redis import ( - CloudRedisAsyncClient, - CloudRedisClient, - pagers, - transports, -) -from google.cloud.redis_v1.types import cloud_redis -from google.longrunning import operations_pb2 # type: ignore -from google.oauth2 import service_account + + CRED_INFO_JSON = { "credential_source": "/path/to/file", diff --git a/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/noxfile.py b/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/noxfile.py index f27f4da6e576..9afec5aeae68 100755 --- a/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/noxfile.py +++ b/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/noxfile.py @@ -17,8 +17,9 @@ import pathlib import re import shutil -import warnings + from typing import Dict, List +import warnings import nox diff --git a/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/tests/unit/gapic/storagebatchoperations_v1/test_storage_batch_operations.py b/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/tests/unit/gapic/storagebatchoperations_v1/test_storage_batch_operations.py index 9e404765cd38..66367a27be49 100755 --- a/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/tests/unit/gapic/storagebatchoperations_v1/test_storage_batch_operations.py +++ b/packages/gapic-generator/tests/integration/goldens/storagebatchoperations/tests/unit/gapic/storagebatchoperations_v1/test_storage_batch_operations.py @@ -13,24 +13,27 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio -import json -import math import os +import asyncio import re -from collections.abc import AsyncIterable, Iterable, Mapping, Sequence from unittest import mock from unittest.mock import AsyncMock import grpc +from grpc.experimental import aio +from collections.abc import Iterable, AsyncIterable +from google.protobuf import json_format +import json +import math import pytest +from collections.abc import Sequence, Mapping from google.api_core import api_core_version -from google.protobuf import json_format -from grpc.experimental import aio -from proto.marshal.rules import wrappers from proto.marshal.rules.dates import DurationRule, TimestampRule -from requests import PreparedRequest, Request, Response +from proto.marshal.rules import wrappers +from requests import Response +from requests import Request, PreparedRequest from requests.sessions import Session +from google.protobuf import json_format try: from google.auth.aio import credentials as ga_credentials_async @@ -38,37 +41,33 @@ except ImportError: # pragma: NO COVER HAS_GOOGLE_AUTH_AIO = False -import google.api_core.operation_async as operation_async # type: ignore -import google.auth -import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore -import google.rpc.code_pb2 as code_pb2 # type: ignore -from google.api_core import ( - client_options, - future, - gapic_v1, - grpc_helpers, - grpc_helpers_async, - operation, - operations_v1, - path_template, -) +from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import future +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers +from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operations_v1 +from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials from google.auth.exceptions import MutualTLSChannelError from google.cloud.location import locations_pb2 -from google.cloud.storagebatchoperations_v1.services.storage_batch_operations import ( - StorageBatchOperationsAsyncClient, - StorageBatchOperationsClient, - pagers, - transports, -) -from google.cloud.storagebatchoperations_v1.types import ( - storage_batch_operations, - storage_batch_operations_types, -) -from google.longrunning import operations_pb2 # type: ignore +from google.cloud.storagebatchoperations_v1.services.storage_batch_operations import StorageBatchOperationsAsyncClient +from google.cloud.storagebatchoperations_v1.services.storage_batch_operations import StorageBatchOperationsClient +from google.cloud.storagebatchoperations_v1.services.storage_batch_operations import pagers +from google.cloud.storagebatchoperations_v1.services.storage_batch_operations import transports +from google.cloud.storagebatchoperations_v1.types import storage_batch_operations +from google.cloud.storagebatchoperations_v1.types import storage_batch_operations_types +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account +import google.api_core.operation_async as operation_async # type: ignore +import google.auth +import google.protobuf.timestamp_pb2 as timestamp_pb2 # type: ignore +import google.rpc.code_pb2 as code_pb2 # type: ignore + + CRED_INFO_JSON = { "credential_source": "/path/to/file", From bd01ac937104965c15592902a0fc028e2b5cff83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Wed, 27 May 2026 21:55:52 +0000 Subject: [PATCH 08/11] format --- packages/bigframes/bigframes/series.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index e3a8b7e49539..72486c1611c6 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -375,7 +375,9 @@ def sql(self) -> str: def transpose(self) -> Series: return self - def _set_internal_query_job(self, query_job: Optional[google.cloud.bigquery.job.QueryJob]): + def _set_internal_query_job( + self, query_job: Optional[google.cloud.bigquery.job.QueryJob] + ): self._query_job = query_job def __len__(self): From a847fc26bbd01e6d271328605108422915d50365 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Thu, 28 May 2026 14:38:16 +0000 Subject: [PATCH 09/11] refactor path constants --- .../scripts/generate_bigframes_bigquery.py | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index 2a1699cdd471..e6437737bd47 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -29,14 +29,20 @@ import jinja2 import yaml +SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.absolute() +PACKAGE_ROOT = SCRIPTS_DIRECTORY.parent +CODE_ROOT = PACKAGE_ROOT / "bigframes" +SCRIPT_PATH_RELATIVE = pathlib.Path(__file__).relative_to(PACKAGE_ROOT) + # Directory containing the YAML files -DATA_DIR = pathlib.Path("scripts/data/sql-functions") +DATA_DIR = SCRIPTS_DIRECTORY / "data" / "sql-functions" # Directory where the generated Python files will be placed -OUTPUT_DIR = pathlib.Path("bigframes/operations/googlesql") +OUTPUT_DIR = CODE_ROOT / "operations" / "googlesql" # Directory where the generated test files will be placed -TEST_OUTPUT_DIR = pathlib.Path("tests/unit/bigquery/generated") +TEST_OUTPUT_DIR = PACKAGE_ROOT / "tests" / "unit" / "bigquery" / "generated" # Directory containing the Jinja2 templates -TEMPLATE_DIR = pathlib.Path("scripts/templates") +TEMPLATE_DIR = SCRIPTS_DIRECTORY / "templates" + RUFF_COMMON_ARGS = [ "--target-version=py310", @@ -415,9 +421,10 @@ def process_yaml_file(yaml_file, templates): # Render and write output_file.parent.mkdir(parents=True, exist_ok=True) ensure_init_py(output_file.parent, OUTPUT_DIR.parent, templates["license"]) + yaml_file_relative = yaml_file.relative_to(PACKAGE_ROOT) content = templates["operation"].render( - yaml_path=str(yaml_file), - script_path="scripts/generate_bigframes_bigquery.py", + yaml_path=yaml_file_relative, + script_path=SCRIPT_PATH_RELATIVE, ops=ops_list, functions=functions_list, ) @@ -438,8 +445,8 @@ def process_yaml_file(yaml_file, templates): test_output_file.parent, TEST_OUTPUT_DIR.parent, templates["license"] ) test_content = templates["test_operation"].render( - yaml_path=str(yaml_file), - script_path="scripts/generate_bigframes_bigquery.py", + yaml_path=yaml_file_relative, + script_path=SCRIPT_PATH_RELATIVE, import_path=import_path, short_name=module_path.name, is_global=is_global, @@ -545,13 +552,13 @@ def generate_series_accessors(functions: list[dict], templates: dict): ) # Render and write core - core_output_file = pathlib.Path("bigframes/extensions/core/series_accessor.py") + core_output_file = CODE_ROOT / "extensions" / "core" / "series_accessor.py" core_output_file.parent.mkdir(parents=True, exist_ok=True) ensure_init_py( - core_output_file.parent, pathlib.Path("bigframes"), templates["license"] + core_output_file.parent, CODE_ROOT, templates["license"] ) core_content = templates["core_series_accessor"].render( - script_path="scripts/generate_bigframes_bigquery.py", + script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs, ) with open(core_output_file, "w") as f: @@ -560,13 +567,13 @@ def generate_series_accessors(functions: list[dict], templates: dict): print(f" Generated {core_output_file}") # Render and write bigframes - bf_output_file = pathlib.Path("bigframes/extensions/bigframes/series_accessor.py") + bf_output_file = CODE_ROOT / "extensions" / "bigframes" / "series_accessor.py" bf_output_file.parent.mkdir(parents=True, exist_ok=True) ensure_init_py( - bf_output_file.parent, pathlib.Path("bigframes"), templates["license"] + bf_output_file.parent, CODE_ROOT, templates["license"] ) bf_content = templates["bigframes_series_accessor"].render( - script_path="scripts/generate_bigframes_bigquery.py", + script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs, ) with open(bf_output_file, "w") as f: @@ -578,10 +585,10 @@ def generate_series_accessors(functions: list[dict], templates: dict): pd_output_file = pathlib.Path("bigframes/extensions/pandas/series_accessor.py") pd_output_file.parent.mkdir(parents=True, exist_ok=True) ensure_init_py( - pd_output_file.parent, pathlib.Path("bigframes"), templates["license"] + pd_output_file.parent, CODE_ROOT, templates["license"] ) pd_content = templates["pandas_series_accessor"].render( - script_path="scripts/generate_bigframes_bigquery.py", + script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs, ) with open(pd_output_file, "w") as f: From 70bab726e54dd25e5359f015bb820dbd9b7f7d4d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Thu, 28 May 2026 14:48:48 +0000 Subject: [PATCH 10/11] one more constant usage --- packages/bigframes/scripts/generate_bigframes_bigquery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index e6437737bd47..39476304d8e9 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -582,7 +582,7 @@ def generate_series_accessors(functions: list[dict], templates: dict): print(f" Generated {bf_output_file}") # Render and write pandas - pd_output_file = pathlib.Path("bigframes/extensions/pandas/series_accessor.py") + pd_output_file = CODE_ROOT / "extensions" / "pandas" / "series_accessor.py" pd_output_file.parent.mkdir(parents=True, exist_ok=True) ensure_init_py( pd_output_file.parent, CODE_ROOT, templates["license"] From 26f4894082425143ee7029505a90625915d595ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Thu, 28 May 2026 14:52:07 +0000 Subject: [PATCH 11/11] ruff format --- .../bigframes/scripts/generate_bigframes_bigquery.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/packages/bigframes/scripts/generate_bigframes_bigquery.py b/packages/bigframes/scripts/generate_bigframes_bigquery.py index 39476304d8e9..124604354205 100755 --- a/packages/bigframes/scripts/generate_bigframes_bigquery.py +++ b/packages/bigframes/scripts/generate_bigframes_bigquery.py @@ -554,9 +554,7 @@ def generate_series_accessors(functions: list[dict], templates: dict): # Render and write core core_output_file = CODE_ROOT / "extensions" / "core" / "series_accessor.py" core_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py( - core_output_file.parent, CODE_ROOT, templates["license"] - ) + ensure_init_py(core_output_file.parent, CODE_ROOT, templates["license"]) core_content = templates["core_series_accessor"].render( script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs, @@ -569,9 +567,7 @@ def generate_series_accessors(functions: list[dict], templates: dict): # Render and write bigframes bf_output_file = CODE_ROOT / "extensions" / "bigframes" / "series_accessor.py" bf_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py( - bf_output_file.parent, CODE_ROOT, templates["license"] - ) + ensure_init_py(bf_output_file.parent, CODE_ROOT, templates["license"]) bf_content = templates["bigframes_series_accessor"].render( script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs, @@ -584,9 +580,7 @@ def generate_series_accessors(functions: list[dict], templates: dict): # Render and write pandas pd_output_file = CODE_ROOT / "extensions" / "pandas" / "series_accessor.py" pd_output_file.parent.mkdir(parents=True, exist_ok=True) - ensure_init_py( - pd_output_file.parent, CODE_ROOT, templates["license"] - ) + ensure_init_py(pd_output_file.parent, CODE_ROOT, templates["license"]) pd_content = templates["pandas_series_accessor"].render( script_path=SCRIPT_PATH_RELATIVE, namespaces=ns_defs,