Aidan/clear has field type aliases (#696)

* Add `_HasFieldNamesType` and `_ClearFieldNamesType` aliases to allow for typing field manipulation functions

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

* Fix shellcheck errors in run_test.sh

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

* Rename *NamesType to *ArgType
Add WhichOneofArgType and WhichOneofReturnType aliases

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

* Update CHANGELOG.md

* Update README.md

---------

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

Author: aidandj

.gitignore

@@ -8,8 +8,8 @@ __pycache__/
 *.pyc
 /test/generated/**/*.py
 !/test/generated/**/__init__.py
-/test/generated-concrete/**/*.py
-!/test/generated-concrete/**/__init__.py
+/test/generated_concrete/**/*.py
+!/test/generated_concrete/**/__init__.py
 .pytest_cache
 /build/
 /dist/

CHANGELOG.md

@@ -7,6 +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 `_HasFieldArgType` and `_ClearFieldArgType` aliases to allow for typing field manipulation functions
+- Add `_WhichOneofArgType_<oneof_name>` and `_WhichOneofReturnType_<oneof_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,6 +371,29 @@ 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
 
+### `_ClearFieldArgType`, `_WhichOneofArgType_<oneof_name>`, `_WhichOneofReturnType_<oneof_name>` and `_HasFieldArgType` aliases
+
+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._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 runtime class)
+
 ### Targeting python2 support
 
 mypy-protobuf's drops support for targeting python2 with version 3.0. If you still need python2 support -

mypy_protobuf/extensions_pb2.pyi

@@ -38,7 +38,8 @@ class FieldOptions(google.protobuf.message.Message):
         keytype: builtins.str = ...,
         valuetype: builtins.str = ...,
     ) -> None: ...
-    def ClearField(self, field_name: typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> 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,31 +578,43 @@ def write_stringly_typed_fields(self, desc: d.DescriptorProto) -> None:
             return
 
         if hf_fields:
+            wl("_HasFieldArgType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), hf_fields_text)
             wl(
-                "def HasField(self, field_name: {}[{}]) -> {}: ...",
-                self._import("typing", "Literal"),
-                hf_fields_text,
+                "def HasField(self, field_name: _HasFieldArgType) -> {}: ...",
                 self._builtin("bool"),
             )
         if cf_fields:
+            wl("_ClearFieldArgType: {} = {}[{}]", self._import("typing_extensions", "TypeAlias"), self._import("typing", "Literal"), cf_fields_text)
             wl(
-                "def ClearField(self, field_name: {}[{}]) -> None: ...",
-                self._import("typing", "Literal"),
-                cf_fields_text,
+                "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,

pyproject.toml

@@ -39,6 +39,6 @@ exclude = [
 executionEnvironments = [
     # Due to how upb is typed, we need to disable incompatible variable override checks
     { root = "test/generated", extraPaths = ["./"], reportIncompatibleVariableOverride = "none" },
-    { root = "test/generated-concrete", extraPaths = ["./"], reportIncompatibleVariableOverride = "none" },
+    { root = "test/generated_concrete", extraPaths = ["./"], reportIncompatibleVariableOverride = "none" },
     { root = "mypy_protobuf/extensions_pb2.pyi", reportIncompatibleVariableOverride = "none" },
 ]

run_test.sh

@@ -132,8 +132,8 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF
     find proto/testproto/grpc -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_grpc_out=test/generated
 
     # Generate with concrete service stubs for testing
-    find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=generate_concrete_servicer_stubs:test/generated-concrete
-    find proto/testproto/grpc -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_grpc_out=generate_concrete_servicer_stubs:test/generated-concrete
+    find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=generate_concrete_servicer_stubs:test/generated_concrete
+    find proto/testproto/grpc -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_grpc_out=generate_concrete_servicer_stubs:test/generated_concrete
 
 
     if [[ -n $VALIDATE ]] && ! diff <(echo "$SHA_BEFORE") <(find test/generated -name "*.pyi" -print0 | xargs -0 sha1sum); then
@@ -142,7 +142,8 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF
     fi
 )
 
-ERRORS=()
+ERROR_FILE=$(mktemp)
+trap 'rm -f "$ERROR_FILE"' EXIT
 
 for PY_VER in $PY_VER_UNIT_TESTS; do
     UNIT_TESTS_VENV=venv_$PY_VER
@@ -159,7 +160,7 @@ for PY_VER in $PY_VER_UNIT_TESTS; do
         source "$MYPY_VENV"/bin/activate
         # Run concrete mypy
         CONCRETE_MODULES=( -m test.test_concrete )
-        MYPYPATH=$MYPYPATH:test/generated-concrete mypy ${CUSTOM_TYPESHED_DIR_ARG:+"$CUSTOM_TYPESHED_DIR_ARG"} --python-executable="$UNIT_TESTS_VENV"/bin/python --python-version="$PY_VER_MYPY_TARGET" "${CONCRETE_MODULES[@]}"
+        MYPYPATH=$MYPYPATH:test/generated_concrete mypy ${CUSTOM_TYPESHED_DIR_ARG:+"$CUSTOM_TYPESHED_DIR_ARG"} --python-executable="$UNIT_TESTS_VENV"/bin/python --python-version="$PY_VER_MYPY_TARGET" "${CONCRETE_MODULES[@]}"
 
         export MYPYPATH=$MYPYPATH:test/generated
 
@@ -202,7 +203,7 @@ for PY_VER in $PY_VER_UNIT_TESTS; do
             cp "$MYPY_OUTPUT/mypy_output.omit_linenos" "test_negative/output.expected.$PY_VER_MYPY_TARGET.omit_linenos"
 
             # Record error instead of echoing and exiting
-            ERRORS+=("test_negative/output.expected.$PY_VER_MYPY_TARGET didnt match. Copying over for you.")
+            echo "test_negative/output.expected.$PY_VER_MYPY_TARGET didnt match. Copying over for you." >> "$ERROR_FILE"
         fi
     )
 
@@ -214,11 +215,11 @@ for PY_VER in $PY_VER_UNIT_TESTS; do
 done
 
 # Report all errors at the end
-if [ ${#ERRORS[@]} -gt 0 ]; then
+if [ -s "$ERROR_FILE" ]; then
     echo -e "\n${RED}===============================================${NC}"
-    for error in "${ERRORS[@]}"; do
+    while IFS= read -r error; do
         echo -e "${RED}$error${NC}"
-    done
+    done < "$ERROR_FILE"
     echo -e "${RED}Now rerun${NC}"
     exit 1
 fi

stubtest_allowlist.txt

@@ -239,3 +239,10 @@ testproto.test_extensions3_pb2.msg_option
 testproto.test_extensions3_pb2.enum_option
 testproto.test_extensions2_pb2.SeparateFileExtension.ext
 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
+.*_HasFieldArgType
+.*_ClearFieldArgType
+.*_WhichOneofReturnType.*
+.*_WhichOneofArgType.*

test/generated/google/protobuf/duration_pb2.pyi

@@ -131,6 +131,7 @@ class Duration(google.protobuf.message.Message, google.protobuf.internal.well_kn
         seconds: builtins.int = ...,
         nanos: builtins.int = ...,
     ) -> None: ...
-    def ClearField(self, field_name: typing.Literal["nanos", b"nanos", "seconds", b"seconds"]) -> 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,7 +38,8 @@ class FieldOptions(google.protobuf.message.Message):
         keytype: builtins.str = ...,
         valuetype: builtins.str = ...,
     ) -> None: ...
-    def ClearField(self, field_name: typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> 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