Add option NANOPB_PB2_TEMP_DIR to store nanopb_pb2.py in a temporary directory (#601)dev_dynamic_nanopb_pb2_2

When environment variable NANOPB_PB2_TEMP_DIR is defined, the generator will
store the autogenerated nanopb_pb2.py in a temporary directory. This can
be the value of the environment variable, if it is valid directory, or
otherwise system-wide temp directory is used.

This avoids polluting the source code folder with autogenerated file, and may avoid some version
incompatibility issues. By default the old behavior of storing nanopb_pb2.py under proto folder is
preserved.

4 files changed, 139 insertions, 79 deletions

diff --git a/generator/nanopb_generator.py b/generator/nanopb_generator.py
index 115d797..dc39246 100755
--- a/generator/nanopb_generator.py
+++ b/generator/nanopb_generator.py
@@ -25,16 +25,6 @@ if not os.getenv("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"):
     os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 try:
-    # Add some dummy imports to keep packaging tools happy.
-    import google # bbfreeze seems to need these
-    import pkg_resources # pyinstaller / protobuf 2.5 seem to need these
-    import proto.nanopb_pb2 as nanopb_pb2 # pyinstaller seems to need this
-    import pkg_resources.py2_warn
-except:
-    # Don't care, we will error out later if it is actually important.
-    pass
-
-try:
     # Make sure grpc_tools gets included in binary package if it is available
     import grpc_tools.protoc
 except:
@@ -48,57 +38,42 @@ try:
     import google.protobuf.descriptor
 except:
     sys.stderr.write('''
-         *************************************************************
-         *** Could not import the Google protobuf Python libraries ***
-         *** Try installing package 'python3-protobuf' or similar.  ***
-         *************************************************************
+         **********************************************************************
+         *** Could not import the Google protobuf Python libraries          ***
+         ***                                                                ***
+         *** Easiest solution is often to install the dependencies via pip: ***
+         ***    pip install protobuf grpcio-tools                           ***
+         **********************************************************************
     ''' + '\n')
     raise
 
-try:
-    from .proto import nanopb_pb2
-    from .proto._utils import invoke_protoc
-except TypeError:
-    sys.stderr.write('''
-         ****************************************************************************
-         *** Got TypeError when importing the protocol definitions for generator. ***
-         *** This usually means that the protoc in your path doesn't match the    ***
-         *** Python protobuf library version.                                     ***
-         ***                                                                      ***
-         *** Please check the output of the following commands:                   ***
-         *** which protoc                                                         ***
-         *** protoc --version                                                     ***
-         *** python3 -c 'import google.protobuf; print(google.protobuf.__file__)'  ***
-         *** If you are not able to find the python protobuf version using the    ***
-         *** above command, use this command.                                     ***
-         *** pip freeze | grep -i protobuf                                        ***
-         ****************************************************************************
-    ''' + '\n')
-    raise
-except (ValueError, SystemError, ImportError):
-    # Probably invoked directly instead of via installed scripts.
-    import proto.nanopb_pb2 as nanopb_pb2
+# Depending on how this script is run, we may or may not have PEP366 package name
+# available for relative imports.
+if not __package__:
+    import proto
     from proto._utils import invoke_protoc
-except:
-    sys.stderr.write('''
-         ********************************************************************
-         *** Failed to import the protocol definitions for generator.     ***
-         *** You have to run 'make' in the nanopb/generator/proto folder. ***
-         ********************************************************************
-    ''' + '\n')
-    raise
+    from proto import TemporaryDirectory
+else:
+    from . import proto
+    from .proto._utils import invoke_protoc
+    from .proto import TemporaryDirectory
+
+if getattr(sys, 'frozen', False):
+    # Binary package, just import the file
+    from proto import nanopb_pb2
+else:
+    # Try to rebuild nanopb_pb2.py if necessary
+    nanopb_pb2 = proto.load_nanopb_pb2()
 
 try:
-    from tempfile import TemporaryDirectory
-except ImportError:
-    class TemporaryDirectory:
-        '''TemporaryDirectory fallback for Python 2'''
-        def __enter__(self):
-            self.dir = tempfile.mkdtemp()
-            return self.dir
-
-        def __exit__(self, *args):
-            shutil.rmtree(self.dir)
+    # Add some dummy imports to keep packaging tools happy.
+    import google # bbfreeze seems to need these
+    import pkg_resources # pyinstaller / protobuf 2.5 seem to need these
+    from proto import nanopb_pb2 # pyinstaller seems to need this
+    import pkg_resources.py2_warn
+except:
+    # Don't care, we will error out later if it is actually important.
+    pass
 
 # ---------------------------------------------------------------------------
 #                     Generation of single fields
diff --git a/generator/proto/__init__.py b/generator/proto/__init__.py
index e2e8000..cfab910 100644
--- a/generator/proto/__init__.py
+++ b/generator/proto/__init__.py
@@ -1,34 +1,37 @@
-'''This file automatically rebuilds the proto definitions for Python.'''
+'''This file dynamically builds the proto definitions for Python.'''
 from __future__ import absolute_import
 
+import os
 import os.path
 import sys
-
+import tempfile
+import shutil
+import traceback
 import pkg_resources
-
 from ._utils import has_grpcio_protoc, invoke_protoc, print_versions
 
-dirname = os.path.dirname(__file__)
-protosrc = os.path.join(dirname, "nanopb.proto")
-protodst = os.path.join(dirname, "nanopb_pb2.py")
-rebuild = False
+# Compatibility layer to make TemporaryDirectory() available on Python 2.
+try:
+    from tempfile import TemporaryDirectory
+except ImportError:
+    class TemporaryDirectory:
+        '''TemporaryDirectory fallback for Python 2'''
+        def __init__(self, prefix = 'tmp', dir = None):
+            self.prefix = prefix
+            self.dir = dir
 
-if os.path.isfile(protosrc):
-    src_date = os.path.getmtime(protosrc)
-    if not os.path.isfile(protodst) or os.path.getmtime(protodst) < src_date:
-        rebuild = True
+        def __enter__(self):
+            self.dir = tempfile.mkdtemp(prefix = self.prefix, dir = self.dir)
+            return self.dir
 
-if not rebuild:
-    try:
-        from . import nanopb_pb2
-    except AttributeError as e:
-        rebuild = True
-        sys.stderr.write("Failed to import nanopb_pb2.py: " + str(e) + "\n"
-                         "Will automatically attempt to rebuild this.\n"
-                         "Verify that python-protobuf and protoc versions match.\n")
-        print_versions()
+        def __exit__(self, *args):
+            shutil.rmtree(self.dir)
+
+def build_nanopb_proto(protosrc, dirname):
+    '''Try to build a .proto file for python-protobuf.
+    Returns True if successful.
+    '''
 
-if rebuild:
     cmd = [
         "protoc",
         "--python_out={}".format(dirname),
@@ -40,11 +43,86 @@ if rebuild:
         # grpcio-tools has an extra CLI argument
         # from grpc.tools.protoc __main__ invocation.
         _builtin_proto_include = pkg_resources.resource_filename('grpc_tools', '_proto')
-
         cmd.append("-I={}".format(_builtin_proto_include))
+
     try:
         invoke_protoc(argv=cmd)
     except:
         sys.stderr.write("Failed to build nanopb_pb2.py: " + ' '.join(cmd) + "\n")
-        raise
+        sys.stderr.write(traceback.format_exc() + "\n")
+        return False
+
+    return True
+
+def load_nanopb_pb2():
+    # To work, the generator needs python-protobuf built version of nanopb.proto.
+    # There are three methods to provide this:
+    #
+    # 1) Load a previously generated generator/proto/nanopb_pb2.py
+    # 2) Use protoc to build it and store it permanently generator/proto/nanopb_pb2.py
+    # 3) Use protoc to build it, but store only temporarily in system-wide temp folder
+    #
+    # By default these are tried in numeric order.
+    # If NANOPB_PB2_TEMP_DIR environment variable is defined, the 2) is skipped.
+    # If the value of the $NANOPB_PB2_TEMP_DIR exists as a directory, it is used instead
+    # of system temp folder.
+
+    build_error = None
+    proto_ok = False
+    tmpdir = os.getenv("NANOPB_PB2_TEMP_DIR")
+    temporary_only = (tmpdir is not None)
+    dirname = os.path.dirname(__file__)
+    protosrc = os.path.join(dirname, "nanopb.proto")
+    protodst = os.path.join(dirname, "nanopb_pb2.py")
+    proto_ok = False
+
+    if tmpdir is not None and not os.path.isdir(tmpdir):
+        tmpdir = None # Use system-wide temp dir
+
+    if os.path.isfile(protosrc):
+        src_date = os.path.getmtime(protosrc)
+        if not os.path.isfile(protodst) or os.path.getmtime(protodst) < src_date:
+            # Outdated, rebuild
+            proto_ok = False
+        else:
+            try:
+                from . import nanopb_pb2 as nanopb_pb2_mod
+                proto_ok = True
+            except Exception as e:
+                sys.stderr.write("Failed to import nanopb_pb2.py: " + str(e) + "\n"
+                                "Will automatically attempt to rebuild this.\n"
+                                "Verify that python-protobuf and protoc versions match.\n")
+                print_versions()
+
+    # Try to rebuild into generator/proto directory
+    if not proto_ok and not temporary_only:
+        proto_ok = build_nanopb_proto(protosrc, dirname)
+
+        try:
+            from . import nanopb_pb2 as nanopb_pb2_mod
+        except:
+            sys.stderr.write("Failed to import generator/proto/nanopb_pb2.py:\n")
+            sys.stderr.write(traceback.format_exc() + "\n")
+
+    # Try to rebuild into temporary directory
+    if not proto_ok:
+        with TemporaryDirectory(prefix = 'nanopb-', dir = tmpdir) as protodir:
+            proto_ok = build_nanopb_proto(protosrc, protodir)
+
+            if protodir not in sys.path:
+                sys.path.insert(0, protodir)
+
+            try:
+                import nanopb_pb2 as nanopb_pb2_mod
+            except:
+                sys.stderr.write("Failed to import %s/nanopb_pb2.py:\n" % protodir)
+                sys.stderr.write(traceback.format_exc() + "\n")
+
+    # If everything fails
+    if not proto_ok:
+        sys.stderr.write("\n\nGenerating nanopb_pb2.py failed.\n")
+        sys.stderr.write("Make sure that a protoc generator is available and matches python-protobuf version.\n")
+        print_versions()
+        sys.exit(1)
 
+    return nanopb_pb2_mod
diff --git a/generator/proto/_utils.py b/generator/proto/_utils.py
index bd49537..f9c8c94 100644
--- a/generator/proto/_utils.py
+++ b/generator/proto/_utils.py
@@ -57,6 +57,7 @@ def print_versions():
 
     try:
         import google.protobuf
+        sys.stderr.write("Python version " + sys.version + "\n")
         sys.stderr.write("Using python-protobuf from " + google.protobuf.__file__ + "\n")
         sys.stderr.write("Python-protobuf version: " + google.protobuf.__version__ + "\n")
     except Exception as e:
diff --git a/generator/protoc b/generator/protoc
index 0d952a3..c259702 100755
--- a/generator/protoc
+++ b/generator/protoc
@@ -6,7 +6,13 @@
 import sys
 import os
 import os.path
-from nanopb_generator import invoke_protoc
+
+# Depending on how this script is run, we may or may not have PEP366 package name
+# available for relative imports.
+if not __package__:
+    from proto._utils import invoke_protoc
+else:
+    from .proto._utils import invoke_protoc
 
 if __name__ == '__main__':
     # Get path of the directory where this script is stored.