_boot.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  1. # Copyright (c) 2025-2026 Buf Technologies, Inc.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import annotations
  15. from typing import TYPE_CHECKING, Any, TypeVar, cast
  16. from protobuf._descriptors import (
  17. DescFieldValueList,
  18. DescFieldValueMap,
  19. DescFieldValueMessage,
  20. DescMessage,
  21. )
  22. from protobuf._file_registry import create_file_registry
  23. from protobuf._message import Message as BaseMessage
  24. from protobuf._native_message import object_setattr
  25. try:
  26. from protobuf_ext import generic_setattr
  27. # Workaround Python <3.13 prevents calling object.__setattr__ on objects with
  28. # a native base class that implements __setattr__.
  29. object_setattr = generic_setattr
  30. except ImportError:
  31. object_setattr = object.__setattr__
  32. if TYPE_CHECKING:
  33. from collections.abc import Mapping
  34. from protobuf._descriptors import DescFile
  35. from protobuf._enum import Enum
  36. from protobuf._extension import Extension
  37. from protobuf.wkt._gen.descriptor_pb import FileDescriptorProto
  38. FieldNamesT = TypeVar("FieldNamesT", bound=str)
  39. def boot(
  40. proto: FileDescriptorProto,
  41. stubs: Mapping[str, type[BaseMessage | Enum] | Extension],
  42. ) -> DescFile:
  43. """Bootstrap google/protobuf/descriptor.proto.
  44. The file descriptor for google/protobuf/descriptor.proto cannot be
  45. embedded in serialized form, because it is required to parse itself.
  46. Instead, the code generator emits descriptor.proto messages as Python
  47. constructor calls (keyword arguments with literal values). This module
  48. provides the primitives that make those calls work before descriptors
  49. exist:
  50. - boot() creates a DescFile from an already-instantiated
  51. FileDescriptorProto, bypassing binary deserialization.
  52. - BootMessage is a base class for descriptor.proto messages that
  53. stores fields as plain attributes during bootstrap, then delegates
  54. to the normal Message once descriptors are attached.
  55. - Unset is a sentinel that distinguishes "field absent with a default
  56. value" from "field explicitly set" in bootstrapped messages.
  57. This module is used exclusively for google/protobuf/descriptor.proto.
  58. All other generated files use file_desc(), which deserializes from
  59. bytes.
  60. """
  61. reg = create_file_registry(proto, _err_resolve, stubs=stubs)
  62. res = reg.file(proto.name)
  63. assert res is not None # noqa: S101 # We just added it.
  64. _fill_presence(proto)
  65. return res
  66. def _fill_presence(msg: Message) -> None:
  67. """Fills presence tracking for bootstrapped descriptor messages.
  68. Bootstrap construction writes field values straight into storage,
  69. bypassing ``Message.__setattr__``, so presence for explicit-presence
  70. fields is never recorded, which causes marshaling to fail. We go through
  71. the message after descriptor initialization to fill presence.
  72. We can fill presence by simply traversing FileDescriptorProto for
  73. descriptor.proto since all bootstrapped protos are children of it.
  74. """
  75. if (boot_unset := msg._boot_unset) is None:
  76. return
  77. for field in type(msg)._desc.fields:
  78. key = field.local_name
  79. if key in boot_unset:
  80. continue
  81. value = getattr(msg, key)
  82. match field.value:
  83. case DescFieldValueMessage():
  84. if value is not None:
  85. _fill_presence(value)
  86. case DescFieldValueList(element=DescMessage()):
  87. for item in value:
  88. _fill_presence(item)
  89. case DescFieldValueMap(value=DescMessage()):
  90. for item in value.values():
  91. _fill_presence(item)
  92. if field._requires_presence:
  93. msg._set_field_number_present(field.number)
  94. object_setattr(msg, "_boot_unset", None)
  95. class Message(BaseMessage[FieldNamesT]):
  96. """A thin wrapper on message to handle bootstrapping."""
  97. __slots__ = ("_boot_unset",)
  98. if TYPE_CHECKING:
  99. _boot_unset: set[str] | None
  100. def __init__(self, **kwargs: Any) -> None:
  101. # If the descriptor is set at init time, we treat it like
  102. # a regular message.
  103. if hasattr(self.__class__, "_desc"):
  104. object_setattr(self, "_boot_unset", None)
  105. super().__init__(**kwargs)
  106. return
  107. # Bootstrap
  108. boot_unset = set()
  109. object_setattr(self, "_boot_unset", boot_unset)
  110. for key, value in kwargs.items():
  111. if value is None or isinstance(value, _Unset):
  112. boot_unset.add(key)
  113. object_setattr(
  114. self, key, value.default if isinstance(value, _Unset) else value
  115. ) # Bypasses Message's __setattr__ which requires a descriptor
  116. def has_field(self, key: FieldNamesT) -> bool:
  117. if (boot_unset := getattr(self, "_boot_unset", None)) is None:
  118. return super().has_field(key)
  119. # Bootstrap
  120. # descriptor.proto doesn't use oneofs, so we can safely ignore them here.
  121. return key not in boot_unset
  122. def unset(default: object) -> Any:
  123. """Used to track unset fields with default values for bootstrapped messages."""
  124. return cast("Any", _Unset(default))
  125. class _Unset:
  126. def __init__(self, default: object) -> None:
  127. self.default = default
  128. def _err_resolve(name: str) -> DescFile:
  129. msg = f"unexpected dependency for google/protobuf/descriptor.proto: {name}"
  130. raise NotImplementedError(msg)