Replace rewrite_imports with map_imports

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`, `escape_module_with_hash`, and `rewrite_imports`.
+The framework reserves certain option names for all plugins, currently: `no_fmt_off`, `escape_module_with_hash`, and `map_imports`.
 If your `Options` dataclass defines a field with a reserved name, `run()` raises a `ValueError`.
 
 ### no_fmt_off
@@ -66,45 +66,48 @@ plugins:
     opt: escape_module_with_hash
 ```
 
-### rewrite_imports
+### map_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.
+By default, generated code imports dependencies from the local output with a relative import.
+For example, the module generated for `foo/bar.proto` imports a message from `buf/validate/validate.proto` as `from ..buf.validate.validate_pb import Rule`.
 
-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:
+If a dependency is provided by a package instead, use `map_imports` to import it from there.
+The option takes the form `map_imports=<pattern>:<target>` and can be given multiple times.
+The pattern is matched against the path of the Protobuf file, and the first matching pattern wins.
+The target is a Python package that is prepended to the module path derived from the Protobuf file:
 
 ```yaml title="buf.gen.yaml"
 plugins:
   - local: protoc-gen-hello
     out: src/gen
-    opt: rewrite_imports=./google/type/**/*_pb.py:mypkg.gen
+    opt: map_imports=google/type/: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.
+With this option, a message from `google/type/date.proto` is imported as `from mypkg.gen.google.type.date_pb import Date`.
+
+Patterns support a subset of glob:
+
+- `*` matches zero or more characters except `/`.
+- `**` matches zero or more characters, including `/`.
+- `**/` matches zero or more directories.
+- A trailing `/` matches every file in the directory and its subdirectories.
 
-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`.
+An empty target imports from the canonical module path: the module path derived from the Protobuf file, relative to the root of `sys.path`.
+This is the layout of generated SDKs installed as separate packages, such as those from the Buf Python registry.
 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:
