DOCSEARCHbyAlgolia
Submit Docs

protoletariat

repository·main·Indexed 42 days ago

https://github.com/cpcloud/protoletariat

A tool to fix broken absolute imports in Python code generated by protoc or buf. It automates the conversion of absolute imports into relative imports by inspecting the FileDescriptorSet, making generated code easier to use in Python packages. Version 3.3.10 provides a CLI (protol) and a Python API including ASTRewriter and ASTImportRewriter for AST-based import transformations.

Tokens
4.4K
Snippets
15
Records
18
Agent score
51%

What's inside protoletariat

  1. How protoletariat converts imports

    main
    At a high level, protoletariat converts absolute imports to relative imports. It specifically targets imports that were generated from .proto files by inspecting the FileDescriptorProtos. The rewrite mechanism uses Python AST pattern matching to identify and apply rewrite rules to matched import patterns.

  2. Use protoletariat to fix protobuf imports

    main

    The protol CLI is designed to be run as a post-processing step after running protoc. It converts absolute imports in generated Python files into relative imports by inspecting the FileDescriptorSet generated from your .proto files.

    To use it, you can wrap your existing protoc command using the protoc subcommand of protol. This allows protol to handle both the generation and the subsequent import rewriting in one step.

    # 1. Generate code with protoc (standard way)
$ mkdir out
$ protoc --python_out=out --proto_path=path/to/protos thing1.proto thing2.proto

# 2. Fix imports using protol (recommended way)
# This wraps the protoc command and applies fixes to the output directory
$ protol \
  --create-package \
  --in-place \
  --python-out out \
  protoc --proto-path=path/to/protos thing1.proto thing2.proto

  3. Install protoletariat

    main

    You can install protoletariat via PyPI, Conda, or by pulling a Docker image. This tool is used to fix broken absolute imports in Python code generated by protoc.

    # Via PyPI
pip install protoletariat

# Via Conda
conda install protoletariat -c conda-forge

# Via Docker
docker pull ghcr.io/cpcloud/protoletariat:latest

  4. Use the protoletariat CLI to rewrite protobuf imports

    main

    The protoletariat CLI rewrites imports in Python code generated by protoc or buf to make them compatible with the protoletariat ecosystem.

    To use the CLI, you must provide a directory containing the generated Python code using the --python-out (or -o) flag. By default, the CLI echoes the modified code to stdout. To modify the files directly, use the --in-place flag.

    # Example: Rewrite imports in-place for generated code in ./gen
protoletariat -o ./gen --in-place

  5. Reference the protol CLI options and subcommands

    main

    The protol command is used to rewrite protoc or buf-generated imports.

    Subcommands

    • protoc: Uses protoc to generate FileDescriptorSet bytes.
    • buf: Uses buf to generate FileDescriptorSet bytes.
    • raw: Uses FileDescriptorSet bytes provided via a file or stdin.

    Options

    • -o, --python-out DIRECTORY (required): Directory containing the generated Python code.
    • --in-place / --not-in-place: Overwrite files under --python-out with adjusted imports (default: --not-in-place).
    • --create-package / --dont-create-package: Recursively create __init__.py files under --python-out (default: --dont-create-package).
    • -s, --module-suffixes TEXT: Suffixes of Python/mypy modules to process (default: _pb2.py, _pb2.pyi, _pb2_grpc.py, _pb2_grpc.pyi).
    • --exclude-google-imports / --dont-exclude-google-imports: Exclude rewriting imports prefixed with google/protobuf.
    • -e, --exclude-imports-glob TEXT: Exclude imports matching a glob pattern (multiple values allowed).
    Usage: protol [OPTIONS] COMMAND [ARGS]...

  Rewrite protoc or buf-generated imports for use by the protoletariat.

Options:
  -o, --python-out DIRECTORY      Directory containing protoc or buf-generated Python code  [required]
  --in-place / --not-in-place     Overwrite all relevant files under `--python-out` with adjusted imports  [default: not-in-place]
  --create-package / --dont-create-package
                                  Recursively create __init__.py files under `--python-out`  [default: dont-create-package]
  -s, --module-suffixes TEXT      Suffixes of Python/mypy modules to process  [default: _pb2.py, _pb2.pyi, _pb2_grpc.py, _pb2_grpc.pyi]
  --exclude-google-imports / --dont-exclude-google-imports
                                  Exclude rewriting imports prefixed with google/protobuf
  -e, --exclude-imports-glob TEXT  Exclude imports matching a glob pattern from being rewritten. Multiple values are allowed
  --help                           Show this message and exit.

Commands:
  buf     Use buf to generate the FileDescriptorSet blob
  protoc  Use protoc to generate the FileDescriptorSet blob
  raw     Rewrite imports using FileDescriptorSet bytes from a file or stdin

  6. Implement FileDescriptorSetGenerator to fix imports

    main

    To use protoletariat to fix imports in generated Python code, you must implement a subclass of FileDescriptorSetGenerator. This base class provides the fix_imports method, which automates the process of rewriting AST imports in files generated by protoc or buf to ensure they work correctly in a Python environment.

    To use fix_imports, you must provide:

    • python_out: The directory where the generated Python files reside.
    • create_package: A boolean indicating whether to create __init__.py files to form Python packages.
    • overwrite_callback: A function Callable[[Path, str], None] that takes the file Path and the new_code (string) and performs the actual write operation.
    • module_suffixes: A sequence of suffixes to look for (e.g., ['.py', '.pyi']).
    • exclude_imports_glob: A sequence of glob patterns to ignore specific imports (e.g., ['google/protobuf/*']).
    class MyGenerator(FileDescriptorSetGenerator):
    def generate_file_descriptor_set_bytes(self) -> bytes:
        # Return bytes of a FileDescriptorSet
        ...

