pw_protobuf_compiler: Iterate over messages
Support iterating over all message types in a proto Library.
Change-Id: Ia95e30842c4ea8a9795762ebb3cf9a7a27597207
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/32881
Reviewed-by: Alexei Frolov <frolv@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: 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 0c4a8cc..9b6f6cd 100644
--- a/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py
+++ b/pw_protobuf_compiler/py/pw_protobuf_compiler/python_protos.py
@@ -308,10 +308,23 @@
for m in modules)
def modules(self) -> Iterable:
- """Allows iterating over all protobuf modules in this library."""
+ """Iterates over all protobuf modules in this library."""
for module_list in self.modules_by_package.values():
yield from module_list
+ def messages(self) -> Iterable:
+ """Iterates over all protobuf messages in this library."""
+ for module in self.modules():
+ yield from _nested_messages(
+ module, module.DESCRIPTOR.message_types_by_name)
+
+
+def _nested_messages(scope, message_names: Iterable[str]) -> Iterator:
+ for name in message_names:
+ msg = getattr(scope, name)
+ yield msg
+ yield from _nested_messages(msg, msg.DESCRIPTOR.nested_types_by_name)
+
def _repr_char(char: int) -> str:
r"""Returns an ASCII char or the \x code for non-printable values."""
diff --git a/pw_protobuf_compiler/py/python_protos_test.py b/pw_protobuf_compiler/py/python_protos_test.py
index 82ed547..7888c55 100755
--- a/pw_protobuf_compiler/py/python_protos_test.py
+++ b/pw_protobuf_compiler/py/python_protos_test.py
@@ -84,6 +84,18 @@
repeated int64 value = 1;
Greeting hi = 2;
}
+
+message NestingMessage {
+ message NestedMessage {
+ message NestedNestedMessage {
+ int32 nested_nested_field = 1;
+ }
+
+ NestedNestedMessage nested_nested_message = 1;
+ }
+
+ NestedMessage nested_message = 1;
+}
"""
@@ -201,6 +213,20 @@
with self.assertRaises(KeyError):
_ = self._library.packages.pw.protobuf_compiler['not here']
+ def test_messages(self):
+ protos = self._library.packages.pw.protobuf_compiler
+ self.assertEqual(
+ set(self._library.messages()), {
+ protos.test1.SomeMessage,
+ protos.test1.AnotherMessage,
+ protos.test2.Request,
+ protos.test2.Response,
+ protos.test2.Hello,
+ protos.test2.NestingMessage,
+ protos.test2.NestingMessage.NestedMessage,
+ protos.test2.NestingMessage.NestedMessage.NestedNestedMessage,
+ })
+
PROTO_FOR_REPR = """\
syntax = "proto3";
