Move TemporaryDirectory() Python 2 compatibility code to generator/proto package.dev_dynamic_nanopb_pb2

2 files changed, 41 insertions, 53 deletions

diff --git a/generator/nanopb_generator.py b/generator/nanopb_generator.py
index e39080d..f25e857 100755
--- a/generator/nanopb_generator.py
+++ b/generator/nanopb_generator.py
@@ -57,33 +57,9 @@ except:
     ''' + '\n')
     raise
 
-try:
-    from .proto import nanopb_pb2
-    from .proto._utils import invoke_protoc
-except:
-    sys.stderr.write('''
-         *******************************************************************************
-         *** Got an exception when importing the protocol definitions for generator. ***
-         *** This usually means that the protoc in your path doesn't match the       ***
-         *** Python protobuf library version.                                        ***
-         ***                                                                         ***
-         *** Easiest solution is often to install the dependencies via pip:          ***
-         ***    pip install protobuf grpcio-tools                                    ***
-         *******************************************************************************
-    ''' + '\n')
-    raise
-
-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)
+from .proto import nanopb_pb2
+from .proto._utils import invoke_protoc
+from .proto import TemporaryDirectory
 
 # ---------------------------------------------------------------------------
 #                     Generation of single fields
diff --git a/generator/proto/__init__.py b/generator/proto/__init__.py
index 479cacf..540aab3 100644
--- a/generator/proto/__init__.py
+++ b/generator/proto/__init__.py
@@ -8,48 +8,60 @@ import shutil
 
 import pkg_resources
 
-from ._utils import has_grpcio_protoc, invoke_protoc
+# Compatibility layer to make TemporaryDirectory() available on Python 2.
 
-dirname = os.path.dirname(__file__)
-protosrc = os.path.join(dirname, "nanopb.proto")
-protodst = os.path.join(dirname, "nanopb_pb2.py")
+try:
+    from tempfile import TemporaryDirectory
+except ImportError:
+    class TemporaryDirectory:
+        '''TemporaryDirectory fallback for Python 2'''
+        def __enter__(self, prefix = 'tmp'):
+            self.dir = tempfile.mkdtemp(prefix = prefix)
+            return self.dir
+
+        def __exit__(self, *args):
+            shutil.rmtree(self.dir)
+
+# The code below tries two methods for importing the definitions from nanopb.proto:
+# 1) If nanopb.proto source file is available, build it using protoc to a temporary directory.
+# 2) Otherwise, try importing a precompiled file.
+
+from ._utils import has_grpcio_protoc, invoke_protoc
 
 # First, try to dynamically build and load the .proto file
-error_build = None
+build_error = None
 proto_ok = False
+dirname = os.path.dirname(__file__)
+protosrc = os.path.join(dirname, "nanopb.proto")
 if os.path.isfile(protosrc):
     tmpdir = tempfile.mkdtemp(prefix="nanopb-")
 
     try:
-        cmd = [
-            "protoc",
-            "--python_out={}".format(tmpdir),
-            protosrc,
-            "-I={}".format(dirname),
-        ]
+        with TemporaryDirectory(prefix = 'nanopb-') as tmpdir:
+            cmd = ["protoc",
+                   "--python_out={}".format(tmpdir),
+                   protosrc,
+                   "-I={}".format(dirname),
+                  ]
 
-        invoke_protoc(argv=cmd)
-
-        sys.path.insert(0, tmpdir)
-        import nanopb_pb2
-
-        proto_ok = True
+            invoke_protoc(argv=cmd)
+            sys.path.insert(0, tmpdir)
+            import nanopb_pb2
+            proto_ok = True
     except e:
-        error_build = e
-    finally:
-        shutil.rmtree(tmpdir)
+        build_error = e
 
 # As a fallback, see if we have a prebuilt file available
-error_import = None
+import_error = None
 if not proto_ok:
     try:
-        import nanopb_pb2
+        from . import nanopb_pb2
         proto_ok = True
+    except ImportError:
+        pass # To be expected
     except:
-        error_import = e
+        import_error = e
 
 if not proto_ok:
     sys.stderr.write("Failed to build/import nanopb_pb2.py.\n")
-    sys.stderr.write("\n\nBuild error: %s\n" % error_build)
-    sys.stderr.write("\n\nDirect import error: %s\n" % error_import)
-    raise Exception("Failed to build/import nanopb_pb2.py.\n")
+    raise build_error or import_error or Exception("Failed to build/import nanopb_pb2.py.\n")