Protobuf Zip Imports in Python

In Python, the most straightforward path to implementing a gRPC server for a Protobuf service is to use protoc to generate code that can be imported in a server, which then defines the service logic.

Let’s take a simple example Protobuf service:

syntax = "proto3";

package simple;

message HelloRequest {
  string name = 1;
}

message HelloResponse {
  string message = 1;
}

service Greeter {
  rpc SayHello (HelloRequest) returns (HelloResponse);
}

Next, we run some variant of python -m grpc_tools.protoc to generate code (assuming we’ve installed grpcio and grpcio-tools). Here’s an example for .proto files in a protos folder:

python -m grpc_tools.protoc --python_out=. --grpc_python_out=. --proto_path=protos protos/*.proto

This command outputs the following files

simple_pb2_grpc.py
simple_pb2.py

Within simple_pb2_grpc.py we see this import

import simple_pb2 as simple__pb2

This import can be problematic because it assumes that the generated code exists at the root of the project. If you want to keep your project structure organized, you probably want to put the generated code into a subfolder and gitignore it. The protoc tool doesn’t seem to support any options for Python code that will write these import statements differently. This limitation leaves us with only a few options:

1. As first mentioned, generate the code at the root of the project and deal with the suboptimal structure

python -m grpc_tools.protoc --python_out=. --grpc_python_out=. --proto_path=protos protos/*.proto

then in src/server.py do the following imports

import simple_pb2
import simple_pb2_grpc

We now have simple_pb2_grpc.py and simple_pb2.py in the project root and the server runs

❯ python -m src.server
Server started, listening on port 50051

2. Re-write the generated code to fix the imports for the package structure we want

❯ sed -i '' 's/import simple_pb2 as simple__pb2/from gen.protos import simple_pb2 as simple__pb2/' gen/protos/simple_pb2_grpc.py

Now the import is

from gen.protos import simple_pb2 as simple__pb2

and we can run the server

❯ python -m src.server
Server started, listening on port 50051

3. Augment the PYTHONPATH to add the target folder of the generated code to allow the generated imports to work

mkdir -p gen/protos
python -m grpc_tools.protoc --python_out=gen/protos --grpc_python_out=gen/protos --proto_path=protos protos/*.proto

If we add these imports

# Import the generated code from gen/protos
from gen.protos import simple_pb2, simple_pb2_grpc

to a src/server.py file then run it

python -m src.server

we still an error like this

python -m src.server
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/danielcorin/dev/toys/protobuf-zip-imports/src/server.py", line 8, in <module>
    from gen.protos import simple_pb2, simple_pb2_grpc
  File "/Users/danielcorin/dev/toys/protobuf-zip-imports/gen/protos/simple_pb2_grpc.py", line 6, in <module>
    import simple_pb2 as simple__pb2
ModuleNotFoundError: No module named 'simple_pb2'
make: *** [serve] Error 1

However, we can get this to work if we augment PYTHONPATH. This approach allows the import in the generated code and the package import in the server to both work.

# from src/server.py
from gen.protos import simple_pb2, simple_pb2_grpc

# from gen/protos/simple_pb2_grpc.py
import simple_pb2 as simple__pb2

It runs.

❯ export PYTHONPATH=gen/protos
❯ python -m src.server
Server started, listening on port 50051

Generating code to a zip archive

All of the aforementioned approaches require some degree of compromise in package structure or environment setup. This approach is no different, but I like it best, because it does not require modifying the generated code, the PYTHONPATH or creating what can feel like a mess in the project root, especially when you have many different protobuf services.

The approach is to create a zip archive of the Python generated code — something protoc supports out of the box.

python -m grpc_tools.protoc -I./protos --python_out=./gen.zip --grpc_python_out=./gen.zip protos/*.proto

This command creates gen.zip at the root of the project. When unarchived, we can see it contains these files:

❯ unzip gen.zip
Archive:  gen.zip
 extracting: simple_pb2.py
 extracting: simple_pb2_grpc.py

To make the imports work in our server, we can use zipimport

import zipimport

# Import the generated code from gen.zip
importer = zipimport.zipimporter("gen.zip")
simple_pb2 = importer.load_module("simple_pb2")
simple_pb2_grpc = importer.load_module("simple_pb2_grpc")

and our server runs

python -m src.server
Server started, listening on port 50051

The final src/server.py code looks like this.

import socket
import sys
from concurrent import futures
import grpc
import zipimport

# Import the generated code from gen.zip
importer = zipimport.zipimporter("gen.zip")
simple_pb2 = importer.load_module("simple_pb2")
simple_pb2_grpc = importer.load_module("simple_pb2_grpc")


class Greeter(simple_pb2_grpc.GreeterServicer):
    def SayHello(self, request, context):
        return simple_pb2.HelloResponse(message=f"Hello, {request.name}!")


def serve(port=50051):
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("localhost", port))
    except socket.error:
        print(f"Error: Port {port} is already in use")
        sys.exit(1)

    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    simple_pb2_grpc.add_GreeterServicer_to_server(Greeter(), server)
    server.add_insecure_port(f"[::]:{port}")
    server.start()
    print(f"Server started, listening on port {port}")
    server.wait_for_termination()


if __name__ == "__main__":
    serve()

You can also find this approach in project form here.