initial edition 2024 support (#689)

Signed-off-by: Aidan Jensen <aidan.jensen@robust.ai>
Co-authored-by: Aidan Jensen <aidan.jensen@robust.ai>

Author: aidandj

CHANGELOG.md

@@ -5,6 +5,7 @@
 - Drop testing support for protobuf <6.32 because they don't support editions
   - With some more work this could be added back in a testing refactor
   - Protobuf <6.32 still had the edition enums and field options, so it *should* still work. But is untested
+- Add support for editions (up to 2024)
 
 ## 3.7.0
 

mypy_protobuf/main.py

@@ -540,8 +540,9 @@ def write_messages(
                         # See https://github.com/nipunn1313/mypy-protobuf/issues/71
                         wl("*,")
                     for field in constructor_fields:
+                        implicit_presence = field.options.features.field_presence == d.FeatureSet.FieldPresence.IMPLICIT
                         field_type = self.python_type(field, generic_container=True)
-                        if self.fd.syntax == "proto3" and is_scalar(field) and field.label != d.FieldDescriptorProto.LABEL_REPEATED and not self.relax_strict_optional_primitives and not field.proto3_optional:
+                        if (implicit_presence and self.fd.syntax == "editions") or (self.fd.syntax == "proto3" and is_scalar(field) and field.label != d.FieldDescriptorProto.LABEL_REPEATED and not self.relax_strict_optional_primitives and not field.proto3_optional):
                             wl(f"{field.name}: {field_type} = ...,")
                         else:
                             wl(f"{field.name}: {field_type} | None = ...,")
@@ -561,7 +562,7 @@ def write_stringly_typed_fields(self, desc: d.DescriptorProto) -> None:
         # HasField only supports singular. ClearField supports repeated as well
         # In proto3, HasField only supports message fields and optional fields
         # HasField always supports oneof fields
-        hf_fields = [f.name for f in desc.field if f.HasField("oneof_index") or (f.label != d.FieldDescriptorProto.LABEL_REPEATED and (self.fd.syntax != "proto3" or f.type == d.FieldDescriptorProto.TYPE_MESSAGE or f.proto3_optional))]
+        hf_fields = [f.name for f in desc.field if f.HasField("oneof_index") or (self.fd.syntax == "editions" and f.options.features.field_presence != d.FeatureSet.FieldPresence.IMPLICIT) or (f.label != d.FieldDescriptorProto.LABEL_REPEATED and (self.fd.syntax in ("proto2", "") or f.type == d.FieldDescriptorProto.TYPE_MESSAGE or f.proto3_optional))]
         cf_fields = [f.name for f in desc.field]
         wo_fields = {oneof.name: [f.name for f in desc.field if f.HasField("oneof_index") and f.oneof_index == idx] for idx, oneof in enumerate(desc.oneof_decl)}
 
@@ -1140,8 +1141,6 @@ def generate_mypy_stubs(
         pkg_writer.write_enums(fd.enum_type, "", [d.FileDescriptorProto.ENUM_TYPE_FIELD_NUMBER])
         pkg_writer.write_messages(fd.message_type, "", [d.FileDescriptorProto.MESSAGE_TYPE_FIELD_NUMBER])
         pkg_writer.write_extensions(fd.extension, [d.FileDescriptorProto.EXTENSION_FIELD_NUMBER])
-        if fd.options.py_generic_services:
-            pkg_writer.write_services(fd.service, [d.FileDescriptorProto.SERVICE_FIELD_NUMBER])
 
         assert name == fd.name
         assert fd.name.endswith(".proto")
@@ -1199,6 +1198,10 @@ def code_generation() -> Iterator[Tuple[plugin_pb2.CodeGeneratorRequest, plugin_
 
     # Declare support for optional proto3 fields
     response.supported_features |= plugin_pb2.CodeGeneratorResponse.FEATURE_PROTO3_OPTIONAL
+    response.supported_features |= plugin_pb2.CodeGeneratorResponse.FEATURE_SUPPORTS_EDITIONS
+
+    response.minimum_edition = d.EDITION_LEGACY
+    response.maximum_edition = d.EDITION_2024
 
     yield request, response
 

proto/testproto/edition2024.proto

@@ -0,0 +1,26 @@
+// Edition version of proto2 file
+edition = "2024";
+
+package test;
+
+option features.utf8_validation = NONE;
+option features.enforce_naming_style = STYLE_LEGACY;
+option features.default_symbol_visibility = EXPORT_ALL;
+
+message Editions2024SubMessage {
+  string thing = 1;
+}
+
+message Editions2024Test {
+  // Expect to be always set
+  string legacy = 1 [features.field_presence = LEGACY_REQUIRED];
+  // Expect HasField generated
+  string explicit_singular = 2 [features.field_presence = EXPLICIT];
+  // Expect HasField generated?
+  Editions2024SubMessage message_field = 3 [features.field_presence = EXPLICIT];
+  // Expect implicit field presence, no HasField generated
+  string implicit_singular = 4 [features.field_presence = IMPLICIT];
+  // Not set, should default to EXPLICIT
+  string default_singular = 5;
+
+}

stubtest_allowlist.txt

@@ -103,6 +103,7 @@ testproto.test_pb2.Simple1.no_package
 testproto.test_pb2.Simple1.outer_message_in_oneof
 testproto.test_pb2.Simple1.rep_inner_enum
 testproto.test_pb2.Simple1.rep_inner_message
+testproto.edition2024_pb2.Editions2024Test.message_field
 
 # All messages now fail with something like this:
 # error: testproto.nested.nested_pb2.Nested is inconsistent, metaclass differs
@@ -150,6 +151,8 @@ testproto.test_pb2.Simple1.InnerMessage
 testproto.test_pb2.Simple2
 testproto.test_pb2.DeprecatedMessage
 testproto.test_pb2.DeprecatedMessageBadComment
+testproto.edition2024_pb2.Editions2024SubMessage
+testproto.edition2024_pb2.Editions2024Test
 
 # All messages now fail with something like this:
 # error: testproto.nopackage_pb2.NoPackage2.DESCRIPTOR variable differs from runtime type google._upb._message.Descriptor
@@ -197,6 +200,8 @@ testproto.test_pb2.Simple1.InnerMessage.DESCRIPTOR
 testproto.test_pb2.Simple2.DESCRIPTOR
 testproto.test_pb2.DeprecatedMessage.DESCRIPTOR
 testproto.test_pb2.DeprecatedMessageBadComment.DESCRIPTOR
+testproto.edition2024_pb2.Editions2024SubMessage.DESCRIPTOR
+testproto.edition2024_pb2.Editions2024Test.DESCRIPTOR
 testproto.test_pb2.DESCRIPTOR
 testproto.test_no_generic_services_pb2.DESCRIPTOR
 testproto.test_extensions2_pb2.DESCRIPTOR
@@ -208,6 +213,7 @@ testproto.nested.nested_pb2.DESCRIPTOR
 testproto.inner.inner_pb2.DESCRIPTOR
 testproto.grpc.import_pb2.DESCRIPTOR
 testproto.grpc.dummy_pb2.DESCRIPTOR
+testproto.edition2024_pb2.DESCRIPTOR
 testproto.dot.com.test_pb2.DESCRIPTOR
 testproto.comment_special_chars_pb2.DESCRIPTOR
 testproto.Capitalized.Capitalized_pb2.DESCRIPTOR

test/generated/testproto/edition2024_pb2.pyi

@@ -0,0 +1,68 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+Edition version of proto2 file"""
+
+import builtins
+import google.protobuf.descriptor
+import google.protobuf.message
+import sys
+import typing
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+@typing.final
+class Editions2024SubMessage(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    THING_FIELD_NUMBER: builtins.int
+    thing: builtins.str
+    def __init__(
+        self,
+        *,
+        thing: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["thing", b"thing"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["thing", b"thing"]) -> None: ...
+
+Global___Editions2024SubMessage: typing_extensions.TypeAlias = Editions2024SubMessage
+
+@typing.final
+class Editions2024Test(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    LEGACY_FIELD_NUMBER: builtins.int
+    EXPLICIT_SINGULAR_FIELD_NUMBER: builtins.int
+    MESSAGE_FIELD_FIELD_NUMBER: builtins.int
+    IMPLICIT_SINGULAR_FIELD_NUMBER: builtins.int
+    DEFAULT_SINGULAR_FIELD_NUMBER: builtins.int
+    legacy: builtins.str
+    """Expect to be always set"""
+    explicit_singular: builtins.str
+    """Expect HasField generated"""
+    implicit_singular: builtins.str
+    """Expect implicit field presence, no HasField generated"""
+    default_singular: builtins.str
+    """Not set, should default to EXPLICIT"""
+    @property
+    def message_field(self) -> Global___Editions2024SubMessage:
+        """Expect HasField generated?"""
+
+    def __init__(
+        self,
+        *,
+        legacy: builtins.str | None = ...,
+        explicit_singular: builtins.str | None = ...,
+        message_field: Global___Editions2024SubMessage | None = ...,
+        implicit_singular: builtins.str = ...,
+        default_singular: builtins.str | None = ...,
+    ) -> None: ...
+    def HasField(self, field_name: typing.Literal["default_singular", b"default_singular", "explicit_singular", b"explicit_singular", "legacy", b"legacy", "message_field", b"message_field"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing.Literal["default_singular", b"default_singular", "explicit_singular", b"explicit_singular", "implicit_singular", b"implicit_singular", "legacy", b"legacy", "message_field", b"message_field"]) -> None: ...
+
+Global___Editions2024Test: typing_extensions.TypeAlias = Editions2024Test

test/test_generated_mypy.py

@@ -68,10 +68,10 @@ def _is_summary(line: str) -> bool:
 
 def test_generate_mypy_matches() -> None:
     proto_files = glob.glob("proto/**/*.proto", recursive=True)
-    assert len(proto_files) == 17  # Just a sanity check that all the files show up
+    assert len(proto_files) == 18  # Just a sanity check that all the files show up
 
     pyi_files = glob.glob("test/generated/**/*.pyi", recursive=True)
-    assert len(pyi_files) == 19  # Should be higher - because grpc files generate extra pyis
+    assert len(pyi_files) == 20  # Should be higher - because grpc files generate extra pyis
 
     failure_check_results = []
     for fn in proto_files:
@@ -121,7 +121,7 @@ def grab_expectations(filename: str, marker: str) -> Generator[Tuple[str, int],
     assert errors_39 == expected_errors_39
 
     # Some sanity checks to make sure we don't mess this up. Please update as necessary.
-    assert len(errors_39) == 81
+    assert len(errors_39) == 83
 
 
 def test_func() -> None:
@@ -524,3 +524,30 @@ def test_reserved_keywords() -> None:
     prk = PythonReservedKeywords(none=none_instance, valid=PythonReservedKeywords.valid_in_finally)
     assert prk.none.valid == 5
     assert prk.valid == PythonReservedKeywords.valid_in_finally
+
+
+def test_editions_2024() -> None:
+    from testproto.edition2024_pb2 import Editions2024SubMessage, Editions2024Test
+
+    submsg = Editions2024SubMessage(thing="example")
+
+    testmsg = Editions2024Test(
+        legacy="legacy value",
+        explicit_singular="explicit value",
+        message_field=submsg,
+        implicit_singular="implicit value",
+        default_singular="default value",
+    )
+    assert testmsg.legacy == "legacy value"
+    assert testmsg.explicit_singular == "explicit value"
+    assert testmsg.message_field == submsg
+    assert testmsg.implicit_singular == "implicit value"
+    assert testmsg.default_singular == "default value"
+
+    assert testmsg.HasField("explicit_singular")
+    assert testmsg.HasField("message_field")
+    assert testmsg.HasField("legacy")
+    assert testmsg.HasField("default_singular")
+
+    with pytest.raises(ValueError):
+        testmsg.HasField("implicit_singular")  # type: ignore

test_negative/negative.py

@@ -11,6 +11,7 @@
 import grpc
 import grpc.aio
 from testproto.dot.com.test_pb2 import TestMessage
+from testproto.edition2024_pb2 import Editions2024Test
 from testproto.grpc.dummy_pb2 import (  # E:3.8
     DeprecatedRequest,
     DummyReply,
@@ -318,3 +319,18 @@ def DeprecatedMethodNotDeprecatedRequest(
 async_stub2 = DeprecatedServiceAsyncStub(grpc.aio.insecure_channel(""))  # Not deprecating async stub at this time
 
 de = DeprecatedEnum.DEPRECATED_ONE
+
+# Edition 2024 tests
+
+ed = Editions2024Test(
+    legacy=None,
+    explicit_singular=None,
+    message_field=None,
+    default_singular=None,
+    implicit_singular=None,  # E:3.8
+)
+ed.HasField("implicit_singular")  # E:3.8
+ed.HasField("legacy")  # Should be generated
+ed.HasField("explicit_singular")  # Should be generated
+ed.HasField("message_field")  # Should be generated
+ed.HasField("default_singular")  # Should be generated