-
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathprotobuf.py
More file actions
451 lines (361 loc) · 17.6 KB
/
Copy pathprotobuf.py
File metadata and controls
451 lines (361 loc) · 17.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
# SPDX-FileCopyrightText: 2019-present The Firebird Projects <www.firebirdsql.org>
#
# SPDX-License-Identifier: MIT
#
# PROGRAM/MODULE: firebird-base
# FILE: firebird/base/protobuf.py
# DESCRIPTION: Registry for Google Protocol Buffer messages and enums
# CREATED: 27.12.2019
#
# The contents of this file are subject to the MIT License
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Copyright (c) 2019 Firebird Project (www.firebirdsql.org)
# All Rights Reserved.
#
# Contributor(s): Pavel Císař (original code)
# ______________________________________.
"""Firebird Base - Registry for Google Protocol Buffer messages and enums
This module provides a central registry for Google Protocol Buffer message types
and enum types generated from `.proto` files. It allows creating message instances
and accessing enum information using their fully qualified names (e.g.,
"my.package.MyMessage", "my.package.MyEnum") without needing to directly import
the corresponding generated `_pb2.py` modules throughout the codebase.
Benefits:
* Decouples code using protobuf messages from the specific generated modules.
* Provides a single point for managing and discovering available message/enum types.
* Facilitates dynamic loading of protobuf definitions via entry points.
Core Features:
* Register message/enum types using their file DESCRIPTOR object.
* Create new message instances by name using `create_message()`.
* Access enum descriptors and values by name using `get_enum_type()`.
* Load protobuf definitions registered by other installed packages via entry points
using `load_registered()`.
* Helpers for common types like `google.protobuf.Struct`.
Example:
# Assume you have my_proto_pb2.py generated from my_proto.proto
# containing:
# message Sample { required string name = 1; }
# enum Status { UNKNOWN = 0; OK = 1; ERROR = 2; }
from firebird.base.protobuf import (
register_descriptor, create_message, get_enum_type, is_msg_registered
)
# Import the generated descriptor (only needed once, e.g., at startup)
try:
from . import my_proto_pb2 # Replace with actual import path
HAS_MY_PROTO = True
except ImportError:
HAS_MY_PROTO = False
# 1. Register the types from the descriptor
if HAS_MY_PROTO:
register_descriptor(my_proto_pb2.DESCRIPTOR)
print(f"Is 'my_proto.Sample' registered? {is_msg_registered('my_proto.Sample')}")
# 2. Create a message instance by name
if HAS_MY_PROTO:
try:
msg = create_message('my_proto.Sample')
msg.name = "Example"
print(f"Created message: {msg}")
# 3. Access enum type and values by name
status_enum = get_enum_type('my_proto.Status')
print(f"Status enum name: {status_enum.name}")
print(f"OK value: {status_enum.OK}") # Access like attribute
print(f"Name for value 2: {status_enum.get_value_name(2)}") # Access via method
print(f"Available status keys: {status_enum.keys()}")
except KeyError as e:
print(f"Error accessing registered proto type: {e}")
"""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from importlib.metadata import entry_points
from typing import Any, cast
from google.protobuf import any_pb2, duration_pb2, empty_pb2, field_mask_pb2, json_format, struct_pb2, timestamp_pb2
from google.protobuf.descriptor import EnumDescriptor
from google.protobuf.message import Message as ProtoMessage
from google.protobuf.struct_pb2 import Struct as StructProto
from .collections import Registry
from .types import Distinct
#: Name of well-known EMPTY protobuf message (for use with `.create_message()`)
PROTO_EMPTY: str = 'google.protobuf.Empty'
#: Name of well-known ANY protobuf message (for use with `.create_message()`)
PROTO_ANY: str = 'google.protobuf.Any'
#: Name of well-known DURATION protobuf message (for use with `.create_message()`)
PROTO_DURATION: str = 'google.protobuf.Duration'
#: Name of well-known TIMESTAMP protobuf message (for use with `.create_message()`)
PROTO_TIMESTAMP: str = 'google.protobuf.Timestamp'
#: Name of well-known STRUCT protobuf message (for use with `.create_message()`)
PROTO_STRUCT: str = 'google.protobuf.Struct'
#: Name of well-known VALUE protobuf message (for use with `.create_message()`)
PROTO_VALUE: str = 'google.protobuf.Value'
#: Name of well-known LISTVALUE protobuf message (for use with `.create_message()`)
PROTO_LISTVALUE: str = 'google.protobuf.ListValue'
#: Name of well-known FIELDMASK protobuf message (for use with `.create_message()`)
PROTO_FIELDMASK: str = 'google.protobuf.FieldMask'
# Classes
@dataclass(eq=True, order=True, frozen=True)
class ProtoMessageType(Distinct):
"""Registry entry representing a registered Protocol Buffer message type.
Stores the fully qualified name and the constructor (the generated class)
for a message type, allowing instantiation via the registry.
Arguments:
name: Fully qualified message type name (e.g., "package.Message").
constructor: The callable (generated message class) used to create instances.
"""
#: Fully qualified message type name (e.g., "package.Message").
name: str
#: The callable (generated message class) used to create instances.
constructor: Callable
def get_key(self) -> str:
"""Returns the message name, used as the key in the registry."""
return self.name
@dataclass(eq=True, order=True, frozen=True)
class ProtoEnumType(Distinct):
"""Registry entry providing access to a registered Protocol Buffer enum type.
Wraps the `EnumDescriptor` and provides an API similar to generated enum
types, allowing access to names and values without direct import of the
generated `_pb2` module.
Arguments:
descriptor: The `google.protobuf.descriptor.EnumDescriptor` for the enum type.
Example::
# Assuming 'my_proto.Status' enum (UNKNOWN=0, OK=1) is registered
status_enum = get_enum_type('my_proto.Status')
print(status_enum.OK) # Output: 1 (Access value by name)
print(status_enum.get_value_name(1)) # Output: 'OK' (Get name by value)
print(status_enum.keys()) # Output: ['UNKNOWN', 'OK']
print(status_enum.values()) # Output: [0, 1]
print(status_enum.items()) # Output: [('UNKNOWN', 0), ('OK', 1)]
try:
print(status_enum.NONEXISTENT)
except AttributeError as e:
print(e) # Output: Enum my_proto.Status has no value with name 'NONEXISTENT'
try:
print(status_enum.get_value_name(99))
except KeyError as e:
print(e) # Output: "Enum my_proto.Status has no name defined for value 99"
"""
#: The `google.protobuf.descriptor.EnumDescriptor` for the enum type.
descriptor: EnumDescriptor
def get_key(self) -> str:
"""Returns the full enum name, used as the key in the registry."""
return self.name
def __getattr__(self, name: str):
"""Return the integer value corresponding to the enum member name `name`.
Arguments:
name: The string name of the enum member.
Returns:
The integer value of the enum member.
Raises:
AttributeError: If `name` is not a valid member name for this enum.
"""
if name in self.descriptor.values_by_name:
return self.descriptor.values_by_name[name].number
raise AttributeError(f"Enum {self.name} has no value with name '{name}'")
def keys(self) -> list[str]:
"""Return a list of the string names in the enum.
These are returned in the order they were defined in the .proto file.
"""
return [value_descriptor.name for value_descriptor in self.descriptor.values]
def values(self) -> list[int]:
"""Return a list of the integer values in the enum.
These are returned in the order they were defined in the .proto file.
"""
return [value_descriptor.number for value_descriptor in self.descriptor.values]
def items(self) -> list[tuple[str, int]]:
"""Return a list of the (name, value) pairs of the enum.
These are returned in the order they were defined in the .proto file.
"""
return [(value_descriptor.name, value_descriptor.number)
for value_descriptor in self.descriptor.values]
def get_value_name(self, number: int) -> str:
"""Return the string name corresponding to the enum member value `number`.
Arguments:
number: The integer value of the enum member.
Returns:
The string name of the enum member.
Raises:
KeyError: If there is no value for specified name.
"""
if number in self.descriptor.values_by_number:
return self.descriptor.values_by_number[number].name
raise KeyError(f"Enum {self.name} has no name defined for value {number}")
@property
def name(self) -> str:
"""The fully qualified name of the enum type (e.g., "package.MyEnum")."""
return self.descriptor.full_name
#: Internal registry storing ProtoMessageType instances.
_msgreg: Registry = Registry()
#: Internal registry storing ProtoEnumType instances.
_enumreg: Registry = Registry()
def struct2dict(struct: StructProto) -> dict[str, Any]:
"""Unpack a `google.protobuf.Struct` message into a Python dictionary.
Uses `google.protobuf.json_format.MessageToDict`.
Arguments:
struct: The `Struct` message instance.
Returns:
A Python dictionary representing the struct's content.
"""
return json_format.MessageToDict(struct)
def dict2struct(value: dict[str, Any]) -> StructProto:
"""Pack a Python dictionary into a `google.protobuf.Struct` message.
Arguments:
value: The Python dictionary.
Returns:
A `Struct` message instance containing the dictionary's data.
"""
struct = StructProto()
struct.update(value)
return struct
def create_message(name: str, serialized: bytes | None=None) -> ProtoMessage:
"""Create a new instance of a registered protobuf message type by name.
Optionally initializes the message by parsing serialized data.
Arguments:
name: Fully qualified name of the registered protobuf message type.
serialized: Optional bytes containing the serialized message data.
Returns:
An instance of the requested protobuf message class.
Raises:
KeyError: If `name` does not correspond to a registered message type.
google.protobuf.message.DecodeError: If `serialized` data is provided
but cannot be parsed correctly for the message type.
"""
if (msg := _msgreg.get(name)) is None:
raise KeyError(f"Unregistered protobuf message '{name}'")
result = cast(ProtoMessageType, msg).constructor()
if serialized is not None:
result.ParseFromString(serialized)
return result
def get_message_factory(name: str) -> Callable:
"""Return the constructor (factory callable) for a registered message type.
Allows creating multiple instances without repeated registry lookups.
Arguments:
name: Fully qualified name of the registered protobuf message type.
Returns:
The callable (message class) used to construct instances.
Raises:
KeyError: If `name` does not correspond to a registered message type.
"""
if (msg := _msgreg.get(name)) is None:
raise KeyError(f"Unregistered protobuf message '{name}'")
return cast(ProtoMessageType, msg).constructor
def is_msg_registered(name: str) -> bool:
"""Check if a protobuf message type with the given name is registered.
Arguments:
name: Fully qualified message type name.
Returns:
True if registered, False otherwise.
"""
return name in _msgreg
def is_enum_registered(name: str) -> bool:
"""Check if a protobuf enum type with the given name is registered.
Arguments:
name: Fully qualified enum type name.
Returns:
True if registered, False otherwise.
"""
return name in _enumreg
def get_enum_type(name: str) -> ProtoEnumType:
"""Return the `ProtoEnumType` wrapper for a registered enum type by name.
Provides access to enum members and values via the wrapper object.
Arguments:
Fully qualified name of the registered protobuf enum type.
Returns:
The `ProtoEnumType` instance for the requested enum.
Raises:
KeyError: If `name` does not correspond to a registered enum type.
"""
if (e := _enumreg.get(name)) is None:
raise KeyError(f"Unregistered protobuf enum type '{name}'")
return e
def get_enum_field_type(msg, field_name: str) -> str:
"""Return the fully qualified name of the enum type for a message field.
Arguments:
msg: An *instance* of a protobuf message.
field_name: The string name of the field within the message.
Returns:
The fully qualified name of the enum type used by the field.
Raises:
KeyError: If `msg` does not have a field named `field_name`.
TypeError: If the specified field is not an enum field.
"""
if (fdesc := msg.DESCRIPTOR.fields_by_name.get(field_name)) is None:
raise KeyError(f"Message does not have field '{field_name}'")
if fdesc.enum_type is None:
raise TypeError(f"Field '{field_name}' in message type '{msg.DESCRIPTOR.full_name}' is not an enum field.")
return fdesc.enum_type.full_name
def get_enum_value_name(enum_type_name: str, value: int) -> str:
"""Return the string name corresponding to a value within a registered enum type.
Convenience function equivalent to `get_enum_type(enum_type_name).get_value_name(value)`.
Arguments:
enum_type_name: Fully qualified name of the registered enum type.
value: The integer value of the enum member.
Returns:
The string name of the enum member.
Raises:
KeyError: If `enum_type_name` is not registered, or if `value` is not
defined within that enum.
"""
return get_enum_type(enum_type_name).get_value_name(value)
def register_decriptor(file_descriptor) -> None:
"""Register all message and enum types defined within a protobuf file descriptor.
This is the primary mechanism for adding types to the registry. The descriptor
object is typically accessed as `DESCRIPTOR` from a generated `_pb2.py` module.
Arguments:
file_descriptor: The `google.protobuf.descriptor.FileDescriptor` object
(e.g., `my_proto_pb2.DESCRIPTOR`).
"""
for msg_desc in file_descriptor.message_types_by_name.values():
if msg_desc.full_name not in _msgreg:
_msgreg.store(ProtoMessageType(msg_desc.full_name, msg_desc._concrete_class))
for enum_desc in file_descriptor.enum_types_by_name.values():
if enum_desc.full_name not in _enumreg:
_enumreg.store(ProtoEnumType(enum_desc))
def load_registered(group: str) -> None: # pragma: no cover
"""Load and register protobuf types defined via package entry points.
Searches for installed packages that register entry points under the specified
`group`. Each entry point should load a `FileDescriptor` object. This allows
packages to automatically make their protobuf types available to the registry
upon installation.
This function is typically called once during application initialization.
Arguments:
group: The name of the entry-point group to scan (e.g., 'firebird.base.protobuf').
Example:
::
# setup.cfg:
[options.entry_points]
firebird.base.protobuf =
firebird.base.lib_a = firebird.base.lib_a_pb2:DESCRIPTOR
firebird.base.lib_b = firebird.base.lib_b_pb2:DESCRIPTOR
firebird.base.lib_c = firebird.base.lib_c_pb2:DESCRIPTOR
# pyproject.toml
[project.entry-points."firebird.base.protobuf"]
"firebird.base.lib_a" = "firebird.base.lib_a_pb2:DESCRIPTOR"
"firebird.base.lib_b" = "firebird.base.lib_b_pb2:DESCRIPTOR"
"firebird.base.lib_c" = "firebird.base.lib_c_pb2:DESCRIPTOR"
Usage::
# In your application's startup code:
load_registered('firebird.base.protobuf')
# Now messages/enums registered via entry points are available
"""
for desc in (entry.load() for entry in entry_points().get(group, [])):
register_decriptor(desc)
for well_known in [any_pb2, struct_pb2, duration_pb2, empty_pb2, timestamp_pb2, field_mask_pb2]:
register_decriptor(well_known.DESCRIPTOR)
del any_pb2, struct_pb2, duration_pb2, empty_pb2, timestamp_pb2, field_mask_pb2