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

Commit 7251258

Browse files
gkevinzhenggkevinzheng
committed
feat: Proto-plus modifications for enforcing strict oneofs
1 parent 67c98f0 commit 7251258

6 files changed

Lines changed: 431 additions & 2 deletions

File tree

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

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
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 collections.abc
17+
import proto
18+
19+
20+
class OneofMessage(proto.Message):
21+
def _get_oneof_field_from_key(self, key):
22+
"""Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
23+
24+
oneof_type = None
25+
26+
try:
27+
oneof_type = self._meta.fields[key].oneof
28+
except KeyError:
29+
# Underscores may be appended to field names
30+
# that collide with python or proto-plus keywords.
31+
# In case a key only exists with a `_` suffix, coerce the key
32+
# to include the `_` suffix. It's not possible to
33+
# natively define the same field with a trailing underscore in protobuf.
34+
# See related issue
35+
# https://github.com/googleapis/python-api-core/issues/227
36+
if f"{key}_" in self._meta.fields:
37+
key = f"{key}_"
38+
oneof_type = self._meta.fields[key].oneof
39+
40+
return oneof_type
41+
42+
def __init__(
43+
self,
44+
mapping=None,
45+
*,
46+
ignore_unknown_fields=False,
47+
**kwargs,
48+
):
49+
# We accept several things for `mapping`:
50+
# * An instance of this class.
51+
# * An instance of the underlying protobuf descriptor class.
52+
# * A dict
53+
# * Nothing (keyword arguments only).
54+
#
55+
#
56+
# Check for oneofs collisions in the parameters provided. Extract a set of
57+
# all fields that are set from the mappings + kwargs combined.
58+
mapping_fields = set(kwargs.keys())
59+
60+
if mapping is None:
61+
pass
62+
elif isinstance(mapping, collections.abc.Mapping):
63+
mapping_fields.update(mapping.keys())
64+
elif isinstance(mapping, self._meta.pb):
65+
mapping_fields.update(field.name for field, _ in mapping.ListFields())
66+
elif isinstance(mapping, type(self)):
67+
mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
68+
else:
69+
# Sanity check: Did we get something not a map? Error if so.
70+
raise TypeError(
71+
"Invalid constructor input for %s: %r"
72+
% (
73+
self.__class__.__name__,
74+
mapping,
75+
)
76+
)
77+
78+
oneofs = set()
79+
80+
for field in mapping_fields:
81+
oneof_field = self._get_oneof_field_from_key(field)
82+
if oneof_field is not None:
83+
if oneof_field in oneofs:
84+
raise ValueError(
85+
"Invalid constructor input for %s: Multiple fields defined for oneof %s"
86+
% (self.__class__.__name__, oneof_field)
87+
)
88+
else:
89+
oneofs.add(oneof_field)
90+
91+
super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
92+
93+
def __setattr__(self, key, value):
94+
# Oneof check: Only set the value of an existing oneof field
95+
# if the field being overridden is the same as the field already set
96+
# for the oneof.
97+
oneof = self._get_oneof_field_from_key(key)
98+
if (
99+
oneof is not None
100+
and self._pb.HasField(oneof)
101+
and self._pb.WhichOneof(oneof) != key
102+
):
103+
raise ValueError(
104+
"Overriding the field set for oneof %s with a different field %s"
105+
% (oneof, key)
106+
)
107+
super().__setattr__(key, value)

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

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
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 collections.abc
17+
import proto
18+
19+
20+
class OneofMessage(proto.Message):
21+
def _get_oneof_field_from_key(self, key):
22+
"""Given a field name, return the corresponding oneof associated with it. If it doesn't exist, return None."""
23+
24+
oneof_type = None
25+
26+
try:
27+
oneof_type = self._meta.fields[key].oneof
28+
except KeyError:
29+
# Underscores may be appended to field names
30+
# that collide with python or proto-plus keywords.
31+
# In case a key only exists with a `_` suffix, coerce the key
32+
# to include the `_` suffix. It's not possible to
33+
# natively define the same field with a trailing underscore in protobuf.
34+
# See related issue
35+
# https://github.com/googleapis/python-api-core/issues/227
36+
if f"{key}_" in self._meta.fields:
37+
key = f"{key}_"
38+
oneof_type = self._meta.fields[key].oneof
39+
40+
return oneof_type
41+
42+
def __init__(
43+
self,
44+
mapping=None,
45+
*,
46+
ignore_unknown_fields=False,
47+
**kwargs,
48+
):
49+
# We accept several things for `mapping`:
50+
# * An instance of this class.
51+
# * An instance of the underlying protobuf descriptor class.
52+
# * A dict
53+
# * Nothing (keyword arguments only).
54+
#
55+
#
56+
# Check for oneofs collisions in the parameters provided. Extract a set of
57+
# all fields that are set from the mappings + kwargs combined.
58+
mapping_fields = set(kwargs.keys())
59+
60+
if mapping is None:
61+
pass
62+
elif isinstance(mapping, collections.abc.Mapping):
63+
mapping_fields.update(mapping.keys())
64+
elif isinstance(mapping, self._meta.pb):
65+
mapping_fields.update(field.name for field, _ in mapping.ListFields())
66+
elif isinstance(mapping, type(self)):
67+
mapping_fields.update(field.name for field, _ in mapping._pb.ListFields())
68+
else:
69+
# Sanity check: Did we get something not a map? Error if so.
70+
raise TypeError(
71+
"Invalid constructor input for %s: %r"
72+
% (
73+
self.__class__.__name__,
74+
mapping,
75+
)
76+
)
77+
78+
oneofs = set()
79+
80+
for field in mapping_fields:
81+
oneof_field = self._get_oneof_field_from_key(field)
82+
if oneof_field is not None:
83+
if oneof_field in oneofs:
84+
raise ValueError(
85+
"Invalid constructor input for %s: Multiple fields defined for oneof %s"
86+
% (self.__class__.__name__, oneof_field)
87+
)
88+
else:
89+
oneofs.add(oneof_field)
90+
91+
super().__init__(mapping, ignore_unknown_fields=ignore_unknown_fields, **kwargs)
92+
93+
def __setattr__(self, key, value):
94+
# Oneof check: Only set the value of an existing oneof field
95+
# if the field being overridden is the same as the field already set
96+
# for the oneof.
97+
oneof = self._get_oneof_field_from_key(key)
98+
if (
99+
oneof is not None
100+
and self._pb.HasField(oneof)
101+
and self._pb.WhichOneof(oneof) != key
102+
):
103+
raise ValueError(
104+
"Overriding the field set for oneof %s with a different field %s"
105+
% (oneof, key)
106+
)
107+
super().__setattr__(key, value)

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
import proto # type: ignore
2121

