From 945fcc4a5d9ac239e36dcdaaa9143e583b7b9667 Mon Sep 17 00:00:00 2001 From: Aidan Jensen Date: Mon, 24 Nov 2025 11:36:07 -0800 Subject: [PATCH] Add option to dedot imports Signed-off-by: Aidan Jensen --- .github/workflows/main.yml | 17 ++++++++++++++++- README.md | 20 ++++++++++++++++++++ mypy_protobuf/main.py | 19 ++++++++++++++++--- run_test.sh | 30 +++++++++++++++++++----------- 4 files changed, 71 insertions(+), 15 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 81c746cd..10094e5b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -9,7 +9,7 @@ on: jobs: run_test: - name: mypy-protobuf - ${{ matrix.py-ver-mypy-protobuf }} - unittest - ${{ matrix.py-ver-unit-tests }} - protobuf ${{ matrix.protobuf-version }} + name: mypy-protobuf - ${{ matrix.py-ver-mypy-protobuf }} - unittest - ${{ matrix.py-ver-unit-tests }} - protobuf ${{ matrix.protobuf-version }} - dedot imports ${{ matrix.dedot-imports }} runs-on: ubuntu-24.04 strategy: matrix: @@ -32,6 +32,20 @@ jobs: protobuf-version: - 6.32.1 - 6.33.1 + dedot-imports: + - 0 + include: + # Dedot imports has a smaller matrix, just validate that things still work + - py-ver-mypy-protobuf: 3.14.0 + py-ver-unit-tests: 3.14.0 + protobuf-version: 6.33.1 + dedot-imports: 1 + - py-ver-mypy-protobuf: 3.14.0 + py-ver-unit-tests: 3.14.0 + protobuf-version: 6.32.1 + dedot-imports: 1 + + steps: - uses: actions/checkout@v4 @@ -52,6 +66,7 @@ jobs: PY_VER_MYPY: ${{matrix.py-ver-mypy-protobuf}} PYTHON_PROTOBUF_VERSION: ${{matrix.protobuf-version}} VALIDATE: 1 + DEDOT_IMPORTS: ${{matrix.dedot-imports}} run: | ./run_test.sh diff --git a/README.md b/README.md index 446f303a..436ba471 100644 --- a/README.md +++ b/README.md @@ -347,6 +347,26 @@ By default mypy-protobuf will output servicer stubs with abstract methods. To ou protoc --python_out=output/location --mypy_grpc_out=generate_concrete_servicer_stubs:output/location ``` +### `dedot_imports` + +This option attempts to match the standard generated python files import pattern by replacing the `.` separated imports. + +``` +protoc --python_out=output/location --mypy_grpc_out=dedot_imports:output/location +``` + +``` +# before: +import testproto.inner.inner_pb2 +def a_inner(self) -> testproto.inner.inner_pb2.Inner: ... + +# after: +import testproto.inner.inner_pb2 as testproto_dot_inner_dot_inner__pb2 +def a_inner(self) -> testproto_dot_inner_dot_inner__pb2.Inner: ... +``` + +*NOTE*: This is currently not compatible with `readable_stubs` + ### Output suppression To suppress output, you can run diff --git a/mypy_protobuf/main.py b/mypy_protobuf/main.py index 487f30f5..e33b6893 100644 --- a/mypy_protobuf/main.py +++ b/mypy_protobuf/main.py @@ -152,6 +152,7 @@ def __init__( relax_strict_optional_primitives: bool, use_default_deprecation_warnings: bool, generate_concrete_servicer_stubs: bool, + dedot_imports: bool, grpc: bool, ) -> None: self.fd = fd @@ -160,6 +161,7 @@ def __init__( self.relax_strict_optional_primitives = relax_strict_optional_primitives self.use_default_depreaction_warnings = use_default_deprecation_warnings self.generate_concrete_servicer_stubs = generate_concrete_servicer_stubs + self.dedot_imports = dedot_imports self.grpc = grpc self.lines: List[str] = [] self.indent = "" @@ -175,7 +177,7 @@ def __init__( # Comments self.source_code_info_by_scl = {tuple(location.path): location for location in fd.source_code_info.location} - def _import(self, path: str, name: str) -> str: + def _import(self, path: str, name: str, dedot: bool = False) -> str: """Imports a stdlib path and returns a handle to it eg. self._import("typing", "Literal") -> "Literal" """ @@ -196,8 +198,13 @@ def _import(self, path: str, name: str) -> str: self.from_imports[imp].add((name, None)) return name else: + dedoted: str | None = None + if dedot: + # Mangle name like protoc does, replace '.' with '_dot_' and '_' with '__' + dedoted = imp.replace("_", "__").replace(".", "_dot_") + imp = f"{imp} as {dedoted}" self.imports.add(imp) - return imp + "." + name + return (dedoted or imp) + "." + name def _import_message(self, name: str) -> str: """Import a referenced message and return a handle""" @@ -226,7 +233,7 @@ def _import_message(self, name: str) -> str: # Not in file. Must import # Python generated code ignores proto packages, so the only relevant factor is # whether it is in the file or not. - import_name = self._import(message_fd.name[:-6].replace("-", "_") + "_pb2", split[0]) + import_name = self._import(message_fd.name[:-6].replace("-", "_") + "_pb2", split[0], dedot=self.dedot_imports) remains = ".".join(split[1:]) if not remains: @@ -1137,6 +1144,7 @@ def generate_mypy_stubs( relax_strict_optional_primitives: bool, use_default_deprecation_warnings: bool, generate_concrete_servicer_stubs: bool, + dedot_imports: bool, ) -> None: for name, fd in descriptors.to_generate.items(): pkg_writer = PkgWriter( @@ -1146,6 +1154,7 @@ def generate_mypy_stubs( relax_strict_optional_primitives, use_default_deprecation_warnings, generate_concrete_servicer_stubs, + dedot_imports=dedot_imports, grpc=False, ) @@ -1171,6 +1180,7 @@ def generate_mypy_grpc_stubs( relax_strict_optional_primitives: bool, use_default_deprecation_warnings: bool, generate_concrete_servicer_stubs: bool, + dedot_imports: bool, ) -> None: for name, fd in descriptors.to_generate.items(): pkg_writer = PkgWriter( @@ -1180,6 +1190,7 @@ def generate_mypy_grpc_stubs( relax_strict_optional_primitives, use_default_deprecation_warnings, generate_concrete_servicer_stubs, + dedot_imports=dedot_imports, grpc=True, ) pkg_writer.write_grpc_async_hacks() @@ -1237,6 +1248,7 @@ def main() -> None: "relax_strict_optional_primitives" in request.parameter, "use_default_deprecation_warnings" in request.parameter, "generate_concrete_servicer_stubs" in request.parameter, + "dedot_imports" in request.parameter, ) @@ -1251,6 +1263,7 @@ def grpc() -> None: "relax_strict_optional_primitives" in request.parameter, "use_default_deprecation_warnings" in request.parameter, "generate_concrete_servicer_stubs" in request.parameter, + "dedot_imports" in request.parameter, ) diff --git a/run_test.sh b/run_test.sh index f6426ac6..78c24fbc 100755 --- a/run_test.sh +++ b/run_test.sh @@ -8,6 +8,7 @@ PY_VER_MYPY_PROTOBUF_SHORT=$(echo "$PY_VER_MYPY_PROTOBUF" | cut -d. -f1-2) PY_VER_MYPY=${PY_VER_MYPY:=3.12.12} PY_VER_UNIT_TESTS="${PY_VER_UNIT_TESTS:=3.9.17 3.10.12 3.11.4 3.12.12 3.13.9 3.14.0}" PYTHON_PROTOBUF_VERSION=${PYTHON_PROTOBUF_VERSION:=6.32.1} +DEDOT_IMPORTS=${DEDOT_IMPORTS:=0} # Confirm UV installed if ! command -v uv &> /dev/null; then @@ -125,9 +126,12 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=relax_strict_optional_primitives:test/generated find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=use_default_deprecation_warnings:test/generated find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=generate_concrete_servicer_stubs:test/generated - # Overwrite w/ run with mypy-protobuf without flags - find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=test/generated - + # Overwrite w/ run with mypy-protobuf without flags unless DEDOT_IMPORTS is set + if [[ "$DEDOT_IMPORTS" -eq 1 ]]; then + find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=dedot_imports:test/generated + else + find proto -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_out=test/generated + fi # Generate grpc protos find proto/testproto/grpc -name "*.proto" -print0 | xargs -0 "$PROTOC" "${PROTOC_ARGS[@]}" --mypy_grpc_out=test/generated @@ -136,13 +140,17 @@ MYPY_PROTOBUF_VENV=venv_$PY_VER_MYPY_PROTOBUF 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 - echo -e "${RED}Some .pyi files did not match. Please commit those files${NC}" - exit 1 + # Don't validate if DEDOT_IMPORTS is set, as it changes the generated files + if [[ $VALIDATE == 1 ]] && [[ $DEDOT_IMPORTS == 0 ]]; then + if ! diff <(echo "$SHA_BEFORE") <(find test/generated -name "*.pyi" -print0 | xargs -0 sha1sum); then + echo -e "${RED}Some .pyi files did not match. Please commit those files${NC}" + exit 1 + fi 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 @@ -202,7 +210,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 +222,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