feat: Proto-plus modifications for enforcing strict oneofs

gapic_templates/README.rst

@@ -0,0 +1,17 @@
+GAPIC Templates
+===============
+
+This directory is intended for inserting handwritten files
+into autogenerated directories (see `owlbot.py`). In particular,
+it is needed for inserting the definition of `OneofMessage` into
+`google/cloud/bigtable/admin_v2/types` to prevent circular import
+issues with having something in that directory import from
+`google/cloud/bigtable/admin_v2/overlay`.
+
+
+Usage
+-----
+
+The contents of this directory will be copied in to `google/cloud/bigtable`.
+As such, create subdirectories in this directory to mirror the final location
+under `google/cloud/bigtable` that your file will be under.

gapic_templates/admin_v2/types/oneof_message.py.j2

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 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.
+#
+#
+# AUTOGENERATED FILE: DO NOT EDIT. The source of truth for this file is under
+# the gapic_templates directory, not this one.
+
+import collections.abc
+import proto
+
+
+class OneofMessage(proto.Message):
+    def _get_oneof_field_from_key(self, key):
+        """Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
+
+        oneof_type = None
+
+        try:
+            oneof_type = self._meta.fields[key].oneof
+        except KeyError:
+            # Underscores may be appended to field names
+            # that collide with python or proto-plus keywords.
+            # In case a key only exists with a `_` suffix, coerce the key
+            # to include the `_` suffix. It's not possible to
+            # natively define the same field with a trailing underscore in protobuf.
+            # See related issue
+            # https://github.com/googleapis/python-api-core/issues/227
+            if f"{key}_" in self._meta.fields:
+                key = f"{key}_"
+                oneof_type = self._meta.fields[key].oneof
+
+        return oneof_type
+
+    def __init__(
+        self,
+        mapping=None,
+        *,
+        ignore_unknown_fields=False,
+        **kwargs,
+    ):
+        # We accept several things for `mapping`:
+        #   * An instance of this class.
+        #   * An instance of the underlying protobuf descriptor class.
+        #   * A dict
+        #   * Nothing (keyword arguments only).
+        #
+        #
+        # Check for oneofs collisions in the parameters provided. Extract a set of
+        # all fields that are set from the mappings + kwargs combined.
+        mapping_fields = set(kwargs.keys())
+
+        if mapping is None:
+            pass
+        elif isinstance(mapping, collections.abc.Mapping):
+            mapping_fields.update(mapping.keys())
+        elif isinstance(mapping, self._meta.pb):
+            mapping_fields.update(field.name for field, _ in mapping.ListFields())
+        elif isinstance(mapping, type(self)):
+            mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
+        else:
+            # Sanity check: Did we get something not a map? Error if so.
+            raise TypeError(
+                "Invalid constructor input for %s: %r"
+                % (
+                    self.__class__.__name__,
+                    mapping,
+                )
+            )
+
+        oneofs = set()
+
+        for field in mapping_fields:
+            oneof_field = self._get_oneof_field_from_key(field)
+            if oneof_field is not None:
+                if oneof_field in oneofs:
+                    raise ValueError(
+                        "Invalid constructor input for %s: Multiple fields defined for oneof %s"
+                        % (self.__class__.__name__, oneof_field)
+                    )
+                else:
+                    oneofs.add(oneof_field)
+
+        super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
+
+    def __setattr__(self, key, value):
+        # Oneof check: Only set the value of an existing oneof field
+        # if the field being overridden is the same as the field already set
+        # for the oneof.
+        oneof = self._get_oneof_field_from_key(key)
+        if (
+            oneof is not None
+            and self._pb.HasField(oneof)
+            and self._pb.WhichOneof(oneof) != key
+        ):
+            raise ValueError(
+                "Overriding the field set for oneof %s with a different field %s"
+                % (oneof, key)
+            )
+        super().__setattr__(key, value)

google/cloud/bigtable/admin_v2/types/oneof_message.py

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 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.
+#
+#
+# AUTOGENERATED FILE: DO NOT EDIT. The source of truth for this file is under
+# the gapic_templates directory, not this one.
+
+import collections.abc
+import proto
+
+
+class OneofMessage(proto.Message):
+    def _get_oneof_field_from_key(self, key):
+        """Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
+
+        oneof_type = None
+
+        try:
+            oneof_type = self._meta.fields[key].oneof
+        except KeyError:
+            # Underscores may be appended to field names
+            # that collide with python or proto-plus keywords.
+            # In case a key only exists with a `_` suffix, coerce the key
+            # to include the `_` suffix. It's not possible to
+            # natively define the same field with a trailing underscore in protobuf.
+            # See related issue
+            # https://github.com/googleapis/python-api-core/issues/227
+            if f"{key}_" in self._meta.fields:
+                key = f"{key}_"
+                oneof_type = self._meta.fields[key].oneof
+
+        return oneof_type
+
+    def __init__(
+        self,
+        mapping=None,
+        *,
+        ignore_unknown_fields=False,
+        **kwargs,
+    ):
+        # We accept several things for `mapping`:
+        #   * An instance of this class.
+        #   * An instance of the underlying protobuf descriptor class.
+        #   * A dict
+        #   * Nothing (keyword arguments only).
+        #
+        #
+        # Check for oneofs collisions in the parameters provided. Extract a set of
+        # all fields that are set from the mappings + kwargs combined.
+        mapping_fields = set(kwargs.keys())
+
+        if mapping is None:
+            pass
+        elif isinstance(mapping, collections.abc.Mapping):
+            mapping_fields.update(mapping.keys())
+        elif isinstance(mapping, self._meta.pb):
+            mapping_fields.update(field.name for field, _ in mapping.ListFields())
+        elif isinstance(mapping, type(self)):
+            mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
+        else:
+            # Sanity check: Did we get something not a map? Error if so.
+            raise TypeError(
+                "Invalid constructor input for %s: %r"
+                % (
+                    self.__class__.__name__,
+                    mapping,
+                )
+            )
+
+        oneofs = set()
+
+        for field in mapping_fields:
+            oneof_field = self._get_oneof_field_from_key(field)
+            if oneof_field is not None:
+                if oneof_field in oneofs:
+                    raise ValueError(
+                        "Invalid constructor input for %s: Multiple fields defined for oneof %s"
+                        % (self.__class__.__name__, oneof_field)
+                    )
+                else:
+                    oneofs.add(oneof_field)
+
+        super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
+
+    def __setattr__(self, key, value):
+        # Oneof check: Only set the value of an existing oneof field
+        # if the field being overridden is the same as the field already set
+        # for the oneof.
+        oneof = self._get_oneof_field_from_key(key)
+        if (
+            oneof is not None
+            and self._pb.HasField(oneof)
+            and self._pb.WhichOneof(oneof) != key
+        ):
+            raise ValueError(
+                "Overriding the field set for oneof %s with a different field %s"
+                % (oneof, key)
+            )
+        super().__setattr__(key, value)

