Automate compilation of protobuf specs into python classes in setup.py

Question

I have a python project that uses google protobufs as a message format for communicating over the network. Generating python files from the .proto files is straight-forward using the protoc program. How can I configure my setup.py file for the project so that it automatically calls the protoc command?

Desalinate

Do you expect the protos to change often? If not, could you just run protoc on your side and distribute the generated .py files? (You'll still need to take a dependency on protobuf for the base files.) – Peculium

Answer 1

In a similar situation, I ended up with this code (setup.py, but written in a way to allow extraction into some external Python module for reuse). Note that I took the generate_proto function and several ideas from the setup.py file of the protobuf source distribution.

from __future__ import print_function

import os
import shutil
import subprocess
import sys

from distutils.command.build_py import build_py as _build_py
from distutils.command.clean import clean as _clean
from distutils.debug import DEBUG
from distutils.dist import Distribution
from distutils.spawn import find_executable
from nose.commands import nosetests as _nosetests
from setuptools import setup

PROTO_FILES = [
    'goobuntu/proto/hoststatus.proto',
    ]

CLEANUP_SUFFIXES = [
    # filepath suffixes of files to remove on "clean" subcommand
    '_pb2.py',
    '.pyc',
    '.so',
    '.o',
    'dependency_links.txt',
    'entry_points.txt',
    'PKG-INFO',
    'top_level.txt',
    'SOURCES.txt',
    '.coverage',
    'protobuf/compiler/__init__.py',
    ]

CLEANUP_DIRECTORIES = [  # subdirectories to remove on "clean" subcommand
    # 'build'  # Note: the build subdirectory is removed if --all is set.
    'html-coverage',
    ]

if 'PROTOC' in os.environ and os.path.exists(os.environ['PROTOC']):
  protoc = os.environ['PROTOC']
else:
  protoc = find_executable('protoc')


def generate_proto(source):
  """Invoke Protocol Compiler to generate python from given source .proto."""
  if not os.path.exists(source):
    sys.stderr.write('Can\'t find required file: %s\n' % source)
    sys.exit(1)

  output = source.replace('.proto', '_pb2.py')
  if (not os.path.exists(output) or
      (os.path.getmtime(source) > os.path.getmtime(output))):
    if DEBUG:
      print('Generating %s' % output)

    if protoc is None:
      sys.stderr.write(
          'protoc not found. Is protobuf-compiler installed? \n'
          'Alternatively, you can point the PROTOC environment variable at a '
          'local version.')
      sys.exit(1)

    protoc_command = [protoc, '-I.', '--python_out=.', source]
    if subprocess.call(protoc_command) != 0:
      sys.exit(1)


class MyDistribution(Distribution):
  # Helper class to add the ability to set a few extra arguments
  # in setup():
  # protofiles : Protocol buffer definitions that need compiling
  # cleansuffixes : Filename suffixes (might be full names) to remove when
  #                   "clean" is called
  # cleandirectories : Directories to remove during cleanup
  # Also, the class sets the clean, build_py, test and nosetests cmdclass
  # options to defaults that compile protobufs, implement test as nosetests
  # and enables the nosetests command as well as using our cleanup class.

  def __init__(self, attrs=None):
    self.protofiles = []  # default to no protobuf files
    self.cleansuffixes = ['_pb2.py', '.pyc']  # default to clean generated files
    self.cleandirectories = ['html-coverage']  # clean out coverage directory
    cmdclass = attrs.get('cmdclass')
    if not cmdclass:
      cmdclass = {}
    # These should actually modify attrs['cmdclass'], as we assigned the
    # mutable dict to cmdclass without copying it.
    if 'nosetests' not in cmdclass:
      cmdclass['nosetests'] = MyNosetests
    if 'test' not in cmdclass:
      cmdclass['test'] = MyNosetests
    if 'build_py' not in cmdclass:
      cmdclass['build_py'] = MyBuildPy
    if 'clean' not in cmdclass:
      cmdclass['clean'] = MyClean
    attrs['cmdclass'] = cmdclass
    # call parent __init__ in old style class
    Distribution.__init__(self, attrs)


class MyClean(_clean):

  def run(self):
    try:
      cleandirectories = self.distribution.cleandirectories
    except AttributeError:
      sys.stderr.write(
          'Error: cleandirectories not defined. MyDistribution not used?')
      sys.exit(1)
    try:
      cleansuffixes = self.distribution.cleansuffixes
    except AttributeError:
      sys.stderr.write(
          'Error: cleansuffixes not defined. MyDistribution not used?')
      sys.exit(1)
    # Remove build and html-coverage directories if they exist
    for directory in cleandirectories:
      if os.path.exists(directory):
        if DEBUG:
          print('Removing directory: "{}"'.format(directory))
        shutil.rmtree(directory)
    # Delete generated files in code tree.
    for dirpath, _, filenames in os.walk('.'):
      for filename in filenames:
        filepath = os.path.join(dirpath, filename)
        for i in cleansuffixes:
          if filepath.endswith(i):
            if DEBUG:
              print('Removing file: "{}"'.format(filepath))
            os.remove(filepath)
    # _clean is an old-style class, so super() doesn't work
    _clean.run(self)


class MyBuildPy(_build_py):

  def run(self):
    try:
      protofiles = self.distribution.protofiles
    except AttributeError:
      sys.stderr.write(
          'Error: protofiles not defined. MyDistribution not used?')
      sys.exit(1)
    for proto in protofiles:
      generate_proto(proto)
    # _build_py is an old-style class, so super() doesn't work
    _build_py.run(self)


class MyNosetests(_nosetests):

  def run(self):
    try:
      protofiles = self.distribution.protofiles
    except AttributeError:
      sys.stderr.write(
          'Error: protofiles not defined. MyDistribution not used?')
    for proto in protofiles:
      generate_proto(proto)
    # _nosetests is an old-style class, so super() doesn't work
    _nosetests.run(self)


setup(
    # MyDistribution automatically enables several extensions, including
    # the compilation of protobuf files.
    distclass=MyDistribution,
    ...
    tests_require=['nose'],
    protofiles=PROTO_FILES,
    cleansuffixes=CLEANUP_SUFFIXES,
    cleandirectories=CLEANUP_DIRECTORIES,
    )

Othilia

Answer 2

Here's the solution that I have used for setup.py. The only thing you need to keep in mind is the version of the protoc compiler is compatible with the installed protobuf version.

    '''
    # here you can specify the proto folder and 
    # output folder or input them as parameters 
    # to script

    protoc_command = [
    "python", "-m", "grpc_tools.protoc",
    f"--proto_path={proto_folder}",
    f"--python_out={output_folder}",
    f"--grpc_python_out={output_folder}",
    ]   
    '''

Chelton