Add rewrite_imports option

docs/writing-plugins/options.md

@@ -38,7 +38,7 @@ Supported field types: `str`, `bool`, `int`, `float`, `Literal[...]`, `StrEnum`,
 
 ## Framework-reserved options
 
-The framework reserves certain option names for all plugins, currently: `no_fmt_off` and `escape_module_with_hash`.
+The framework reserves certain option names for all plugins, currently: `no_fmt_off`, `escape_module_with_hash`, and `rewrite_imports`.
 If your `Options` dataclass defines a field with a reserved name, `run()` raises a `ValueError`.
 
 ### no_fmt_off
@@ -66,6 +66,46 @@ plugins:
     opt: escape_module_with_hash
 ```
 
+### rewrite_imports
+
+Rewrites imports of generated modules that match a glob pattern to an absolute package.
+This makes it possible to reference generated code published from a separate package.
+
+The option takes the form `rewrite_imports=<pattern>:<target>` and can be given multiple times; the first matching pattern wins.
+The pattern is a very reduced subset of glob:
+
+- `*` matches zero or more characters except `/`.
+- `**/` matches zero or more path elements, where an element is one or more characters with a trailing `/`.
+
+The pattern is matched against the import path of the module before it is made relative to the file importing it.
+A generated module such as `.google.type.foo_pb` is matched as the file path `./google/type/foo_pb.py`, relative to the generation root.
+On a match, the target package is prepended:
+
+```yaml title="buf.gen.yaml"
+plugins:
+  - local: protoc-gen-hello
+    out: src/gen
+    opt: rewrite_imports=./google/type/**/*_pb.py:mypkg.gen
+```
+
+With this option, `from .google.type.foo_pb import Foo` is generated as `from mypkg.gen.google.type.foo_pb import Foo` instead.
+References to symbols defined in the file being generated are not imports and are never rewritten.
+
+An empty target rewrites matching imports to the canonical import path, the module path derived from the proto file name, relative to the root of `sys.path`.
+For example, when `buf/validate/validate.proto` is provided by an installed package:
+
+```yaml title="buf.gen.yaml"
+plugins:
+  - local: protoc-gen-hello
+    out: src/gen
+    opt: rewrite_imports=./buf/validate/**/*_pb.py:
+```
+
+This generates `from buf.validate import validate_pb` instead of a relative import that points at a location that does not exist in the output directory.
+A rewritten module import that lands at the top level (for example a proto file at the root of the module) is written as a plain `import foo_pb` statement.
+
+Absolute imports (such as `protobuf.wkt`) are matched without a leading `./` and `.py` extension (for example `protobuf/wkt`), and are replaced by the target entirely.
+
 ## Example: Sensitive fields plugin
 
 Here is a plugin that generates a `_sensitive.py` file for each proto file, listing all fields marked with the `sensitive` custom option from [extensions](../extensions.md#extensions-in-custom-options):

pyproject.toml

@@ -51,6 +51,7 @@ dev-required = [
   "upstream-protobuf",
 
   "buf-bin==1.72.0",
+  "googleapis-googleapis-bufbuild-py==0.3.0.2.20260414192239+c17df5b2beca",
   "hypothesis==6.167.1",
   "license-header==0.0.1",
   "poethepoet==0.48.0",
@@ -476,6 +477,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"]
 
@@ -488,6 +494,8 @@ example-plugin = { 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

@@ -22,10 +22,13 @@
 
 from protobuf import DescEnum, DescExtension, DescFile, DescMessage, ScalarType
 from protobuf.plugin._ident import Ident, Module
+from protobuf.plugin._rewrite_imports import rewrite_module_path
 
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterable, Iterator
 
+    from protobuf.plugin._rewrite_imports import RewriteImports
+
 _INDENT = " " * 4
 
 _TYPING = Module("typing")
@@ -220,6 +223,7 @@ def __init__(
         parameter: str,
         *,
         escape_module_with_hash: bool = False,
+        rewrite_imports: RewriteImports = (),
     ) -> None:
         self.path = path
         self.module = module
@@ -228,6 +232,7 @@ def __init__(
         self._plugin_version = plugin_version
         self._parameter = parameter
         self._escape_module_with_hash = escape_module_with_hash
+        self._rewrite_imports = rewrite_imports
         self._indent = 0
         self._type_checking = False
         self._in_doc = False
@@ -329,19 +334,48 @@ def _to_el(self, v: object) -> str | Ident:
                 )
             case _:
                 return repr(v)
-        ident = self._relativize(ident)
+        ident = self._relativize(self._rewrite_import(ident))
         if ident.type_only:
             self._type_imports[ident.module].add(ident)
         else:
             self._runtime_imports[ident.module].add(ident)
         return ident
 
+    def _rewrite_import(self, ident: Ident) -> Ident:
+        if not self._rewrite_imports:
+            return ident
+        # A DescFile ident imports the module itself, so its import
+        # path includes the ident name.
+        is_module_import = isinstance(ident._desc, DescFile)
+        module_path = ident.module.path
+        if is_module_import:
+            sep = "" if module_path.endswith(".") else "."
+            module_path = f"{module_path}{sep}{ident.name}"
+        elif _is_relative(ident.module) and _module_segments(
+            ident.module
+        ) == _module_segments(self.module):
+            # References to symbols in this file are not imports.
+            return ident
+        rewritten = rewrite_module_path(module_path, self._rewrite_imports)
+        if rewritten is None:
+            return ident
+        if is_module_import:
+            parent, _, name = rewritten.rpartition(".")
+            # A module import that lands at the top level is written as
+            # a plain `import X` statement, keyed by the module itself.
+            module = Module(parent) if parent else Module(name)
+            return Ident(name, module, type_only=ident.type_only, _desc=ident._desc)
+        return Ident(
+            ident.name, Module(rewritten), type_only=ident.type_only, _desc=ident._desc
+        )
+
     def _relativize(self, ident: Ident) -> Ident:
         if not _is_relative(ident.module):
             return Ident(
                 ident.name,
                 ident.module,
                 type_only=ident.type_only or self._type_checking,
+                _desc=ident._desc,
             )
 
         self_segments = _module_segments(self.module)
@@ -486,8 +520,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("")
 
@@ -600,6 +650,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 []

src/protobuf/plugin/_rewrite_imports.py

@@ -0,0 +1,64 @@
+# 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
+
+import re
+
+RewriteImports = tuple[tuple[re.Pattern[str], str], ...]
+
+_OPTION_NAME = "rewrite_imports"
+
+
+def compile_rewrite_imports(rewrites: dict[str, str]) -> RewriteImports:
+    return tuple(
+        (_glob_to_regex(pattern), target.rstrip("."))
+        for pattern, target in rewrites.items()
+    )
+
+
+def rewrite_module_path(module_path: str, rewrites: RewriteImports) -> str | None:
+    relative = module_path.startswith(".")
+    segments = module_path.removeprefix(".").split(".")
+    if not all(segments):
+        return None
+    import_path = f"./{'/'.join(segments)}.py" if relative else "/".join(segments)
+    for pattern, target in rewrites:
+        if pattern.fullmatch(import_path):
+            if relative:
+                dotted = ".".join(segments)
+                return f"{target}.{dotted}" if target else dotted
+            if not target:
+                msg = f"option '{_OPTION_NAME}': cannot rewrite absolute import '{module_path}' to an empty target"
+                raise ValueError(msg)
+            return target
+    return None
+
+
+def _glob_to_regex(pattern: str) -> re.Pattern[str]:
+    parts: list[str] = []
+    i = 0
+    while i < len(pattern):
+        char = pattern[i]
+        if char == "*":
+            if pattern[i + 1 : i + 3] == "*/":
+                parts.append(r"([^/]+/)*")
+                i += 3
+                continue
+            parts.append(r"[^/]*")
+        else:
+            parts.append(re.escape(char))
+        i += 1
+    return re.compile("".join(parts))

src/protobuf/plugin/_run.py

@@ -21,6 +21,7 @@
 from protobuf import maximum_supported_edition, minimum_supported_edition
 from protobuf.plugin._file import write
 from protobuf.plugin._options import parse_options
+from protobuf.plugin._rewrite_imports import compile_rewrite_imports
 from protobuf.plugin._schema import _Schema
 from protobuf.wkt import CodeGeneratorRequest, CodeGeneratorResponse
 
@@ -167,6 +168,7 @@ def run(
             name=name,
             version=version,
             escape_module_with_hash=fw_opts.escape_module_with_hash,
+            rewrite_imports=compile_rewrite_imports(fw_opts.rewrite_imports),
         )
 
         generate(schema)
@@ -230,3 +232,4 @@ def _parse_plugin_options(
 class _FrameworkOptions:
     no_fmt_off: bool = False
     escape_module_with_hash: bool = False
+    rewrite_imports: dict[str, str] = dataclasses.field(default_factory=dict)

src/protobuf/plugin/_schema.py

@@ -27,6 +27,8 @@
 
     from protobuf.wkt import CodeGeneratorRequest
 
+    from ._rewrite_imports import RewriteImports
+
 T_co = TypeVar("T_co", covariant=True)
 
 
@@ -96,13 +98,15 @@ def __init__(
         name: str,
         version: str,
         escape_module_with_hash: bool,
+        rewrite_imports: RewriteImports = (),
     ) -> None:
         self._options = options
         self._name = name
         self._version = version
         self._parameter = req.parameter
         self._generated_files: dict[str, _File] = {}
         self._escape_module_with_hash = escape_module_with_hash
+        self._rewrite_imports = rewrite_imports
 
         file_to_generate = frozenset(req.file_to_generate)
         source_by_name = {s.name: s for s in req.source_file_descriptors}
@@ -166,6 +170,7 @@ def generate_file(
             self._version,
             self._parameter,
             escape_module_with_hash=self._escape_module_with_hash,
+            rewrite_imports=self._rewrite_imports,
         )
         self._generated_files[path] = f
         return f

tests/buf.gen.yaml

@@ -3,3 +3,4 @@ clean: true
 plugins:
   - local: protoc-gen-py
     out: gen_buf
+    opt: "rewrite_imports=./google/**/*_pb.py:"

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