Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 3d714e2

Browse files
gkevinzhenggkevinzheng
authored
feat: Proto-plus modifications for enforcing strict oneofs (#1126)
* feat: Proto-plus modifications for enforcing strict oneofs * Added template directory + changed unit tests to pytest * Finished README * linting * Added source of truth comment
1 parent 67c98f0 commit 3d714e2

7 files changed

Lines changed: 481 additions & 2 deletions

File tree

‎gapic_templates/README.rst‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
GAPIC Templates
2+
===============
3+
4+
This directory is intended for inserting handwritten files
5+
into autogenerated directories (see `owlbot.py`). In particular,
6+
it is needed for inserting the definition of `OneofMessage` into
7+
`google/cloud/bigtable/admin_v2/types` to prevent circular import
8+
issues with having something in that directory import from
9+
`google/cloud/bigtable/admin_v2/overlay`.
10+
11+
12+
Usage
13+
-----
14+
15+
The contents of this directory will be copied in to `google/cloud/bigtable`.
16+
As such, create subdirectories in this directory to mirror the final location
17+
under `google/cloud/bigtable` that your file will be under.

‎gapic_templates/admin_v2/types/oneof_message.py.j2‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
#
17+
# AUTOGENERATED FILE: DO NOT EDIT. The source of truth for this file is under
18+
# the gapic_templates directory, not this one.
19+
20+
import collections.abc
21+
import proto
22+
23+
24+
class OneofMessage(proto.Message):
25+
def _get_oneof_field_from_key(self, key):
26+
"""Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
27+
28+
oneof_type = None
29+
30+
try:
31+
oneof_type = self._meta.fields[key].oneof
32+
except KeyError:
33+
# Underscores may be appended to field names
34+
# that collide with python or proto-plus keywords.
35+
# In case a key only exists with a `_` suffix, coerce the key
36+
# to include the `_` suffix. It's not possible to
37+
# natively define the same field with a trailing underscore in protobuf.
38+
# See related issue
39+
# https://github.com/googleapis/python-api-core/issues/227
40+
if f"{key}_" in self._meta.fields:
41+
key = f"{key}_"
42+
oneof_type = self._meta.fields[key].oneof
43+
44+
return oneof_type
45+
46+
def __init__(
47+
self,
48+
mapping=None,
49+
*,
50+
ignore_unknown_fields=False,
51+
**kwargs,
52+
):
53+
# We accept several things for `mapping`:
54+
# * An instance of this class.
55+
# * An instance of the underlying protobuf descriptor class.
56+
# * A dict
57+
# * Nothing (keyword arguments only).
58+
#
59+
#
60+
# Check for oneofs collisions in the parameters provided. Extract a set of
61+
# all fields that are set from the mappings + kwargs combined.
62+
mapping_fields = set(kwargs.keys())
63+
64+
if mapping is None:
65+
pass
66+
elif isinstance(mapping, collections.abc.Mapping):
67+
mapping_fields.update(mapping.keys())
68+
elif isinstance(mapping, self._meta.pb):
69+
mapping_fields.update(field.name for field, _ in mapping.ListFields())
70+
elif isinstance(mapping, type(self)):
71+
mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
72+
else:
73+
# Sanity check: Did we get something not a map? Error if so.
74+
raise TypeError(
75+
"Invalid constructor input for %s: %r"
76+
% (
77+
self.__class__.__name__,
78+
mapping,
79+
)
80+
)
81+
82+
oneofs = set()
83+
84+
for field in mapping_fields:
85+
oneof_field = self._get_oneof_field_from_key(field)
86+
if oneof_field is not None:
87+
if oneof_field in oneofs:
88+
raise ValueError(
89+
"Invalid constructor input for %s: Multiple fields defined for oneof %s"
90+
% (self.__class__.__name__, oneof_field)
91+
)
92+
else:
93+
oneofs.add(oneof_field)
94+
95+
super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
96+
97+
def __setattr__(self, key, value):
98+
# Oneof check: Only set the value of an existing oneof field
99+
# if the field being overridden is the same as the field already set
100+
# for the oneof.
101+
oneof = self._get_oneof_field_from_key(key)
102+
if (
103+
oneof is not None
104+
and self._pb.HasField(oneof)
105+
and self._pb.WhichOneof(oneof) != key
106+
):
107+
raise ValueError(
108+
"Overriding the field set for oneof %s with a different field %s"
109+
% (oneof, key)
110+
)
111+
super().__setattr__(key, value)

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

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
#
17+
# AUTOGENERATED FILE: DO NOT EDIT. The source of truth for this file is under
18+
# the gapic_templates directory, not this one.
19+
20+
import collections.abc
21+
import proto
22+
23+
24+
class OneofMessage(proto.Message):
25+
def _get_oneof_field_from_key(self, key):
26+
"""Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
27+
28+
oneof_type = None
29+
30+
try:
31+
oneof_type = self._meta.fields[key].oneof
32+
except KeyError:
33+
# Underscores may be appended to field names
34+
# that collide with python or proto-plus keywords.
35+
# In case a key only exists with a `_` suffix, coerce the key
36+
# to include the `_` suffix. It's not possible to
37+
# natively define the same field with a trailing underscore in protobuf.
38+
# See related issue
39+
# https://github.com/googleapis/python-api-core/issues/227
40+
if f"{key}_" in self._meta.fields:
41+
key = f"{key}_"
42+
oneof_type = self._meta.fields[key].oneof
43+
44+
return oneof_type
45+
46+
def __init__(
47+
self,
48+
mapping=None,
49+
*,
50+
ignore_unknown_fields=False,
51+
**kwargs,
52+
):
53+
# We accept several things for `mapping`:
54+
# * An instance of this class.
55+
# * An instance of the underlying protobuf descriptor class.
56+
# * A dict
57+
# * Nothing (keyword arguments only).
58+
#
59+
#
60+
# Check for oneofs collisions in the parameters provided. Extract a set of
61+
# all fields that are set from the mappings + kwargs combined.
62+
mapping_fields = set(kwargs.keys())
63+
64+
if mapping is None:
65+
pass
66+
elif isinstance(mapping, collections.abc.Mapping):
67+
mapping_fields.update(mapping.keys())
68+
elif isinstance(mapping, self._meta.pb):
69+
mapping_fields.update(field.name for field, _ in mapping.ListFields())
70+
elif isinstance(mapping, type(self)):
71+
mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
72+
else:
73+
# Sanity check: Did we get something not a map? Error if so.
74+
raise TypeError(
75+
"Invalid constructor input for %s: %r"
76+
% (
77+
self.__class__.__name__,
78+
mapping,
79+
)
80+
)
81+
82+
oneofs = set()
83+
84+
for field in mapping_fields:
85+
oneof_field = self._get_oneof_field_from_key(field)
86+
if oneof_field is not None:
87+
if oneof_field in oneofs:
88+
raise ValueError(
89+
"Invalid constructor input for %s: Multiple fields defined for oneof %s"
90+
% (self.__class__.__name__, oneof_field)
91+
)
92+
else:
93+
oneofs.add(oneof_field)
94+
95+
super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
96+
97+
def __setattr__(self, key, value):
98+
# Oneof check: Only set the value of an existing oneof field
99+
# if the field being overridden is the same as the field already set
100+
# for the oneof.
101+
oneof = self._get_oneof_field_from_key(key)
102+
if (
103+
oneof is not None
104+
and self._pb.HasField(oneof)
105+
and self._pb.WhichOneof(oneof) != key
106+
):
107+
raise ValueError(
108+
"Overriding the field set for oneof %s with a different field %s"
109+
% (oneof, key)
110+
)
111+
super().__setattr__(key, value)

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import proto # type: ignore
2121

2222
from google.cloud.bigtable.admin_v2.types import types
23+
from google.cloud.bigtable.admin_v2.types import oneof_message
2324
from google.protobuf import duration_pb2 # type: ignore
2425
from google.protobuf import timestamp_pb2 # type: ignore
2526
from google.rpc import status_pb2 # type: ignore
@@ -584,7 +585,7 @@ class ColumnFamily(proto.Message):
584585
)
585586