+    opt: "map_imports=buf/validate/:"
 ```
 
-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.
+This generates `from buf.validate import validate_pb` and `from buf.validate.validate_pb import Rule`.
+A mapped module at the top level, such as one generated for a Protobuf file at the root, 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.
+Mapping applies to imports derived from descriptors.
+Identifiers constructed directly from a `Module` are not mapped, and well-known types are always imported from `protobuf.wkt`.
 
 ## Example: Sensitive fields plugin
 

src/protobuf/plugin/_file.py

@@ -22,12 +22,12 @@
 
 from protobuf import DescEnum, DescExtension, DescFile, DescMessage, ScalarType
 from protobuf.plugin._ident import Ident, Module
-from protobuf.plugin._rewrite_imports import rewrite_module_path
+from protobuf.plugin._map_imports import map_import_target
 
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterable, Iterator
 
-    from protobuf.plugin._rewrite_imports import RewriteImports
+    from protobuf.plugin._map_imports import MapImports
 
 _INDENT = " " * 4
 
@@ -223,7 +223,7 @@ def __init__(
         parameter: str,
         *,
         escape_module_with_hash: bool = False,
-        rewrite_imports: RewriteImports = (),
+        map_imports: MapImports = (),
     ) -> None:
         self.path = path
         self.module = module
@@ -232,7 +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._map_imports = map_imports
         self._indent = 0
         self._type_checking = False
         self._in_doc = False
@@ -334,39 +334,45 @@ def _to_el(self, v: object) -> str | Ident:
                 )
             case _:
                 return repr(v)
-        ident = self._relativize(self._rewrite_import(ident))
+        ident = self._relativize(self._map_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:
+    def _map_import(self, ident: Ident) -> Ident:
+        # Only imports derived from a descriptor know their Protobuf file.
+        if not self._map_imports or ident._desc is None:
+            return ident
+        if not _is_relative(ident.module):
             return ident
-        # A DescFile ident imports the module itself, so its import
-        # path includes the ident name.
         is_module_import = isinstance(ident._desc, DescFile)
+        if not is_module_import and _module_segments(ident.module) == _module_segments(
+            self.module
+        ):
+            # References to symbols in this file are not imports.
+            return ident
+        file = ident._desc if isinstance(ident._desc, DescFile) else ident._desc.file
+        target = map_import_target(file.name, self._map_imports)
+        if target is None:
+            return ident
         module_path = ident.module.path
         if is_module_import:
+            # A DescFile ident imports the module itself, so the module
+            # path includes the ident name.
             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
+        dotted = module_path.removeprefix(".")
+        mapped = f"{target}.{dotted}" if target else dotted
         if is_module_import:
-            parent, _, name = rewritten.rpartition(".")
+            parent, _, name = mapped.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
+            ident.name, Module(mapped), type_only=ident.type_only, _desc=ident._desc
         )
 
     def _relativize(self, ident: Ident) -> Ident:

src/protobuf/plugin/_map_imports.py

@@ -0,0 +1,68 @@
+# 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
+
+MapImports = tuple[tuple[re.Pattern[str], str], ...]
+
+_OPTION_NAME = "map_imports"
+
+_TARGET = re.compile(r"[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*")
+
+
+def compile_map_imports(mappings: dict[str, str]) -> MapImports:
+    compiled: list[tuple[re.Pattern[str], str]] = []
+    for pattern, target in mappings.items():
+        if not pattern:
+            msg = f"option '{_OPTION_NAME}': pattern must not be empty"
+            raise ValueError(msg)
+        # An empty target (or ".") maps to the canonical module path.
+        normalized = target.removesuffix(".")
+        if normalized and not _TARGET.fullmatch(normalized):
+            msg = f"option '{_OPTION_NAME}': target '{target}' must be a Python package path"
+            raise ValueError(msg)
+        compiled.append((_glob_to_regex(pattern), normalized))
+    return tuple(compiled)
+
+
+def map_import_target(proto_name: str, map_imports: MapImports) -> str | None:
+    for pattern, target in map_imports:
+        if pattern.fullmatch(proto_name):
+            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 + 2] == "*":
+                if pattern[i + 2 : i + 3] == "/":
+                    parts.append(r"([^/]+/)*")
+                    i += 3
+                    continue
+                parts.append(".*")
+                i += 2
+                continue
+            parts.append(r"[^/]*")
+        elif char == "/" and i == len(pattern) - 1:
+            # A trailing slash matches everything in the directory.
+            parts.append("/.*")
+        else:
+            parts.append(re.escape(char))
+        i += 1
+    return re.compile("".join(parts))

src/protobuf/plugin/_run.py

@@ -20,8 +20,8 @@
 
 from protobuf import maximum_supported_edition, minimum_supported_edition
 from protobuf.plugin._file import write
+from protobuf.plugin._map_imports import compile_map_imports
 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
 
@@ -168,7 +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),
+            map_imports=compile_map_imports(fw_opts.map_imports),
         )
 
         generate(schema)
@@ -232,4 +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)
+    map_imports: dict[str, str] = dataclasses.field(default_factory=dict)

src/protobuf/plugin/_schema.py

@@ -27,7 +27,7 @@
 
     from protobuf.wkt import CodeGeneratorRequest
 
-    from ._rewrite_imports import RewriteImports
+    from ._map_imports import MapImports
 
 T_co = TypeVar("T_co", covariant=True)
 
@@ -98,15 +98,15 @@ def __init__(
         name: str,
         version: str,
         escape_module_with_hash: bool,
-        rewrite_imports: RewriteImports = (),
+        map_imports: MapImports = (),
     ) -> 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
+        self._map_imports = map_imports
 
         file_to_generate = frozenset(req.file_to_generate)
         source_by_name = {s.name: s for s in req.source_file_descriptors}
@@ -170,7 +170,7 @@ def generate_file(
             self._version,
             self._parameter,
             escape_module_with_hash=self._escape_module_with_hash,
-            rewrite_imports=self._rewrite_imports,
+            map_imports=self._map_imports,
         )
         self._generated_files[path] = f
         return f

tests/buf.gen.yaml

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