1'''This file dynamically builds the proto definitions for Python.'''
2from __future__ import absolute_import
3
4import os
5import os.path
6import sys
7import tempfile
8import shutil
9import traceback
10import pkg_resources
11from ._utils import has_grpcio_protoc, invoke_protoc, print_versions
12
13# Compatibility layer to make TemporaryDirectory() available on Python 2.
14try:
15    from tempfile import TemporaryDirectory
16except ImportError:
17    class TemporaryDirectory:
18        '''TemporaryDirectory fallback for Python 2'''
19        def __init__(self, prefix = 'tmp', dir = None):
20            self.prefix = prefix
21            self.dir = dir
22
23        def __enter__(self):
24            self.dir = tempfile.mkdtemp(prefix = self.prefix, dir = self.dir)
25            return self.dir
26
27        def __exit__(self, *args):
28            shutil.rmtree(self.dir)
29
30def build_nanopb_proto(protosrc, dirname):
31    '''Try to build a .proto file for python-protobuf.
32    Returns True if successful.
33    '''
34
35    cmd = [
36        "protoc",
37        "--python_out={}".format(dirname),
38        protosrc,
39        "-I={}".format(dirname),
40    ]
41
42    if has_grpcio_protoc():
43        # grpcio-tools has an extra CLI argument
44        # from grpc.tools.protoc __main__ invocation.
45        _builtin_proto_include = pkg_resources.resource_filename('grpc_tools', '_proto')
46        cmd.append("-I={}".format(_builtin_proto_include))
47
48    try:
49        invoke_protoc(argv=cmd)
50    except:
51        sys.stderr.write("Failed to build nanopb_pb2.py: " + ' '.join(cmd) + "\n")
52        sys.stderr.write(traceback.format_exc() + "\n")
53        return False
54
55    return True
56
57def load_nanopb_pb2():
58    # To work, the generator needs python-protobuf built version of nanopb.proto.
59    # There are three methods to provide this:
60    #
61    # 1) Load a previously generated generator/proto/nanopb_pb2.py
62    # 2) Use protoc to build it and store it permanently generator/proto/nanopb_pb2.py
63    # 3) Use protoc to build it, but store only temporarily in system-wide temp folder
64    #
65    # By default these are tried in numeric order.
66    # If NANOPB_PB2_TEMP_DIR environment variable is defined, the 2) is skipped.
67    # If the value of the $NANOPB_PB2_TEMP_DIR exists as a directory, it is used instead
68    # of system temp folder.
69
70    build_error = None
71    proto_ok = False
72    tmpdir = os.getenv("NANOPB_PB2_TEMP_DIR")
73    temporary_only = (tmpdir is not None)
74    dirname = os.path.dirname(__file__)
75    protosrc = os.path.join(dirname, "nanopb.proto")
76    protodst = os.path.join(dirname, "nanopb_pb2.py")
77    proto_ok = False
78
79    if tmpdir is not None and not os.path.isdir(tmpdir):
80        tmpdir = None # Use system-wide temp dir
81
82    if os.path.isfile(protosrc):
83        src_date = os.path.getmtime(protosrc)
84        if not os.path.isfile(protodst) or os.path.getmtime(protodst) < src_date:
85            # Outdated, rebuild
86            proto_ok = False
87        else:
88            try:
89                from . import nanopb_pb2 as nanopb_pb2_mod
90                proto_ok = True
91            except Exception as e:
92                sys.stderr.write("Failed to import nanopb_pb2.py: " + str(e) + "\n"
93                                "Will automatically attempt to rebuild this.\n"
94                                "Verify that python-protobuf and protoc versions match.\n")
95                print_versions()
96
97    # Try to rebuild into generator/proto directory
98    if not proto_ok and not temporary_only:
99        proto_ok = build_nanopb_proto(protosrc, dirname)
100
101        try:
102            from . import nanopb_pb2 as nanopb_pb2_mod
103        except:
104            sys.stderr.write("Failed to import generator/proto/nanopb_pb2.py:\n")
105            sys.stderr.write(traceback.format_exc() + "\n")
106
107    # Try to rebuild into temporary directory
108    if not proto_ok:
109        with TemporaryDirectory(prefix = 'nanopb-', dir = tmpdir) as protodir:
110            proto_ok = build_nanopb_proto(protosrc, protodir)
111
112            if protodir not in sys.path:
113                sys.path.insert(0, protodir)
114
115            try:
116                import nanopb_pb2 as nanopb_pb2_mod
117            except:
118                sys.stderr.write("Failed to import %s/nanopb_pb2.py:\n" % protodir)
119                sys.stderr.write(traceback.format_exc() + "\n")
120
121    # If everything fails
122    if not proto_ok:
123        sys.stderr.write("\n\nGenerating nanopb_pb2.py failed.\n")
124        sys.stderr.write("Make sure that a protoc generator is available and matches python-protobuf version.\n")
125        print_versions()
126        sys.exit(1)
127
128    return nanopb_pb2_mod
129