22-
from google.cloud.bigtable.admin_v2.types import types
22+
from google.cloud.bigtable.admin_v2.types import types, oneof_message
2323
from google.protobuf import duration_pb2 # type: ignore
2424
from google.protobuf import timestamp_pb2 # type: ignore
2525
from google.rpc import status_pb2 # type: ignore
@@ -584,7 +584,7 @@ class ColumnFamily(proto.Message):
584584
)
585585

586586

587-
class GcRule(proto.Message):
587+
class GcRule(oneof_message.OneofMessage):
588588
r"""Rule for determining which cells to delete during garbage
589589
collection.
590590

‎owlbot.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,4 +238,18 @@ def add_overlay_to_init_py(init_py_location, import_statements):
238238
"""
239239
)
240240

241+
# Add the oneof_message import into table.py for GcRule
242+
s.replace(
243+
"google/cloud/bigtable/admin_v2/types/table.py",
244+
r"from google\.cloud\.bigtable\.admin_v2\.types import types",
245+
"from google.cloud.bigtable.admin_v2.types import types, oneof_message",
246+
)
247+
248+
# Re-subclass GcRule in table.py
249+
s.replace(
250+
"google/cloud/bigtable/admin_v2/types/table.py",
251+
r"class GcRule\(proto\.Message\)\:",
252+
"class GcRule(oneof_message.OneofMessage):",
253+
)
254+
241255
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)