[feat] Support google.protobuf.Struct (#10)

This commit makes the denizens of `google/protobuf/struct.proto`
be represented as native Python objects (primitives, sequences,
and maps).

Author: lukesneeringer

packages/proto-plus/proto/marshal/collections/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .map import MapComposite
+from .maps import MapComposite
 from .repeated import Repeated
 from .repeated import RepeatedComposite
 

…to-plus/proto/marshal/collections/map.py …o-plus/proto/marshal/collections/maps.pypackages/proto-plus/proto/marshal/collections/map.py renamed to packages/proto-plus/proto/marshal/collections/maps.py packages/proto-plus/proto/marshal/collections/map.py renamed to packages/proto-plus/proto/marshal/collections/maps.py

@@ -102,6 +102,11 @@ def _pb_type(self):
         canary = copy.deepcopy(self.pb).add()
         return type(canary)
 
+    def __eq__(self, other):
+        if super().__eq__(other):
+            return True
+        return tuple([i for i in self]) == tuple(other)
+
     def __getitem__(self, key):
         return self._marshal.to_python(self._pb_type, self.pb[key])
 

packages/proto-plus/proto/marshal/collections/repeated.py

@@ -17,13 +17,15 @@
 from google.protobuf import message
 from google.protobuf import duration_pb2
 from google.protobuf import timestamp_pb2
+from google.protobuf import struct_pb2
 from google.protobuf import wrappers_pb2
 
 from proto.marshal import compat
 from proto.marshal.collections import MapComposite
 from proto.marshal.collections import Repeated
 from proto.marshal.collections import RepeatedComposite
 from proto.marshal.rules import dates
+from proto.marshal.rules import struct
 from proto.marshal.rules import wrappers
 
 
@@ -130,6 +132,15 @@ def reset(self):
         self.register(wrappers_pb2.UInt32Value, wrappers.UInt32ValueRule())
         self.register(wrappers_pb2.UInt64Value, wrappers.UInt64ValueRule())
 
+        # Register the google.protobuf.Struct wrappers.
+        #
+        # These are aware of the marshal that created them, because they
+        # create RepeatedComposite and MapComposite instances directly and
+        # need to pass the marshal to them.
+        self.register(struct_pb2.Value, struct.ValueRule(marshal=self))
+        self.register(struct_pb2.ListValue, struct.ListValueRule(marshal=self))
+        self.register(struct_pb2.Struct, struct.StructRule(marshal=self))
+
     def to_python(self, proto_type, value, *, absent: bool = None):
         # Internal protobuf has its own special type for lists of values.
         # Return a view around it that implements MutableSequence.
@@ -147,14 +158,21 @@ def to_python(self, proto_type, value, *, absent: bool = None):
         return rule.to_python(value, absent=absent)
 
     def to_proto(self, proto_type, value, *, strict: bool = False):
-        # For our repeated and map view objects, simply return the
-        # underlying pb.
-        if isinstance(value, (Repeated, MapComposite)):
-            return value.pb
-
-        # Convert lists and tuples recursively.
-        if isinstance(value, (list, tuple)):
-            return type(value)([self.to_proto(proto_type, i) for i in value])
+        # The protos in google/protobuf/struct.proto are exceptional cases,
+        # because they can and should represent themselves as lists and dicts.
+        # These cases are handled in their rule classes.
+        if proto_type not in (struct_pb2.Value,
+                struct_pb2.ListValue, struct_pb2.Struct):
+            # For our repeated and map view objects, simply return the
+            # underlying pb.
+            if isinstance(value, (Repeated, MapComposite)):
+                return value.pb
+
+            # Convert lists and tuples recursively.
+            if isinstance(value, (list, tuple)):
+                return type(value)(
+                    [self.to_proto(proto_type, i) for i in value],
+                )
 
         # Convert dictionaries recursively when the proto type is a map.
         # This is slightly more complicated than converting a list or tuple

packages/proto-plus/proto/marshal/marshal.py

@@ -0,0 +1,130 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections.abc
+
+from google.protobuf import struct_pb2
+
+from proto.marshal.collections import maps
+from proto.marshal.collections import repeated
+
+
+class ValueRule:
+    """A rule to marshal between google.protobuf.Value and Python values."""
+
+    def __init__(self, *, marshal):
+        self._marshal = marshal
+
+    def to_python(self, value, *, absent: bool = None):
+        """Coerce the given value to the appropriate Python type.
+
+        Note that setting ``null_value`` is distinct from not setting
+        a value, and the absent value will raise an exception.
+        """
+        kind = value.WhichOneof('kind')
+        if kind == 'null_value':
+            return None
+        if kind == 'bool_value':
+            return bool(value.bool_value)
+        if kind == 'number_value':
+            return float(value.number_value)
+        if kind == 'string_value':
+            return str(value.string_value)
+        if kind == 'struct_value':
+            return self._marshal.to_python(
+                struct_pb2.Struct,
+                value.struct_value,
+                absent=False,
+            )
+        if kind == 'list_value':
+            return self._marshal.to_python(
+                struct_pb2.ListValue,
+                value.list_value,
+                absent=False,
+            )
+        raise AttributeError
+
+    def to_proto(self, value) -> struct_pb2.Value:
+        """Return a protobuf Value object representing this value."""
+        if isinstance(value, struct_pb2.Value):
+            return value
+        if value is None:
+            return struct_pb2.Value(null_value=0)
+        if isinstance(value, bool):
+            return struct_pb2.Value(bool_value=value)
+        if isinstance(value, (int, float)):
+            return struct_pb2.Value(number_value=float(value))
+        if isinstance(value, str):
+            return struct_pb2.Value(string_value=value)
+        if isinstance(value, collections.abc.Sequence):
+            return struct_pb2.Value(
+                list_value=self._marshal.to_proto(struct_pb2.ListValue, value),
+            )
+        if isinstance(value, collections.abc.Mapping):
+            return struct_pb2.Value(
+                struct_value=self._marshal.to_proto(struct_pb2.Struct, value),
+            )
+        raise ValueError('Unable to coerce value: %r' % value)
+
+
+class ListValueRule:
+    """A rule translating google.protobuf.ListValue and list-like objects."""
+
+    def __init__(self, *, marshal):
+        self._marshal = marshal
+
+    def to_python(self, value, *, absent: bool = None):
+        """Coerce the given value to a Python sequence."""
+        return repeated.RepeatedComposite(value.values, marshal=self._marshal)
+
+    def to_proto(self, value) -> struct_pb2.ListValue:
+        # We got a proto, or else something we sent originally.
+        # Preserve the instance we have.
+        if isinstance(value, struct_pb2.ListValue):
+            return value
+        if isinstance(value, repeated.RepeatedComposite):
+            return struct_pb2.ListValue(values=[v for v in value.pb])
+
+        # We got a list (or something list-like); convert it.
+        return struct_pb2.ListValue(values=[
+            self._marshal.to_proto(struct_pb2.Value, v) for v in value
+        ])
+
+
+class StructRule:
+    """A rule translating google.protobuf.Struct and dict-like objects."""
+
+    def __init__(self, *, marshal):
+        self._marshal = marshal
+
+    def to_python(self, value, *, absent: bool = None):
+        """Coerce the given value to a Python mapping."""
+        return maps.MapComposite(value.fields, marshal=self._marshal)
+
+    def to_proto(self, value) -> struct_pb2.Struct:
+        # We got a proto, or else something we sent originally.
+        # Preserve the instance we have.
+        if isinstance(value, struct_pb2.Struct):
+            return value
+        if isinstance(value, maps.MapComposite):
+            return struct_pb2.Struct(
+                fields={k: v for k, v in value.pb.items()},
+            )
+
+        # We got a dict (or something dict-like); convert it.
+        answer = struct_pb2.Struct(fields={
+            k: self._marshal.to_proto(struct_pb2.Value, v)
+            for k, v in value.items()
+        })
+        return answer

packages/proto-plus/proto/marshal/rules/struct.py

@@ -489,7 +489,7 @@ def __setattr__(self, key, value):
         pb_type = self._meta.fields[key].pb_type
         pb_value = marshal.to_proto(pb_type, value)
 
-        # We *always* clear the existing field.
+        # Clear the existing field.
         # This is the only way to successfully write nested falsy values,
         # because otherwise MergeFrom will no-op on them.
         self._pb.ClearField(key)

packages/proto-plus/proto/message.py

@@ -57,14 +57,10 @@ def pytest_runtest_setup(item):
             module = getattr(item.module, name)
             pool.AddSerializedFile(module.DESCRIPTOR.serialized_pb)
             fd = pool.FindFileByName(module.DESCRIPTOR.name)
-            for message_name, descriptor in fd.message_types_by_name.items():
-                new_message = reflection.GeneratedProtocolMessageType(
-                    message_name,
-                    (message.Message,),
-                    {'DESCRIPTOR': descriptor, '__module__': None},
-                )
-                sym_db.RegisterMessage(new_message)
-                setattr(module, message_name, new_message)
+
+            # Register all the messages to the symbol database and the
+            # module. Do this recursively if there are nested messages.
+            _register_messages(module, fd.message_types_by_name, sym_db)
 
             # Track which modules had new message classes loaded.
             # This is used below to wire the new classes into the marshal.
@@ -74,10 +70,25 @@ def pytest_runtest_setup(item):
     # then reload the appropriate modules so the marshal is using the new ones.
     if 'wrappers_pb2' in reloaded:
         imp.reload(rules.wrappers)
+    if 'struct_pb2' in reloaded:
+        imp.reload(rules.struct)
     if reloaded.intersection({'timestamp_pb2', 'duration_pb2'}):
         imp.reload(rules.dates)
 
 
 def pytest_runtest_teardown(item):
     Marshal._instances.clear()
     [i.stop() for i in item._mocks]
+
+
+def _register_messages(scope, iterable, sym_db):
+    """Create and register messages from the file descriptor."""
+    for name, descriptor in iterable.items():
+        new_msg = reflection.GeneratedProtocolMessageType(
+            name,
+            (message.Message,),
+            {'DESCRIPTOR': descriptor, '__module__': None},
+        )
+        sym_db.RegisterMessage(new_msg)
+        setattr(scope, name, new_msg)
+        _register_messages(new_msg, descriptor.nested_types_by_name, sym_db)

packages/proto-plus/tests/conftest.py

@@ -31,9 +31,21 @@ class Baz(proto.Message):
 
     baz = Baz(foos=[Foo(bar=42)])
     assert len(baz.foos) == 1
+    assert baz.foos == baz.foos
     assert baz.foos[0].bar == 42
 
 
+def test_repeated_composite_equality():
+    class Foo(proto.Message):
+        bar = proto.Field(proto.INT32, number=1)
+
+    class Baz(proto.Message):
+        foos = proto.RepeatedField(proto.MESSAGE, message=Foo, number=1)
+
+    baz = Baz(foos=[Foo(bar=42)])
+    assert baz.foos == baz.foos
+
+
 def test_repeated_composite_init_struct():
     class Foo(proto.Message):
         bar = proto.Field(proto.INT32, number=1)