Skip to content
Browse files

Initial commit

  • Loading branch information...
0 parents commit 02e957f68487bb78f14ae30f1199f7952d45b110 @bdarnell committed May 18, 2012
2 .gitignore
@@ -0,0 +1,2 @@
+*.pyc
+/env/
0 codegenloader/__init__.py
No changes.
151 codegenloader/base.py
@@ -0,0 +1,151 @@
+import imp
+import itertools
+import os
+import sys
+
+_counter = itertools.count()
+
+class CodeGenLoader(object):
+ """Abstract base class for code generation import hooks.
+
+ The entry point for applications is to define a subclass
+ and use the `register` class method to set __path__.
+ This will make the module where __path__ was assigned
+ a pseudo-package from which the generated code can be imported.
+
+ The interface between this class and the python interpreter
+ is defined in PEP 302:
+ http://www.python.org/dev/peps/pep-0302/
+ """
+ # Class entry points
+ @classmethod
+ def register(cls, *args, **kwargs):
+ """Registers an import hook.
+
+ Arguments are passed (eventually) to `initialize`.
+ """
+ cls._install_hook()
+ hook_key = '__codegenloader_%d' % _counter.next()
+ if not hasattr(cls, "_register_args"):
+ cls._register_args = {}
+ cls._register_args[hook_key] = (args, kwargs)
+ return [hook_key]
+
+ @classmethod
+ def _install_hook(cls):
+ """Installs this loader as a ``path_hook``."""
+ if hasattr(cls, "_hook_installed"):
+ return
+ sys.path_hooks.append(cls)
+ cls._hook_installed = True
+
+ # PEP302 implementation
+ def __init__(self, path):
+ """Constructs a CodeGenLoader.
+
+ Implements the hook protocol from PEP 302: it is called with a
+ "path", and returns a loader if we can handle that path
+ (i.e. if the path is actually a unique token we created in
+ `register`), or raises an ImportError if not.
+ """
+ if path in self._register_args:
+ self.hook_key = path
+ args, kwargs = self._register_args[path]
+ self.initialize(*args, **kwargs)
+ else:
+ raise ImportError("not my path")
+
+ def initialize(self, modname, basedir):
+ """Real initialization function, independent of PEP302 requirements.
+
+ ``modname`` is the module name relative to which the generated code
+ will be imported.
+ ``basedir`` is the directory in which the source files for generation
+ can be found. If it is not an absolute path, it is interpreted
+ as relative to the file containing ``modname``
+ """
+ self.basename = modname + '.'
+ if os.path.isabs(basedir):
+ self.basedir = basedir
+ else:
+ self.basedir = os.path.join(
+ os.path.dirname(sys.modules[modname].__file__), basedir)
+ self.contents = {}
+
+ def find_module(self, fullname):
+ """Returns a loader object for the module ``fullname``, if it exists.
+
+ Implements the "finder" portion of the PEP 302 interface.
+ """
+ relname = self.get_relname(fullname)
+ try:
+ self.get_contents(relname)
+ # No error: we can load the module
+ return self
+ except KeyError:
+ if not self.can_generate(relname):
+ return None
+ self.generate(relname)
+ # Try again after generating
+ try:
+ self.get_contents(relname)
+ return self
+ except KeyError:
+ # Still not there
+ return None
+
+ def load_module(self, fullname):
+ """Returns the module named ``fullname``.
+
+ Implements the "loader" portion of the PEP 302 interface.
+ """
+ relname = self.get_relname(fullname)
+ # This should never fail since python always calls find_module first
+ is_pkg, contents = self.get_contents(relname)
+ mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
+ mod.__file__ = '<codegenloader %s>' % fullname
+ mod.__loader__ = self
+ if is_pkg:
+ mod.__path__ = [self.hook_key]
+ exec contents in mod.__dict__
+ return mod
+
+ # Internal methods
+ def get_relname(self, fullname):
+ """Converts a fully-qualified module name to a relative one."""
+ assert fullname.startswith(self.basename)
+ return fullname[len(self.basename):]
+
+ def get_contents(self, relname):
+ """Return a tuple (is_pkg, contents) if code is stored for this module.
+
+ If the code is not found, raises KeyError.
+ """
+ relpath = relname.replace('.', '/')
+ init_path = relpath + '/__init__.py'
+ if init_path in self.contents:
+ return True, self.contents[init_path]
+ mod_path = relpath + '.py'
+ if mod_path in self.contents:
+ return False, self.contents[mod_path]
+ raise KeyError("code not found for %s" % relname)
+
+ def store_contents(self, relpath, contents):
+ """Store the contents of a file at relpath.
+
+ To be called from subclasses after code has been generated.
+ """
+ assert relpath not in self.contents
+ self.contents[relpath] = contents
+
+ # Methods for overriding in subclass
+ def can_generate(self, relname):
+ """Should return True if we can generate a module named ``relname``."""
+ raise NotImplementedError()
+
+ def generate(self, relname):
+ """Generate code for module ``relname``.
+
+ Should call `store_contents` for any files generated.
+ """
+ raise NotImplementedError()
50 codegenloader/protobuf.py
@@ -0,0 +1,50 @@
+from __future__ import with_statement
+import os
+import shutil
+import subprocess
+import tempfile
+from codegenloader.base import CodeGenLoader
+
+class ProtobufLoader(CodeGenLoader):
+ def protoname(self, relname):
+ assert relname.endswith("_pb2")
+ relname = relname[:-len("_pb2")]
+ return os.path.abspath(os.path.join(self.basedir, relname + '.proto'))
+
+ def can_generate(self, relname):
+ if relname.endswith("_pb2"):
+ return os.path.exists(self.protoname(relname))
+ else:
+ return False
+
+ def generate(self, relname):
+ tempdir = tempfile.mkdtemp(prefix='codegenloader')
+ try:
+ protodir, protofile = os.path.split(self.protoname(relname))
+ subprocess.check_call(
+ ["protoc",
+ "--python_out=.",
+ "--proto_path=%s" % protodir,
+ self.protoname(relname)],
+ cwd=tempdir)
+ relpath = relname + ".py"
+ with open(os.path.join(tempdir, relpath)) as f:
+ self.store_contents(relpath, f.read())
+ finally:
+ shutil.rmtree(tempdir)
+
+
+def make_path(modname, basedir):
+ """Returns a object to be set as __path__.
+
+ This is the visible entry point to this module. To use it,
+ assign the result of this function to ``__path__``:
+
+ import dropbox.codegenloader.protobuf
+ __path__ = dropbox.codegenloader.protobuf.make_path(__name__, "proto")
+
+ The first argument should always be __name__; the second is a
+ directory name that contains the .proto files. (relative to the
+ file where make_path is called).
+ """
+ return ProtobufLoader.register(modname, basedir)
0 codegenloader/test/__init__.py
No changes.
2 codegenloader/test/proto/__init__.py
@@ -0,0 +1,2 @@
+import codegenloader.protobuf
+__path__ = codegenloader.protobuf.make_path(__name__, ".")
2 codegenloader/test/proto/simple.proto
@@ -0,0 +1,2 @@
+message Simple {
+}
7 codegenloader/test/protobuf_test.py
@@ -0,0 +1,7 @@
+import unittest
+from google.protobuf.message import Message
+
+class ProtobufTest(unittest.TestCase):
+ def test_simple_import(self):
+ from codegenloader.test.proto.simple_pb2 import Simple
+ self.assertTrue(issubclass(Simple, Message))
14 codegenloader/test/runtests.py
@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+
+import unittest
+
+TEST_MODULES = [
+ 'codegenloader.test.thrift_test',
+ 'codegenloader.test.protobuf_test',
+ ]
+
+def all():
+ return unittest.defaultTestLoader.loadTestsFromNames(TEST_MODULES)
+
+if __name__ == '__main__':
+ unittest.main(defaultTest='all')
2 codegenloader/test/thrift/simple.thrift
@@ -0,0 +1,2 @@
+struct Simple {
+}
6 codegenloader/test/thrift_test.py
@@ -0,0 +1,6 @@
+import unittest
+
+class ThriftTest(unittest.TestCase):
+ def test_simple_import(self):
+ from codegenloader.test.thriftgen.simple.ttypes import Simple
+ self.assertTrue(isinstance(Simple, type))
2 codegenloader/test/thriftgen.py
@@ -0,0 +1,2 @@
+import codegenloader.thrift
+__path__ = codegenloader.thrift.make_path(__name__, "thrift")
45 codegenloader/thrift.py
@@ -0,0 +1,45 @@
+from __future__ import with_statement
+import glob
+import os
+import shutil
+import subprocess
+import tempfile
+from codegenloader.base import CodeGenLoader
+
+class ThriftLoader(CodeGenLoader):
+ def thriftname(self, relname):
+ return os.path.abspath(os.path.join(self.basedir, relname + '.thrift'))
+
+ def can_generate(self, relname):
+ return os.path.exists(self.thriftname(relname))
+
+ def generate(self, relname):
+ tempdir = tempfile.mkdtemp(prefix='codegenloader')
+ try:
+ subprocess.check_call(
+ ["thrift", "--gen", "py:new_style", self.thriftname(relname)],
+ cwd=tempdir)
+ outdir = os.path.join(tempdir, 'gen-py')
+ for fn in glob.glob("%s/%s/*.py" % (outdir, relname)):
+ assert fn.startswith(outdir + '/')
+ relpath = fn[len(outdir) + 1:]
+ with open(fn) as f:
+ self.store_contents(relpath, f.read())
+ finally:
+ shutil.rmtree(tempdir)
+
+
+def make_path(modname, basedir):
+ """Returns a object to be set as __path__.
+
+ This is the visible entry point to this module. To use it,
+ assign the result of this function to ``__path__``:
+
+ import dropbox.codegenloader.thrift
+ __path__ = dropbox.codegenloader.thrift.make_path(__name__, "thrift")
+
+ The first argument should always be __name__; the second is a
+ directory name that contains the thrift files. (relative to the
+ file where make_path is called).
+ """
+ return ThriftLoader.register(modname, basedir)

0 comments on commit 02e957f

Please sign in to comment.
Something went wrong with that request. Please try again.