pw_protobuf_compiler: Simplify proto imports
Simplify importing protos from strings, which is a common use case in
tests.
Change-Id: Ib4b32f5e2d18009ec91fe2e1c30c1627dc844c9a
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/14260
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Wyatt Hepler <hepler@google.com>
diff --git a/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py b/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py
index 4c55bd3..9499520 100644
--- a/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py
+++ b/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py
@@ -13,6 +13,7 @@
# the License.
"""Tools for compiling and importing Python protos on the fly."""
+import importlib.util
import logging
import os
from pathlib import Path
@@ -20,8 +21,7 @@
import shlex
import tempfile
from types import ModuleType
-from typing import Dict, Iterable, List, Set, Tuple, Union
-import importlib.util
+from typing import Dict, Iterable, Iterator, List, Set, Tuple, Union
_LOG = logging.getLogger(__name__)
@@ -68,7 +68,7 @@
return module
-def import_modules(directory: PathOrStr) -> Iterable[ModuleType]:
+def import_modules(directory: PathOrStr) -> Iterator[ModuleType]:
"""Imports modules in a directory and yields them."""
parent = os.path.dirname(directory)
@@ -85,7 +85,7 @@
def compile_and_import(proto_files: Iterable[PathOrStr],
includes: Iterable[PathOrStr] = (),
- output_dir: PathOrStr = None) -> Iterable[ModuleType]:
+ output_dir: PathOrStr = None) -> Iterator[ModuleType]:
"""Compiles protos and imports their modules; yields the proto modules.
Args:
@@ -102,7 +102,7 @@
compile_protos(output_dir, proto_files, includes)
yield from import_modules(output_dir)
else:
- with tempfile.TemporaryDirectory(prefix='protos_') as tempdir:
+ with tempfile.TemporaryDirectory(prefix='compiled_protos_') as tempdir:
compile_protos(tempdir, proto_files, includes)
yield from import_modules(tempdir)
@@ -114,6 +114,28 @@
return next(iter(compile_and_import([proto_file], includes, output_dir)))
+def compile_and_import_strings(
+ contents: Iterable[str],
+ includes: Iterable[PathOrStr] = (),
+ output_dir: PathOrStr = None) -> Iterator[ModuleType]:
+ """Compiles protos in one or more strings."""
+
+ if isinstance(contents, str):
+ contents = [contents]
+
+ with tempfile.TemporaryDirectory(prefix='proto_sources_') as path:
+ protos = []
+
+ for proto in contents:
+ # Use a hash of the proto so the same contents map to the same file
+ # name. The protobuf package complains if it seems the same contents
+ # in files with different names.
+ protos.append(Path(path, f'protobuf_{hash(proto):x}.proto'))
+ protos[-1].write_text(proto)
+
+ yield from compile_and_import(protos, includes, output_dir)
+
+
class _ProtoPackage:
"""Used by the Library class for accessing protocol buffer modules."""
def __init__(self, package: str):
@@ -152,6 +174,13 @@
the list of modules in a particular package, and the modules() generator
for iterating over all modules.
"""
+ @classmethod
+ def from_strings(cls,
+ contents: Iterable[str],
+ includes: Iterable[PathOrStr] = (),
+ output_dir: PathOrStr = None) -> 'Library':
+ return cls(compile_and_import_strings(contents, includes, output_dir))
+
def __init__(self, modules: Iterable[ModuleType]):
"""Constructs a Library from an iterable of modules.
diff --git a/pw_protobuf_compiler/py/python_protos_test.py b/pw_protobuf_compiler/py/python_protos_test.py
index 2408fb3..f7dd75a 100755
--- a/pw_protobuf_compiler/py/python_protos_test.py
+++ b/pw_protobuf_compiler/py/python_protos_test.py
@@ -20,7 +20,7 @@
from pw_protobuf_compiler import python_protos
-PROTO_1 = b"""\
+PROTO_1 = """\
syntax = "proto3";
package pw.protobuf_compiler.test1;
@@ -48,7 +48,7 @@
}
"""
-PROTO_2 = b"""\
+PROTO_2 = """\
syntax = "proto2";
package pw.protobuf_compiler.test2;
@@ -69,7 +69,7 @@
}
"""
-PROTO_3 = b"""\
+PROTO_3 = """\
syntax = "proto3";
package pw.protobuf_compiler.test2;
@@ -93,7 +93,7 @@
for i, contents in enumerate([PROTO_1, PROTO_2, PROTO_3], 1):
self._protos.append(Path(self._proto_dir.name, f'test_{i}.proto'))
- self._protos[-1].write_bytes(contents)
+ self._protos[-1].write_text(contents)
def tearDown(self):
self._proto_dir.cleanup()
@@ -158,6 +158,24 @@
with self.assertRaises(KeyError):
_ = self._library.modules_by_package['pw.not_real']
+ def test_library_from_strings(self):
+ # Replace the package to avoid conflicts with the other proto imports
+ new_protos = [
+ p.replace('pw.protobuf_compiler', 'proto.library.test')
+ for p in [PROTO_1, PROTO_2, PROTO_3]
+ ]
+
+ library = python_protos.Library.from_strings(new_protos)
+
+ # Make sure we can safely import the same proto contents multiple times.
+ library = python_protos.Library.from_strings(new_protos)
+
+ msg = library.packages.proto.library.test.test2.Request
+ self.assertEqual(msg(magic_number=50).magic_number, 50)
+
+ val = library.packages.proto.library.test.test2.YO
+ self.assertEqual(val, 0)
+
if __name__ == '__main__':
unittest.main()
