Add option to dedot imports

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

Author: aidandj

.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
 

README.md

@@ -347,6 +347,24 @@ 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: ...
+```
+
 ### Output suppression
 
 To suppress output, you can run

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

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