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?

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,
    )