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
 
@@ -372,7 +373,7 @@ If you need stubs for grpc internal code we suggest using this package https://g
 
 ### `_ClearFieldNamesType` and `_HasFieldNamesType` 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
@@ -381,6 +382,14 @@ Where applicable, type aliases are generated for the arguments to `ClearField` a
       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)

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, members 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