# Usage
generator = MyGenerator()
generator.fix_imports(
    python_out=Path("out/"),
    create_package=True,
    overwrite_callback=lambda path, code: path.write_text(code),
    module_suffixes=".py",
    exclude_imports_glob=["google/protobuf/*"]
)

  7. Use ASTRewriter to perform AST pattern matching and rewriting

    main

    The ASTRewriter class allows you to define custom rewrite rules for Python Abstract Syntax Trees (AST). You register a pattern (an AST node) using the @register decorator and provide a function that returns a new AST node to replace the matched pattern.

    To apply the rules, call rewrite(node) on the AST node you wish to transform. If no pattern matches, the original node is returned.

    import ast
from protoletariat.rewrite import ASTRewriter

rewriter = ASTRewriter()
pattern = ast.parse("x = 1").body[0]

@rewriter.register(pattern)
def make_it_two(node):
    return ast.parse("x = 2").body[0]

node = ast.parse("x = 1").body[0]
rewritten = rewriter.rewrite(node)
# rewritten is now the AST for 'x = 2'

  8. Match AST nodes with the matches() function

    main

    The matches(value, pattern) function checks if a given AST value matches a pattern node.

    Matching logic:

    • Types must match exactly.
    • If both are non-string iterables, it recursively checks all elements.
    • If both are primitive values (like None, True, False), it uses standard equality.
    • If both are AST nodes, it recursively checks all fields defined in the pattern's _fields.
    import ast
from protoletariat.rewrite import matches

node = ast.parse("x = 1").body[0]
pattern = ast.parse("x = 1").body[0]

print(matches(node, pattern))  # True

  9. Generate protobuf import replacements with build_rewrites()

    main

    The build_rewrites(proto, dep, *, is_public=False) function calculates the necessary Replacement mappings to convert absolute protobuf imports into relative imports based on a target path.

    • proto: The dependent protobuf path (e.g., "a/b/c").
    • dep: The name of the dependency's FileDescriptorProto message, without the .proto suffix (e.g., "foo/bar").
    • is_public: If True, generates additional replacements for star imports (from ... import *).

    Returns a list of Replacement(old, new) objects.

    from protoletariat.rewrite import build_rewrites, Replacement
from pprint import pprint

proto = "a/b/c"
dep = "foo/bar"
replacements = build_rewrites(proto, dep)
pprint(replacements)
# [Replacement(old='from foo import bar_pb2 as foo_dot_bar__pb2', new='from ...foo import bar_pb2 as foo_dot_bar__pb2'),
#  Replacement(old='import foo.bar_pb2', new='from ... import foo_pb2')]

  10. Rewrite source code imports with ASTImportRewriter

    main

    The ASTImportRewriter is a high-level utility designed specifically to transform import statements in a source string using Replacement rules.

    1. Initialize ASTImportRewriter().
    2. Register rules using register_rewrite(replacement) where replacement is a Replacement(old, new) object containing valid Python import statements.
    3. Apply the transformations to a string of code using rewrite(src).

    It uses an internal ImportNodeTransformer to ensure that duplicate imports are not introduced during the rewriting process.

    from protoletariat.rewrite import ASTImportRewriter, Replacement

rewriter = ASTImportRewriter()

# Define a rule to change an absolute import to a relative one
rule = Replacement(old='from foo import bar_pb2 as foo_dot_bar__pb2', 
                   new='from ...foo import bar_pb2 as foo_dot_bar__pb2')

rewriter.register_rewrite(rule)

source = "from foo import bar_pb2 as foo_dot_bar__pb2"
new_source = rewriter.rewrite(source)
print(new_source)
# from ...foo import bar_pb2 as foo_dot_bar__pb2

  11. Configure protoletariat global options

    main

    The following options are available for the main protoletariat command and affect all subcommands (protoc, buf, raw).

    OptionFlagDefaultDescription
    --python-out-oRequiredDirectory containing protoc or buf-generated Python code
    --in-place--in-placeFalseOverwrite all relevant files under --python-out with adjusted imports
    --create-package--create-packageFalseRecursively create __init__.py files under --python-out
    --module-suffixes-s_pb2.py, _pb2.pyi, _pb2_grpc.py, _pb2_grpc.pyiSuffixes of Python/mypy modules to process
    --exclude-google-imports--exclude-google-importsTrueExclude rewriting imports prefixed with google/protobuf
    --exclude-imports-glob-e[]Exclude imports matching a glob pattern from being rewritten (multiple allowed)

  12. Use the `buf` subcommand to generate and fix imports

    main

    The buf subcommand uses the buf executable to generate the FileDescriptorSet and then applies the import rewriting logic.

    Arguments:

    • input: The input path for buf (defaults to current directory).

    Optional Arguments/Options:

    • --buf-path: Path to the buf executable. Can be set via BUF_PATH environment variable. Defaults to buf.

    Example:

    protoletariat -o ./gen --in-place buf ./buf.yaml