Add Registry.extensions_for

Author: anuraaga

src/protobuf/_registry.py

@@ -27,6 +27,7 @@
     from ._enum import Enum
 
 _Types = DescMessage | DescEnum | DescExtension | DescService
+_MessageTypeInfo = DescMessage | str | Message
 _T = TypeVar("_T", bound=_Types)
 
 
@@ -176,8 +177,21 @@ def extension(self, type_name: str) -> DescExtension | None:
         msg = self._types.get(type_name)
         return msg if isinstance(msg, DescExtension) else None
 
+    def extensions_for(self, type_info: _MessageTypeInfo) -> list[DescExtension]:
+        """Look up all extensions for a given message type.
+
+        Args:
+            type_info: The extended message, either as a
+                DescMessage, its fully qualified name, or a Message instance.
+
+        Returns:
+            A list of extension descriptors for the given message type.
+        """
+        msg_type_name = _resolve_type_name(type_info)
+        return list(self._extendees.get(msg_type_name, {}).values())
+
     def extension_for(
-        self, type_info: DescMessage | str | Message, number: int
+        self, type_info: _MessageTypeInfo, number: int
     ) -> DescExtension | None:
         """Look up an extension by the message it extends and field number.
 
@@ -189,14 +203,7 @@ def extension_for(
         Returns:
             The descriptor for the extension, or `None` if not found.
         """
-        match type_info:
-            case DescMessage():
-                msg_type_name = type_info.type_name
-            case Message():
-                msg_type_name = type_info.desc().type_name
-            case str():
-                msg_type_name = type_info
-
+        msg_type_name = _resolve_type_name(type_info)
         return self._extendees[msg_type_name].get(number)
 
     def _get_type(self, type_name: str, typ: type[_T]) -> None | _T:
@@ -220,3 +227,13 @@ def __iter__(
         yield from self._types.values()
 
     __slots__ = "_extendees", "_files", "_types"
+
+
+def _resolve_type_name(type_info: _MessageTypeInfo) -> str:
+    match type_info:
+        case DescMessage():
+            return type_info.type_name
+        case Message():
+            return type_info.desc().type_name
+        case str():
+            return type_info

src/protobuf/_to_json.py

@@ -183,10 +183,9 @@ def _message_to_json_value(message: Message, opts: ToJsonOptions) -> JsonValue:
         result[json_key] = json_value
 
     # Extension fields
-    if opts.registry and (uf := message._unknown_fields):
-        for field_number in uf:
-            ext_desc = opts.registry.extension_for(message._desc, field_number)
-            if not ext_desc:
+    if opts.registry:
+        for ext_desc in opts.registry.extensions_for(message):
+            if ext_desc.type not in message:
                 continue
             value = message[ext_desc.type]
             match field_value := ext_desc.value:

tests/test_registry.py

@@ -32,6 +32,7 @@ def test_empty_registry() -> None:
     assert reg.extension("foo") is None
     assert reg.service("foo") is None
     assert reg.extension_for("foo", 1) is None
+    assert reg.extensions_for("foo") == []
 
 
 @pytest.fixture
@@ -100,6 +101,8 @@ def test_extension_without_message(file: DescFile) -> None:
     assert reg.message("P.M") is None
     assert reg.extension_for("P.M", 100) is file.extensions[0]
     assert reg.extension_for(file.messages[0], 100) is file.extensions[0]
+    assert reg.extensions_for("P.M") == [file.extensions[0]]
+    assert reg.extensions_for(file.messages[0]) == [file.extensions[0]]
 
 
 def test_multiple(protoc: Protoc, file: DescFile) -> None:
@@ -130,6 +133,7 @@ def test_multiple(protoc: Protoc, file: DescFile) -> None:
     assert reg.enum("O.E") is other.enums[0]
     assert reg.extension("O.ext") is other.extensions[0]
     assert reg.extension_for("O.M", 100) is other.extensions[0]
+    assert reg.extensions_for("O.M") == [other.extensions[0]]
 
 
 def test_last_win(protoc: Protoc, file: DescFile) -> None:
@@ -183,3 +187,4 @@ def assert_all(file: DescFile, reg: Registry) -> None:
     assert reg.enum("P.E") is file.enums[0]
     assert reg.extension("P.ext") is file.extensions[0]
     assert reg.extension_for("P.M", 100) is file.extensions[0]
+    assert file.extensions[0] in reg.extensions_for("P.M")