google/cloud/bigtable/admin_v2/types/table.py

@@ -20,6 +20,7 @@
 import proto  # type: ignore
 
 from google.cloud.bigtable.admin_v2.types import types
+from google.cloud.bigtable.admin_v2.types import oneof_message
 from google.protobuf import duration_pb2  # type: ignore
 from google.protobuf import timestamp_pb2  # type: ignore
 from google.rpc import status_pb2  # type: ignore
@@ -584,7 +585,7 @@ class ColumnFamily(proto.Message):
     )
 
 
-class GcRule(proto.Message):
+class GcRule(oneof_message.OneofMessage):
     r"""Rule for determining which cells to delete during garbage
     collection.
 

owlbot.py

@@ -20,8 +20,9 @@
 from typing import List, Optional
 
 import synthtool as s
-from synthtool import gcp
+from synthtool import gcp, _tracked_paths
 from synthtool.languages import python
+from synthtool.sources import templates
 
 common = gcp.CommonTemplates()
 
@@ -238,4 +239,33 @@ def add_overlay_to_init_py(init_py_location, import_statements):
 """
 )
 
+# Oneofs work:
+
+# Move the definition of a oneof message into the autogenerated types
+# directory. This is needed to prevent circular import issues in admin_v2/overlay.
+# This can also be used to insert other files into other autogenerated directories.
+gapic_templates = templates.TemplateGroup("gapic_templates")
+_tracked_paths.add(gapic_templates.dir)
+gapic_templated_files = gapic_templates.render()
+s.move(
+    gapic_templated_files,
+    destination="google/cloud/bigtable",
+    excludes=["README.rst"],
+)
+
+# Add the oneof_message import into table.py for GcRule
+s.replace(
+    "google/cloud/bigtable/admin_v2/types/table.py",
+    r"^(from google\.cloud\.bigtable\.admin_v2\.types import .+)$",
+    r"""\1
+from google.cloud.bigtable.admin_v2.types import oneof_message""",
+)
+
+# Re-subclass GcRule in table.py
+s.replace(
+    "google/cloud/bigtable/admin_v2/types/table.py",
+    r"class GcRule\(proto\.Message\)\:",
+    "class GcRule(oneof_message.OneofMessage):",
+)
+
 s.shell.run(["nox", "-s", "blacken"], hide_output=False)

tests/unit/admin_overlay/my_oneof_message.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 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.
+#
+import proto
+
+from google.cloud.bigtable.admin_v2.types import oneof_message
+
+__protobuf__ = proto.module(
+    package="test.oneof.v1",
+    manifest={
+        "MyOneofMessage",
+    },
+)
+
+
+# Foo and Bar belong to oneof foobar, and baz is independent.
+class MyOneofMessage(oneof_message.OneofMessage):
+    foo: int = proto.Field(
+        proto.INT32,
+        number=1,
+        oneof="foobar",
+    )
+
+    bar: int = proto.Field(
+        proto.INT32,
+        number=2,
+        oneof="foobar",
+    )
+
+    baz: int = proto.Field(
+        proto.INT32,
+        number=3,
+    )