586587

587-
class GcRule(proto.Message):
588+
class GcRule(oneof_message.OneofMessage):
588589
r"""Rule for determining which cells to delete during garbage
589590
collection.
590591

‎owlbot.py‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,9 @@
2020
from typing import List, Optional
2121

2222
import synthtool as s
23-
from synthtool import gcp
23+
from synthtool import gcp, _tracked_paths
2424
from synthtool.languages import python
25+
from synthtool.sources import templates
2526

2627
common = gcp.CommonTemplates()
2728

@@ -238,4 +239,33 @@ def add_overlay_to_init_py(init_py_location, import_statements):
238239
"""
239240
)
240241

242+
# Oneofs work:
243+
244+
# Move the definition of a oneof message into the autogenerated types
245+
# directory. This is needed to prevent circular import issues in admin_v2/overlay.
246+
# This can also be used to insert other files into other autogenerated directories.
247+
gapic_templates = templates.TemplateGroup("gapic_templates")
248+
_tracked_paths.add(gapic_templates.dir)
249+
gapic_templated_files = gapic_templates.render()
250+
s.move(
251+
gapic_templated_files,
252+
destination="google/cloud/bigtable",
253+
excludes=["README.rst"],
254+
)
255+
256+
# Add the oneof_message import into table.py for GcRule
257+
s.replace(
258+
"google/cloud/bigtable/admin_v2/types/table.py",
259+
r"^(from google\.cloud\.bigtable\.admin_v2\.types import .+)$",
260+
r"""\1
261+
from google.cloud.bigtable.admin_v2.types import oneof_message""",
262+
)
263+
264+
# Re-subclass GcRule in table.py
265+
s.replace(
266+
"google/cloud/bigtable/admin_v2/types/table.py",
267+
r"class GcRule\(proto\.Message\)\:",
268+
"class GcRule(oneof_message.OneofMessage):",
269+
)
270+
241271
s.shell.run(["nox", "-s", "blacken"], hide_output=False)

‎tests/unit/admin_overlay/my_oneof_message.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
import proto
17+
18+
from google.cloud.bigtable.admin_v2.types import oneof_message
19+
20+
__protobuf__ = proto.module(
21+
package="test.oneof.v1",
22+
manifest={
23+
"MyOneofMessage",
24+
},
25+
)
26+
27+
28+
# Foo and Bar belong to oneof foobar, and baz is independent.
29+
class MyOneofMessage(oneof_message.OneofMessage):
30+
foo: int = proto.Field(
31+
proto.INT32,
32+
number=1,
33+
oneof="foobar",
34+
)
35+
36+
bar: int = proto.Field(
37+
proto.INT32,
38+
number=2,
39+
oneof="foobar",
40+
)
41+
42+
baz: int = proto.Field(
43+
proto.INT32,
44+
number=3,
45+
)

0 commit comments

Comments
 (0)