Rename *NamesType to *ArgType

Add WhichOneofArgType and WhichOneofReturnType aliases

Signed-off-by: Aidan Jensen <aidandj.github@gmail.com>

Author: aidandj

CHANGELOG.md

@@ -7,7 +7,8 @@
   - 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)
 - Add `generate_concrete_servicer_stubs` option to generate concrete instead of abstract servicer stubs
-- Add `_HasFieldNamesType` and `_ClearFieldNamesType` aliases to allow for typing field manipulation functions
+- Add `_HasFieldArgType` and `_ClearFieldArgType` aliases to allow for typing field manipulation functions
+- Add `_WhichOneofArgType_<field_name>` and `_WhichOneofReturnType_<field_name>` type aliases
 
 ## 3.7.0
 

README.md

@@ -101,6 +101,7 @@ See [Changelog](CHANGELOG.md) for full listing
 * `mypy-protobuf` generates correctly typed constructors dependinding on field presence.
 * `mypy-protobuf` generates correctly typed `HasField`, `WhichOneof`, and `ClearField` methods.
 * There are differences in how `mypy-protobuf` and `pyi_out` generate enums. See [this issue](https://github.com/protocolbuffers/protobuf/issues/8175) for details
+* Type aliases exported for `HasField`, `WhichOneof` and `ClearField` arguments
 
 #### Examples
 
@@ -370,20 +371,28 @@ protoc \
 Note that generated code for grpc will work only together with code for python and locations should be the same.
 If you need stubs for grpc internal code we suggest using this package https://github.com/shabbyrobe/grpc-stubs
 
-### `_ClearFieldNamesType` and `_HasFieldNamesType` aliases
+### `_ClearFieldArgType` and `_HasFieldArgType` aliases
 
-Where applicable, type aliases are generated for the arguments to `ClearField` and `HasField`. These can be used to create typed functions for field manipulation:
+Where applicable, type aliases are generated for the arguments to `ClearField`, `WhichOneof` and `HasField`. These can be used to create typed functions for field manipulation:
 
 ```python
   from testproto.edition2024_pb2 import Editions2024Test
 
-  def test_hasfield_alias(msg: Editions2024Test, field: "Editions2024Test._HasFieldNamesType") -> bool:
+  def test_hasfield_alias(msg: Editions2024Test, field: "Editions2024Test._HasFieldArgType") -> bool:
       return msg.HasField(field)
 
   test_hasfield_alias(Editions2024Test(), "legacy")
+
+  def test_whichoneof_alias(
+      msg: SimpleProto3,
+      oneof: "SimpleProto3._WhichOneofArgType_a_oneof",
+  ) -> "SimpleProto3._WhichOneofReturnType_a_oneof | None":
+      return msg.WhichOneof(oneof)
+
+  test_whichoneof_alias(SimpleProto3(), "a_oneof")
 ```
 
-Note the deferred evaluation (string reference, or `from __future__ import annotations`. This bypasses the fact that the alias does not exist on the stub)
+Note the deferred evaluation (string reference, or `from __future__ import annotations`. This bypasses the fact that the alias does not exist on the runtime class)
 
 ### Targeting python2 support
 

mypy_protobuf/extensions_pb2.pyi

@@ -38,8 +38,8 @@ class FieldOptions(google.protobuf.message.Message):
         keytype: builtins.str = ...,
         valuetype: builtins.str = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___FieldOptions: typing_extensions.TypeAlias = FieldOptions
 

mypy_protobuf/main.py

@@ -578,29 +578,43 @@ def write_stringly_typed_fields(self, desc: d.DescriptorProto) -> None:
             return
 
         if hf_fields:
-            wl("_HasFieldNamesType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), hf_fields_text)
+            wl("_HasFieldArgType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), hf_fields_text)
             wl(
-                "def HasField(self, field_name: _HasFieldNamesType) -> {}: ...",
+                "def HasField(self, field_name: _HasFieldArgType) -> {}: ...",
                 self._builtin("bool"),
             )
         if cf_fields:
-            wl("_ClearFieldNamesType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), cf_fields_text)
+            wl("_ClearFieldArgType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), cf_fields_text)
             wl(
-                "def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...",
+                "def ClearField(self, field_name: _ClearFieldArgType) -> None: ...",
             )
 
+        # Write type aliases first so overloads are not interrupted
         for wo_field, members in sorted(wo_fields.items()):
-            if len(wo_fields) > 1:
-                wl("@{}", self._import("typing", "overload"))
             wl(
-                "def WhichOneof(self, oneof_group: {}[{}]) -> {}[{}] | None: ...",
-                self._import("typing", "Literal"),
-                # Accepts both str and bytes
-                f'"{wo_field}", b"{wo_field}"',
+                "_WhichOneofReturnType_{}: {} = {}[{}]",
+                wo_field,
+                self._import("typing_extensions", "TypeAlias"),
                 self._import("typing", "Literal"),
                 # Returns `str`
                 ", ".join(f'"{m}"' for m in members),
             )
+            wl(
+                "_WhichOneofArgType_{}: {} = {}[{}]",
+                wo_field,
+                self._import("typing_extensions", "TypeAlias"),
+                self._import("typing", "Literal"),
+                # Accepts both str and bytes
+                f'"{wo_field}", b"{wo_field}"',
+            )
+        for wo_field, _ in sorted(wo_fields.items()):
+            if len(wo_fields) > 1:
+                wl("@{}", self._import("typing", "overload"))
+            wl(
+                "def WhichOneof(self, oneof_group: {}) -> {} | None: ...",
+                f"_WhichOneofArgType_{wo_field}",
+                f"_WhichOneofReturnType_{wo_field}",
+            )
 
     def write_extensions(
         self,

stubtest_allowlist.txt

@@ -242,5 +242,7 @@ testproto.test_pb2.Extensions1.ext
 
 
 # Generated type aliases for HasField and ClearField. These do not exist on a message, but are also just type aliases
-.*_HasFieldNamesType
-.*_ClearFieldNamesType
+.*_HasFieldArgType
+.*_ClearFieldArgType
+.*_WhichOneofReturnType.*
+.*_WhichOneofArgType.*

test/generated/google/protobuf/duration_pb2.pyi

@@ -131,7 +131,7 @@ class Duration(google.protobuf.message.Message, google.protobuf.internal.well_kn
         seconds: builtins.int = ...,
         nanos: builtins.int = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["nanos", b"nanos", "seconds", b"seconds"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["nanos", b"nanos", "seconds", b"seconds"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___Duration: typing_extensions.TypeAlias = Duration

test/generated/mypy_protobuf/extensions_pb2.pyi

@@ -38,8 +38,8 @@ class FieldOptions(google.protobuf.message.Message):
         keytype: builtins.str = ...,
         valuetype: builtins.str = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___FieldOptions: typing_extensions.TypeAlias = FieldOptions
 

test/generated/testproto/Capitalized/Capitalized_pb2.pyi

@@ -27,8 +27,8 @@ class lower(google.protobuf.message.Message):
         *,
         a: builtins.int = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["a", b"a"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["a", b"a"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___lower: typing_extensions.TypeAlias = lower
 
@@ -44,10 +44,10 @@ class Upper(google.protobuf.message.Message):
         *,
         Lower: Global___lower | None = ...,
     ) -> None: ...
-    _HasFieldNamesType: typing_extensions.TypeAlias = typing.Literal["Lower", b"Lower"]
-    def HasField(self, field_name: _HasFieldNamesType) -> builtins.bool: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["Lower", b"Lower"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _HasFieldArgType: typing_extensions.TypeAlias = typing.Literal["Lower", b"Lower"]
+    def HasField(self, field_name: _HasFieldArgType) -> builtins.bool: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["Lower", b"Lower"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___Upper: typing_extensions.TypeAlias = Upper
 
@@ -63,9 +63,9 @@ class lower2(google.protobuf.message.Message):
         *,
         upper: Global___Upper | None = ...,
     ) -> None: ...
-    _HasFieldNamesType: typing_extensions.TypeAlias = typing.Literal["upper", b"upper"]
-    def HasField(self, field_name: _HasFieldNamesType) -> builtins.bool: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["upper", b"upper"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _HasFieldArgType: typing_extensions.TypeAlias = typing.Literal["upper", b"upper"]
+    def HasField(self, field_name: _HasFieldArgType) -> builtins.bool: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["upper", b"upper"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___lower2: typing_extensions.TypeAlias = lower2

test/generated/testproto/comment_special_chars_pb2.pyi

@@ -77,7 +77,7 @@ class Test(google.protobuf.message.Message):
         j: builtins.str = ...,
         k: builtins.str = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["a", b"a", "b", b"b", "c", b"c", "d", b"d", "e", b"e", "f", b"f", "g", b"g", "h", b"h", "i", b"i", "j", b"j", "k", b"k"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["a", b"a", "b", b"b", "c", b"c", "d", b"d", "e", b"e", "f", b"f", "g", b"g", "h", b"h", "i", b"i", "j", b"j", "k", b"k"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___Test: typing_extensions.TypeAlias = Test

test/generated/testproto/dot/com/test_pb2.pyi

@@ -27,7 +27,7 @@ class TestMessage(google.protobuf.message.Message):
         *,
         foo: builtins.str = ...,
     ) -> None: ...
-    _ClearFieldNamesType: typing_extensions.TypeAlias = typing.Literal["foo", b"foo"]
-    def ClearField(self, field_name: _ClearFieldNamesType) -> None: ...
+    _ClearFieldArgType: typing_extensions.TypeAlias = typing.Literal["foo", b"foo"]
+    def ClearField(self, field_name: _ClearFieldArgType) -> None: ...
 
 Global___TestMessage: typing_extensions.TypeAlias = TestMessage