feat: add copy_from method for field assignment (#215)

software-dov · web-flow · commit 11c3e58a9ba5 · 2021-03-16T12:14:34.000-07:00
Because of implementation detatils of the underlying protocol buffers
runtime, assigning to a proto-plus message field is achieved by
copying, not by updating references. This can lead to surprising and
unintuitive behavior for developers who are expecting python object
behavior, e.g. reference aliases.

This PR adds a 'Message.copy_from' method that is semantically
identical to regular assignment. This can be used at the discretion of
the developer to clarify expected behavior.
diff --git a/docs/messages.rst b/docs/messages.rst
@@ -80,6 +80,74 @@ Instantiate messages using either keyword arguments or a :class:`dict`
     >>> song.title
     'Canon in D'
 
+
+Assigning to Fields
+-------------------
+
+One of the goals of proto-plus is to make protobufs feel as much like regular python
+objects as possible. It is possible to update a message's field by assigning to it,
+just as if it were a regular python object.
+
+.. code-block:: python
+
+   song = Song()
+   song.composer = Composer(given_name="Johann", family_name="Bach")
+
+   # Can also assign from a dictionary as a convenience.
+   song.composer = {"given_name": "Claude", "family_name": "Debussy"}
+
+   # Repeated fields can also be assigned
+   class Album(proto.Message):
+       songs = proto.RepeatedField(Song, number=1)
+
+   a = Album()
+   songs = [Song(title="Canon in D"), Song(title="Little Fugue")]
+   a.songs = songs
+
+.. note::
+
+   Assigning to a proto-plus message field works by making copies, not by updating references.
+   This is necessary because of memory layout requirements of protocol buffers.
+   These memory constraints are maintained by the protocol buffers runtime.
+   This behavior can be surprising under certain circumstances, e.g. trying to save
+   an alias to a nested field.
+
+   :class:`proto.Message` defines a helper message, :meth:`~.Message.copy_from` to
+   help make the distinction clear when reading code.
+   The semantics of :meth:`~.Message.copy_from` are identical to the field assignment behavior described above.
+
+   .. code-block:: python
+
+      composer = Composer(given_name="Johann", family_name="Bach")
+      song = Song(title="Tocatta and Fugue in D Minor", composer=composer)
+      composer.given_name = "Wilhelm"
+
+      # 'composer' is NOT a reference to song.composer
+      assert song.composer.given_name == "Johann"
+
+      # We CAN update the song's composer by assignment.
+      song.composer = composer
+      composer.given_name = "Carl"
+
+      # 'composer' is STILL not a referene to song.composer.
+      assert song.composer.given_name == "Wilhelm"
+
+      # It does work in reverse, though,
+      # if we want a reference we can access then update.
+      composer = song.composer
+      composer.given_name = "Gottfried"
+
+      assert song.composer.given_name == "Gottfried"
+
+      # We can use 'copy_from' if we're concerned that the code
+      # implies that assignment involves references.
+      composer = Composer(given_name="Elisabeth", family_name="Bach")
+      # We could also do Message.copy_from(song.composer, composer) instead.
+      Composer.copy_from(song.composer, composer)
+
+      assert song.composer.given_name == "Elisabeth"
+
+
 Enums
 -----
 
diff --git a/docs/reference/message.rst b/docs/reference/message.rst
@@ -11,7 +11,7 @@ Message and Field
     .. automethod:: to_json
     .. automethod:: from_json
     .. automethod:: to_dict
-
+    .. automethod:: copy_from
 
 .. automodule:: proto.fields
   :members:
diff --git a/proto/message.py b/proto/message.py
@@ -401,6 +401,36 @@ def to_dict(
             use_integers_for_enums=use_integers_for_enums,
         )
 
+    def copy_from(cls, instance, other):
+        """Equivalent for protobuf.Message.CopyFrom
+
+        Args:
+            instance: An instance of this message type
+            other: (Union[dict, ~.Message):
+                A dictionary or message to reinitialize the values for this message.
+        """
+        if isinstance(other, cls):
+            # Just want the underlying proto.
+            other = Message.pb(other)
+        elif isinstance(other, cls.pb()):
+            # Don't need to do anything.
+            pass
+        elif isinstance(other, collections.abc.Mapping):
+            # Coerce into a proto
+            other = cls._meta.pb(**other)
+        else:
+            raise TypeError(
+                "invalid argument type to copy to {}: {}".format(
+                    cls.__name__, other.__class__.__name__
+                )
+            )
+
+        # Note: we can't just run self.__init__ because this may be a message field
+        # for a higher order proto; the memory layout for protos is NOT LIKE the
+        # python memory model. We cannot rely on just setting things by reference.
+        # Non-trivial complexity is (partially) hidden by the protobuf runtime.
+        cls.pb(instance).CopyFrom(other)
+
 
 class Message(metaclass=MessageMeta):
     """The abstract base class for a message.
@@ -436,7 +466,7 @@ def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
             #
             # The `wrap` method on the metaclass is the public API for taking
             # ownership of the passed in protobuf objet.
-            mapping = copy.copy(mapping)
+            mapping = copy.deepcopy(mapping)
             if kwargs:
                 mapping.MergeFrom(self._meta.pb(**kwargs))
 
diff --git a/tests/test_message.py b/tests/test_message.py
@@ -317,3 +317,27 @@ class Squid(proto.Message):
 
     s = Squid({"mass_kg": 20, "length_cm": 100}, ignore_unknown_fields=True)
     assert not hasattr(s, "length_cm")
+
+
+def test_copy_from():
+    class Mollusc(proto.Message):
+        class Squid(proto.Message):
+            mass_kg = proto.Field(proto.INT32, number=1)
+
+        squid = proto.Field(Squid, number=1)
+
+    m = Mollusc()
+    s = Mollusc.Squid(mass_kg=20)
+    Mollusc.Squid.copy_from(m.squid, s)
+    assert m.squid is not s
+    assert m.squid == s
+
+    s.mass_kg = 30
+    Mollusc.Squid.copy_from(m.squid, Mollusc.Squid.pb(s))
+    assert m.squid == s
+
+    Mollusc.Squid.copy_from(m.squid, {"mass_kg": 10})
+    assert m.squid.mass_kg == 10
+
+    with pytest.raises(TypeError):
+        Mollusc.Squid.copy_from(m.squid, (("mass_kg", 20)))
