Relativize imports for files in generated set only

pyproject.toml

@@ -40,6 +40,7 @@ dev = [
   "upstream-protobuf",
 
   "buf-bin==1.71.0",
+  "googleapis-googleapis-bufbuild-py==0.1.1.1.20260414192239+c17df5b2beca",
   "poethepoet==0.46.0",
   "protoc-runner==35.1",
   "pyright==1.1.410",
@@ -441,6 +442,11 @@ exclude = [
 reinstall-package = ["protobuf-py-ext"]
 no-build-isolation-package = ["protobuf-py-ext"]
 
+[[tool.uv.index]]
+name = "buf"
+url = "https://buf.build/gen/python"
+explicit = true
+
 [tool.uv.build-backend]
 module-name = ["protobuf"]
 
@@ -454,6 +460,8 @@ license-header = { workspace = true }
 protobuf-py-ext = { workspace = true }
 protobuf-py-bench = { workspace = true }
 
+googleapis-googleapis-bufbuild-py = { index = "buf" }
+
 [tool.uv.workspace]
 members = [
   "packages/protobuf-py-ext",

src/protobuf/plugin/_file.py

@@ -334,12 +334,15 @@ def _to_el(self, v: object) -> str | Ident:
         return ident
 
     def _relativize(self, ident: Ident) -> Ident:
+        type_only = ident.type_only or self._type_checking
+        if (
+            ident._desc is not None
+            and _desc_file(ident._desc).name not in self._file_to_generate
+        ):
+            return _absolute_ident(ident, type_only=type_only)
+
         if not _is_relative(ident.module):
-            return Ident(
-                ident.name,
-                ident.module,
-                type_only=ident.type_only or self._type_checking,
-            )
+            return Ident(ident.name, ident.module, type_only=type_only)
 
         self_segments = _module_segments(self.module)
         import_segments = _module_segments(ident.module)
@@ -352,9 +355,7 @@ def _relativize(self, ident: Ident) -> Ident:
             shared = _shared_prefix_len(package_segments, import_segments)
             leading_dots = len(package_segments) - shared + 1
             path = "." * leading_dots + ".".join(import_segments[shared:])
-        return Module(path).ident(
-            ident.name, type_only=ident.type_only or self._type_checking
-        )
+        return Module(path).ident(ident.name, type_only=type_only)
 
 
 def write(file: _File, path: str, *, no_fmt_off: bool = False) -> str:
@@ -465,6 +466,19 @@ def _desc_ident(
     return ident
 
 
+def _desc_file(desc: DescEnum | DescMessage | DescExtension | DescFile) -> DescFile:
+    return desc if isinstance(desc, DescFile) else desc.file
+
+
+def _absolute_ident(ident: Ident, *, type_only: bool) -> Ident:
+    module_path = ident.module.path.removeprefix(".")
+    if module_path == "" and isinstance(ident._desc, DescFile):
+        module_path = ident.name
+    return Ident(
+        ident.name, Module(module_path), type_only=type_only, _desc=ident._desc
+    )
+
+
 def _use_wkt_module(desc: DescFile, file_to_generate: frozenset[str]) -> bool:
     """Return True if the descriptor should be imported from protobuf.wkt."""
     # Well-known types are imported from protobuf.wkt unless the
@@ -483,8 +497,24 @@ def _write_imports(
     # Ruff splits the imports into groups of std, global, and relative. We also do the same:
     for group in _group_and_sort_imports(imports):
         for module, idents in group.items():
-            deduped = sorted({aliases.resolve_import(ident) for ident in idents})
-            lines.append(f"{indent}from {module.path} import {', '.join(deduped)}")
+            module_imports = [
+                ident for ident in idents if _is_module_import(module, ident)
+            ]
+            lines.extend(
+                f"{indent}import {import_}"
+                for import_ in sorted(
+                    {aliases.resolve_import(ident) for ident in module_imports}
+                )
+            )
+
+            from_imports = [
+                ident for ident in idents if not _is_module_import(module, ident)
+            ]
+            if from_imports:
+                deduped = sorted(
+                    {aliases.resolve_import(ident) for ident in from_imports}
+                )
+                lines.append(f"{indent}from {module.path} import {', '.join(deduped)}")
 
         lines.append("")
 
@@ -597,6 +627,11 @@ def _is_relative(module: Module) -> bool:
     return module.path.startswith(".")
 
 
+def _is_module_import(module: Module, ident: Ident) -> bool:
+    """Return True if the identifier imports a top-level module itself."""
+    return isinstance(ident._desc, DescFile) and module.path == ident.name
+
+
 def _module_segments(module: Module) -> list[str]:
     path = module.path.removeprefix(".")
     return path.split(".") if path else []

tests/buf.lock

@@ -0,0 +1,6 @@
+# Generated by buf. DO NOT EDIT.
+version: v2
+deps:
+  - name: buf.build/googleapis/googleapis
+    commit: c17df5b2beca46928cc87d5656bd5343
+    digest: b5:648a01e0170d4512dea7d564016165decd1ed6e34bef79fe54753e51ad7e27545709ad9157d7551270147d551155c595a2fb0bf5bb33b1c83040ddbce915c604

tests/buf.yaml

@@ -1,3 +1,5 @@
 version: v2
 modules:
   - path: proto_buf
+deps:
+  - buf.build/googleapis/googleapis

tests/gen_buf/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025-2026 Buf Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations

tests/gen_buf/local_dep/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025-2026 Buf Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations

tests/gen_buf/local_import/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025-2026 Buf Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations