Repository: OpenMined/PySyft
Pull request: Update `Role.register_action()` to allow registering `CommunicationActions`
## Description

This mostly standardizes the fields with `ComputationAction`, which makes it possible to register both action types interchangeably in `Roles`.

Depends on OpenMined/syft-proto#100

Resolves #3424

## Type of change

Please mark the options that are relevant.

- [ ] Added/Modified tutorials
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)

## Checklist:

* [X] I have added tests for my changes
* [ ] I did follow the [contribution guidelines](https://github.com/OpenMined/PySyft/blob/master/CONTRIBUTING.md)
* [ ] I have commented my code following [Google style](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html).


Problem statement:
Register `CommunicationActions` in `Roles` (for `Protocols`)
**Is your feature request related to a problem? Please describe.**
`Role.register_action()` should be able to register both `ComputationActions` and `CommunicationActions`. Some of the assumptions about which fields are present only make sense for `ComputationActions`, which causes errors when trying to register `CommunicationActions`.

**Describe the solution you'd like**
Make the change easy:
* Add tests for `Role.register_action()` with `ComputationActions` to guard the existing functionality against regressions.
* Move fields shared by both action types into the `Action` base class. These fields can be provided as standard `args` to `register_action()`.
* Provide fields that are unique to `ComputationActions` as `kwargs` instead of `args`

Make the easy change:
* Add tests for `Role.register_action()` with `CommunicationActions` to drive the new functionality
* Provide fields that are unique to `CommunicationActions` as `kwargs`
* Use the `action_type` param to decide which `kwargs` are expected to be present
* Refactor any private methods used by `register_action()` (like `_store_placeholders`) to work with both `Action` types

**Describe alternatives you've considered**
We could have separate methods to register the two types of `Action`. That would push the decision about which `Action` type to register out into the tracing code in the hooks, tracing wrappers, and `Placeholders` though. Might be better to make that decision in one centralized place to avoid duplication and discrepancies.

**Additional context**
This is an early step toward tracing `CommunicationActions` into `Protocols`.

Hints from issue discussion:

Commit message: Remove unnecessary/untested support for lists/dicts as `Action` args

Patch to review:
diff --git a/pip-dep/requirements.txt b/pip-dep/requirements.txt
index 60a812ac690..3b25958a32c 100644
--- a/pip-dep/requirements.txt
+++ b/pip-dep/requirements.txt
@@ -7,7 +7,8 @@ phe~=1.4.0
 Pillow~=6.2.2
 requests~=2.22.0
 scipy~=1.4.1
-syft-proto~=0.4.1
+# syft-proto~=0.4.1
+git+git://github.com/karlhigley/syft-proto@refactor/register-comms-actions#egg=syft-proto
 tblib~=1.6.0
 torchvision~=0.5.0
 torch~=1.4.0
diff --git a/syft/execution/action.py b/syft/execution/action.py
index 626737b28ca..9eb18df4a63 100644
--- a/syft/execution/action.py
+++ b/syft/execution/action.py
@@ -1,4 +1,14 @@
 from abc import ABC
+from abc import abstractmethod
+
+import syft as sy
+from syft.execution.placeholder import PlaceHolder
+from syft.execution.placeholder_id import PlaceholderId
+from syft.workers.abstract import AbstractWorker
+
+from syft_proto.execution.v1.communication_action_pb2 import (
+    CommunicationAction as CommunicationActionPB,
+)
 
 
 class Action(ABC):
@@ -10,4 +20,258 @@ class Action(ABC):
     tensors and add them together is an Action. Sending an object from one worker to another is
     also an Action."""
 
-    pass
+    def __init__(
+        self, name: str, target, args_: tuple, kwargs_: dict, return_ids: tuple, return_value=False
+    ):
+        """Initialize an action
+
+        Args:
+            name (String): The name of the method to be invoked (e.g. "__add__")
+            target (Tensor): The object to invoke the method on
+            args_ (Tuple): The arguments to the method call
+            kwargs_ (Dictionary): The keyword arguments to the method call
+            return_ids (Tuple): primarily for our async infrastructure (Plan, Protocol, etc.), the id of
+                action results are set by the client. This allows the client to be able to predict where
+                the results will be ahead of time. Importantly, this allows the client to pre-initalize the
+                pointers to the future data, regardless of whether the action has yet executed. It also
+                reduces the size of the response from the action (which is very often empty).
+            return_value (boolean): return the result or not. If true, the result is directly returned,
+                if not, the command sender will create a pointer to the remote result using the return_ids
+                and will need to do .get() later to get the result.
+
+        """
+
+        # call the parent constructor - setting the type integer correctly
+        super().__init__()
+
+        self.name = name
+        self.target = target
+        self.args = args_
+        self.kwargs = kwargs_
+        self.return_ids = return_ids
+        self.return_value = return_value
+
+        self._type_check("name", str)
+        self._type_check("args", tuple)
+        self._type_check("kwargs", dict)
+        self._type_check("return_ids", tuple)
+
+    def __eq__(self, other):
+        return (
+            self.name == other.name
+            and self.target == other.target
+            and self.args == other.args
+            and self.kwargs == other.kwargs
+            and self.return_ids == other.return_ids
+        )
+
+    def code(self, var_names=None) -> str:
+        """Returns pseudo-code representation of computation action"""
+
+        def stringify(obj):
+            if isinstance(obj, PlaceholderId):
+                id = obj.value
+                if var_names is None:
+                    ret = f"var_{id}"
+                else:
+                    if id in var_names:
+                        ret = var_names[id]
+                    else:
+                        idx = sum("var_" in k for k in var_names.values())
+                        name = f"var_{idx}"
+                        var_names[id] = name
+                        ret = name
+            elif isinstance(obj, PlaceHolder):
+                ret = stringify(obj.id)
+            elif isinstance(obj, (tuple, list)):
+                ret = ", ".join(stringify(o) for o in obj)
+            else:
+                ret = str(obj)
+
+            return ret
+
+        out = ""
+        if self.return_ids is not None:
+            out += stringify(self.return_ids) + " = "
+        if self.target is not None:
+            out += stringify(self.target) + "."
+        out += self.name + "("
+        out += stringify(self.args)
+        if self.kwargs:
+            if len(self.args) > 0:
+                out += ", "
+            out += ", ".join(f"{k}={w}" for k, w in self.kwargs.items())
+        out += ")"
+
+        return out
+
+    def __str__(self) -> str:
+        """Returns string representation of this action"""
+        return f"{type(self).__name__}[{self.code()}]"
+
+    def _type_check(self, field_name, expected_type):
+        actual_value = getattr(self, field_name)
+        assert actual_value is None or isinstance(
+            actual_value, expected_type
+        ), f"{field_name} must be {expected_type.__name__}, but was {type(actual_value).__name__}: {actual_value}."
+
+    # These methods must be implemented by child classes in order to return the correct type
+    # and to be detected by the serdes as serializable. They are therefore marked as abstract
+    # methods even though implementations are provided. Child classes may delegate to these
+    # implementations, but must implement their own conversions to the appropriate classes.
+
+    @staticmethod
+    @abstractmethod
+    def simplify(worker: AbstractWorker, action: "Action") -> tuple:
+        """
+        This function takes the attributes of a CommunicationAction and saves them in a tuple
+        Args:
+            worker (AbstractWorker): a reference to the worker doing the serialization
+            action (CommunicationAction): a CommunicationAction
+        Returns:
+            tuple: a tuple holding the unique attributes of the CommunicationAction
+        Examples:
+            data = simplify(worker, action)
+        """
+        return (
+            sy.serde.msgpack.serde._simplify(worker, action.name),
+            sy.serde.msgpack.serde._simplify(worker, action.target),
+            sy.serde.msgpack.serde._simplify(worker, action.args),
+            sy.serde.msgpack.serde._simplify(worker, action.kwargs),
+            sy.serde.msgpack.serde._simplify(worker, action.return_ids),
+            sy.serde.msgpack.serde._simplify(worker, action.return_value),
+        )
+
+    @staticmethod
+    @abstractmethod
+    def detail(worker: AbstractWorker, action_tuple: tuple) -> "Action":
+        """
+        This function takes the simplified tuple version of this message and converts
+        it into a CommunicationAction. The simplify() method runs the inverse of this method.
+
+        Args:
+            worker (AbstractWorker): a reference to the worker necessary for detailing. Read
+                syft/serde/serde.py for more information on why this is necessary.
+            communication_tuple (Tuple): the raw information being detailed.
+        Returns:
+            communication (CommunicationAction): a CommunicationAction.
+        Examples:
+            communication = detail(sy.local_worker, communication_tuple)
+        """
+        name, target, args_, kwargs_, return_ids, return_value = action_tuple
+
+        return (
+            sy.serde.msgpack.serde._detail(worker, name),
+            sy.serde.msgpack.serde._detail(worker, target),
+            sy.serde.msgpack.serde._detail(worker, args_),
+            sy.serde.msgpack.serde._detail(worker, kwargs_),
+            sy.serde.msgpack.serde._detail(worker, return_ids),
+            sy.serde.msgpack.serde._detail(worker, return_value),
+        )
+
+    @staticmethod
+    @abstractmethod
+    def bufferize(worker: AbstractWorker, action: "Action", protobuf_action):
+        """
+        This function takes the attributes of a Action and saves them in Protobuf
+        Args:
+            worker (AbstractWorker): a reference to the worker doing the serialization
+            action (Action): an Action
+        Returns:
+            protobuf_obj: a Protobuf message holding the unique attributes of the message
+        Examples:
+            data = bufferize(message)
+        """
+        protobuf_action.command = action.name
+
+        protobuf_target = None
+        if isinstance(action.target, sy.generic.pointers.pointer_tensor.PointerTensor):
+            protobuf_target = protobuf_action.target_pointer
+        elif isinstance(action.target, sy.execution.placeholder_id.PlaceholderId):
+            protobuf_target = protobuf_action.target_placeholder_id
+        elif isinstance(action.target, (int, str)):
+            sy.serde.protobuf.proto.set_protobuf_id(protobuf_action.target_id, action.target)
+        elif action.target is not None:
+            protobuf_target = protobuf_action.target_tensor
+
+        if protobuf_target is not None:
+            protobuf_target.CopyFrom(sy.serde.protobuf.serde._bufferize(worker, action.target))
+
+        if action.args:
+            protobuf_action.args.extend(sy.serde.protobuf.serde.bufferize_args(worker, action.args))
+
+        if action.kwargs:
+            for key, value in action.kwargs.items():
+                protobuf_action.kwargs.get_or_create(key).CopyFrom(
+                    sy.serde.protobuf.serde.bufferize_arg(worker, value)
+                )
+
+        if action.return_ids is not None:
+            if not isinstance(action.return_ids, (list, tuple)):
+                return_ids = (action.return_ids,)
+            else:
+                return_ids = action.return_ids
+
+            for return_id in return_ids:
+                if isinstance(return_id, PlaceholderId):
+                    # NOTE to know when we have a PlaceholderId, we store it
+                    # in return_placeholder_ids and not in return_ids
+                    protobuf_action.return_placeholder_ids.append(
+                        sy.serde.protobuf.serde._bufferize(worker, return_id)
+                    )
+                else:
+                    sy.serde.protobuf.proto.set_protobuf_id(
+                        protobuf_action.return_ids.add(), return_id
+                    )
+
+        return protobuf_action
+
+    @staticmethod
+    @abstractmethod
+    def unbufferize(worker: AbstractWorker, protobuf_obj):
+        """
+        This function takes the Protobuf version of this message and converts
+        it into an Action. The bufferize() method runs the inverse of this method.
+
+        Args:
+            worker (AbstractWorker): a reference to the worker necessary for detailing. Read
+                syft/serde/serde.py for more information on why this is necessary.
+            protobuf_obj (ActionPB): the Protobuf message
+
+        Returns:
+            obj (tuple): a tuple of the args required to instantiate an Action object
+
+        Examples:
+            message = unbufferize(sy.local_worker, protobuf_msg)
+        """
+        command = protobuf_obj.command
+        protobuf_target = protobuf_obj.WhichOneof("target")
+        if protobuf_target:
+            target = sy.serde.protobuf.serde._unbufferize(
+                worker, getattr(protobuf_obj, protobuf_obj.WhichOneof("target"))
+            )
+        else:
+            target = None
+        args_ = sy.serde.protobuf.serde.unbufferize_args(worker, protobuf_obj.args)
+
+        kwargs_ = {}
+        for key in protobuf_obj.kwargs:
+            kwargs_[key] = sy.serde.protobuf.serde.unbufferize_arg(worker, protobuf_obj.kwargs[key])
+
+        return_ids = tuple(
+            [sy.serde.protobuf.proto.get_protobuf_id(pb_id) for pb_id in protobuf_obj.return_ids]
+        )
+
+        return_placeholder_ids = tuple(
+            [
+                sy.serde.protobuf.serde._unbufferize(worker, placeholder)
+                for placeholder in protobuf_obj.return_placeholder_ids
+            ]
+        )
+
+        if return_placeholder_ids:
+            action = (command, target, args_, kwargs_, return_placeholder_ids)
+        else:
+            action = (command, target, args_, kwargs_, return_ids)
+
+        return action
diff --git a/syft/execution/communication.py b/syft/execution/communication.py
index 4262dd1c6ea..7d6c3157d64 100644
--- a/syft/execution/communication.py
+++ b/syft/execution/communication.py
@@ -9,66 +9,62 @@
 )
 
 
+COMMUNICATION_METHODS = [
+    "get",
+    "mid_get",
+    "move",
+    "remote_get",
+    "remote_send",
+    "send",
+    "share",
+    "share_",
+]
+
+
 class CommunicationAction(Action):
     """Describes communication actions performed on tensors"""
 
-    def __init__(
-        self,
-        obj_id: Union[str, int],  # id of target object (almost always) pointer_tensor
-        name: str,  # the pointer_tensor method # @TODO: reactor to enum
-        source: Union[str, int],  # id of the sender (worker)
-        destinations: List[Union[str, int]],  # id of the receiver(s) (worker(s))
-        kwargs_: dict,  # key word args needed for the pointer tensor method == self.name
-    ):
-        """Initialize an communication action
+    def __init__(self, name, target, args_, kwargs_, return_ids, return_value=False):
+        """Initialize an action
 
         Args:
-        """
-        super().__init__()
+            name (String): The name of the method to be invoked (e.g. "send")
+            target (Tensor): The object to invoke the method on
+            args_ (Tuple): The arguments to the method call
+            kwargs_ (Dictionary): The keyword arguments to the method call
+            return_ids (Tuple): primarily for our async infrastructure (Plan, Protocol, etc.), the id of
+                action results are set by the client. This allows the client to be able to predict where
+                the results will be ahead of time. Importantly, this allows the client to pre-initalize the
+                pointers to the future data, regardless of whether the action has yet executed. It also
+                reduces the size of the response from the action (which is very often empty).
+            return_value (boolean): return the result or not. If true, the result is directly returned,
+                if not, the command sender will create a pointer to the remote result using the return_ids
+                and will need to do .get() later to get the result.
 
-        self.obj_id = obj_id
-        if name in ["move", "remote_send", "mid_get", "remote_get", "get", "share", "share_"]:
-            #  float_prec, fix_prec => should be computation actions (they modify tensors)?
-            self.name = name
-        else:
+        """
+        if name not in COMMUNICATION_METHODS:
             raise ValueError(
-                f"name `{name}` for CommunicationAction is not in the list of supported actions"
+                f"Method `{name}` is not supported by CommunicationActions. Consider using ComputationAction instead."
             )
-        self.source = source
-        self.destinations = destinations
-        self.kwargs = kwargs_
-
-    def __eq__(self, other):
-        return (
-            self.obj_id == other.obj_id
-            and self.name == other.name
-            and self.source == other.source
-            and self.destinations == other.destinations
-            and self.kwargs == other.kwargs
-        )
+
+        super().__init__(name, target, args_, kwargs_, return_ids, return_value=return_value)
 
     @staticmethod
-    def simplify(worker: AbstractWorker, communication: "CommunicationAction") -> tuple:
+    def simplify(worker: AbstractWorker, action: "Action") -> tuple:
         """
         This function takes the attributes of a CommunicationAction and saves them in a tuple
         Args:
             worker (AbstractWorker): a reference to the worker doing the serialization
-            communication (CommunicationAction): a CommunicationAction
+            action (CommunicationAction): a CommunicationAction
         Returns:
             tuple: a tuple holding the unique attributes of the CommunicationAction
         Examples:
-            data = simplify(worker, communication)
+            data = simplify(worker, action)
         """
-        return (
-            sy.serde.msgpack.serde._simplify(worker, communication.obj_id),
-            sy.serde.msgpack.serde._simplify(worker, communication.name),
-            sy.serde.msgpack.serde._simplify(worker, communication.source),
-            sy.serde.msgpack.serde._simplify(worker, communication.destinations),
-            sy.serde.msgpack.serde._simplify(worker, communication.kwargs),
-        )
+        return Action.simplify(worker, action)
 
     @staticmethod
-    def detail(worker: AbstractWorker, communication_tuple: tuple) -> "CommunicationAction":
+    def detail(worker: AbstractWorker, action_tuple: tuple) -> "Action":
         """
         This function takes the simplified tuple version of this message and converts
         it into a CommunicationAction. The simplify() method runs the inverse of this method.
@@ -82,18 +78,9 @@ def detail(worker: AbstractWorker, communication_tuple: tuple) -> "Communication
         Examples:
             communication = detail(sy.local_worker, communication_tuple)
         """
+        attrs = Action.detail(worker, action_tuple)
 
-        (obj_id, name, source, destinations, kwargs_) = communication_tuple
-
-        detailed_obj = sy.serde.msgpack.serde._detail(worker, obj_id)
-        detailed_name = sy.serde.msgpack.serde._detail(worker, name)
-        detailed_source = sy.serde.msgpack.serde._detail(worker, source)
-        detailed_destinations = sy.serde.msgpack.serde._detail(worker, destinations)
-        detailed_kwargs = sy.serde.msgpack.serde._detail(worker, kwargs_)
-
-        return CommunicationAction(
-            detailed_obj, detailed_name, detailed_source, detailed_destinations, detailed_kwargs
-        )
+        return CommunicationAction(*attrs)
 
     @staticmethod
     def bufferize(
@@ -109,22 +96,9 @@ def bufferize(
         Examples:
             data = bufferize(sy.local_worker, communication)
         """
-        protobuf_obj = CommunicationActionPB()
-        protobuf_obj.name = communication.name
-
-        sy.serde.protobuf.proto.set_protobuf_id(protobuf_obj.obj_id, communication.obj_id)
-        sy.serde.protobuf.proto.set_protobuf_id(protobuf_obj.source, communication.source)
-
-        for destination in communication.destinations:
-            sy.serde.protobuf.proto.set_protobuf_id(protobuf_obj.destinations.add(), destination)
+        protobuf_action = CommunicationActionPB()
 
-        if communication.kwargs:
-            for key, value in communication.kwargs.items():
-                protobuf_obj.kwargs.get_or_create(key).CopyFrom(
-                    sy.serde.protobuf.serde.bufferize_arg(worker, value)
-                )
-
-        return protobuf_obj
+        return Action.bufferize(worker, communication, protobuf_action)
 
     @staticmethod
     def unbufferize(
@@ -132,7 +106,7 @@ def unbufferize(
     ) -> "CommunicationAction":
         """
         This function takes the Protobuf version of this message and converts
-        it into a CommunicationAction. The bufferize() method runs the inverse of this method.
+        it into an Action. The bufferize() method runs the inverse of this method.
 
         Args:
             worker (AbstractWorker): a reference to the worker necessary for detailing. Read
@@ -140,22 +114,11 @@ def unbufferize(
             protobuf_obj (CommunicationActionPB): the Protobuf message
 
         Returns:
-            obj_id (CommunicationAction): a CommunicationAction
+            obj (CommunicationAction): a CommunicationAction
 
         Examples:
             message = unbufferize(sy.local_worker, protobuf_msg)
         """
-        name = protobuf_obj.name
-
-        obj_id = sy.serde.protobuf.proto.get_protobuf_id(protobuf_obj.obj_id)
-        source = sy.serde.protobuf.proto.get_protobuf_id(protobuf_obj.source)
-        destinations = [
-            sy.serde.protobuf.proto.get_protobuf_id(pb_id) for pb_id in protobuf_obj.destinations
-        ]
-
-        kwargs_ = {
-            key: sy.serde.protobuf.serde.unbufferize_arg(worker, kwarg)
-            for key, kwarg in protobuf_obj.kwargs.items()
-        }
+        attrs = Action.unbufferize(worker, protobuf_obj)
 
-        return CommunicationAction(obj_id, name, source, destinations, kwargs_)
+        return CommunicationAction(*attrs)
diff --git a/syft/execution/computation.py b/syft/execution/computation.py
index 2f46e118f6f..4918fbf8db0 100644
--- a/syft/execution/computation.py
+++ b/syft/execution/computation.py
@@ -29,120 +29,58 @@ def __init__(self, name, target, args_, kwargs_, return_ids, return_value=False)
                 and will need to do .get() later to get the result.
 
         """
-
-        # call the parent constructor - setting the type integer correctly
-        super().__init__()
-
-        self.name = name
-        self.target = target
-        self.args = args_
-        self.kwargs = kwargs_
-        self.return_ids = return_ids
-        self.return_value = return_value
+        super().__init__(name, target, args_, kwargs_, return_ids, return_value=return_value)
 
     @staticmethod
-    def simplify(worker: AbstractWorker, action: "ComputationAction") -> tuple:
+    def simplify(worker: AbstractWorker, action: "Action") -> tuple:
         """
-        This function takes the attributes of a Action and saves them in a tuple
+        This function takes the attributes of a ComputationAction and saves them in a tuple
         Args:
             worker (AbstractWorker): a reference to the worker doing the serialization
-            action (ComputationAction): the ComputationAction object to simplify
+            action (ComputationAction): a ComputationAction
         Returns:
-            tuple: a tuple holding the unique attributes of the message
+            tuple: a tuple holding the unique attributes of the ComputationAction
         Examples:
-            data = simplify(sy.local_worker, action)
+            data = simplify(worker, action)
         """
-        # NOTE: we can skip calling _simplify on return_ids because they should already be
-        # a list of simple types.
-        message = (action.name, action.target, action.args, action.kwargs)
-
-        return (
-            sy.serde.msgpack.serde._simplify(worker, message),
-            sy.serde.msgpack.serde._simplify(worker, action.return_ids),
-            sy.serde.msgpack.serde._simplify(worker, action.return_value),
-        )
+        return Action.simplify(worker, action)
 
     @staticmethod
-    def detail(worker: AbstractWorker, msg_tuple: tuple) -> "ComputationAction":
+    def detail(worker: AbstractWorker, action_tuple: tuple) -> "Action":
         """
         This function takes the simplified tuple version of this message and converts
-        it into a Action. The simplify() method runs the inverse of this method.
+        it into a ComputationAction. The simplify() method runs the inverse of this method.
 
         Args:
             worker (AbstractWorker): a reference to the worker necessary for detailing. Read
                 syft/serde/serde.py for more information on why this is necessary.
-            msg_tuple (Tuple): the raw information being detailed.
+            communication_tuple (Tuple): the raw information being detailed.
         Returns:
-            action (ComputationAction): a ComputationAction.
+            communication (ComputationAction): a ComputationAction.
         Examples:
-            action = detail(sy.local_worker, msg_tuple)
+            communication = detail(sy.local_worker, communication_tuple)
         """
-        message = msg_tuple[0]
-        return_ids = msg_tuple[1]
-        return_value = msg_tuple[2]
-
-        detailed_msg = sy.serde.msgpack.serde._detail(worker, message)
-        detailed_ids = sy.serde.msgpack.serde._detail(worker, return_ids)
-        detailed_return_value = sy.serde.msgpack.serde._detail(worker, return_value)
+        attrs = Action.detail(worker, action_tuple)
 
-        name, target, args_, kwargs_ = detailed_msg
-
-        return ComputationAction(name, target, args_, kwargs_, detailed_ids, detailed_return_value)
+        return ComputationAction(*attrs)
 
     @staticmethod
-    def bufferize(worker: AbstractWorker, action: "ComputationAction") -> "ComputationActionPB":
+    def bufferize(
+        worker: AbstractWorker, communication: "ComputationAction"
+    ) -> "ComputationActionPB":
         """
-        This function takes the attributes of a Action and saves them in Protobuf
+        This function takes the attributes of a ComputationAction and saves them in Protobuf
         Args:
             worker (AbstractWorker): a reference to the worker doing the serialization
-            action (Action): an Action
+            communication (ComputationAction): a ComputationAction
         Returns:
-            protobuf_obj: a Protobuf message holding the unique attributes of the message
+            protobuf_obj: a Protobuf message holding the unique attributes of the communication
         Examples:
-            data = bufferize(message)
+            data = bufferize(sy.local_worker, communication)
         """
-        protobuf_op = ComputationActionPB()
-        protobuf_op.command = action.name
-
-        protobuf_target = None
-        if isinstance(action.target, sy.generic.pointers.pointer_tensor.PointerTensor):
-            protobuf_target = protobuf_op.target_pointer
-        elif isinstance(action.target, sy.execution.placeholder_id.PlaceholderId):
-            protobuf_target = protobuf_op.target_placeholder_id
-        elif isinstance(action.target, (int, str)):
-            sy.serde.protobuf.proto.set_protobuf_id(protobuf_op.target_id, action.target)
-        elif action.target is not None:
-            protobuf_target = protobuf_op.target_tensor
-
-        if protobuf_target is not None:
-            protobuf_target.CopyFrom(sy.serde.protobuf.serde._bufferize(worker, action.target))
-
-        if action.args:
-            protobuf_op.args.extend(sy.serde.protobuf.serde.bufferize_args(worker, action.args))
-
-        if action.kwargs:
-            for key, value in action.kwargs.items():
-                protobuf_op.kwargs.get_or_create(key).CopyFrom(
-                    sy.serde.protobuf.serde.bufferize_arg(worker, value)
-                )
-
-        if action.return_ids is not None:
-            if not isinstance(action.return_ids, (list, tuple)):
-                return_ids = (action.return_ids,)
-            else:
-                return_ids = action.return_ids
-
-            for return_id in return_ids:
-                if isinstance(return_id, PlaceholderId):
-                    # NOTE to know when we have a PlaceholderId, we store it
-                    # in return_placeholder_ids and not in return_ids
-                    protobuf_op.return_placeholder_ids.append(
-                        sy.serde.protobuf.serde._bufferize(worker, return_id)
-                    )
-                else:
-                    sy.serde.protobuf.proto.set_protobuf_id(protobuf_op.return_ids.add(), return_id)
-
-        return protobuf_op
+        protobuf_action = ComputationActionPB()
+
+        return Action.bufferize(worker, communication, protobuf_action)
 
     @staticmethod
     def unbufferize(
@@ -163,78 +101,6 @@ def unbufferize(
         Examples:
             message = unbufferize(sy.local_worker, protobuf_msg)
         """
-        command = protobuf_obj.command
-        protobuf_target = protobuf_obj.WhichOneof("target")
-        if protobuf_target:
-            target = sy.serde.protobuf.serde._unbufferize(
-                worker, getattr(protobuf_obj, protobuf_obj.WhichOneof("target"))
-            )
-        else:
-            target = None
-        args_ = sy.serde.protobuf.serde.unbufferize_args(worker, protobuf_obj.args)
-
-        kwargs_ = {}
-        for key in protobuf_obj.kwargs:
-            kwargs_[key] = sy.serde.protobuf.serde.unbufferize_arg(worker, protobuf_obj.kwargs[key])
-
-        return_ids = [
-            sy.serde.protobuf.proto.get_protobuf_id(pb_id) for pb_id in protobuf_obj.return_ids
-        ]
-
-        return_placeholder_ids = [
-            sy.serde.protobuf.serde._unbufferize(worker, placeholder)
-            for placeholder in protobuf_obj.return_placeholder_ids
-        ]
-
-        if return_placeholder_ids:
-            action = ComputationAction(
-                command, target, tuple(args_), kwargs_, return_placeholder_ids
-            )
-        else:
-            action = ComputationAction(command, target, tuple(args_), kwargs_, tuple(return_ids))
-
-        return action
-
-    def code(self, var_names=None) -> str:
-        """Returns pseudo-code representation of computation action"""
-
-        def stringify(obj):
-            if isinstance(obj, PlaceholderId):
-                id = obj.value
-                if var_names is None:
-                    ret = f"var_{id}"
-                else:
-                    if id in var_names:
-                        ret = var_names[id]
-                    else:
-                        idx = sum("var_" in k for k in var_names.values())
-                        name = f"var_{idx}"
-                        var_names[id] = name
-                        ret = name
-            elif isinstance(obj, PlaceHolder):
-                ret = stringify(obj.id)
-            elif isinstance(obj, (tuple, list)):
-                ret = ", ".join(stringify(o) for o in obj)
-            else:
-                ret = str(obj)
-
-            return ret
-
-        out = ""
-        if self.return_ids is not None:
-            out += stringify(self.return_ids) + " = "
-        if self.target is not None:
-            out += stringify(self.target) + "."
-        out += self.name + "("
-        out += stringify(self.args)
-        if self.kwargs:
-            if len(self.args) > 0:
-                out += ", "
-            out += ", ".join(f"{k}={w}" for k, w in self.kwargs.items())
-        out += ")"
-
-        return out
-
-    def __str__(self) -> str:
-        """Returns string representation of computation action"""
-        return f"{type(self).__name__}[{self.code()}]"
+        attrs = Action.unbufferize(worker, protobuf_obj)
+
+        return ComputationAction(*attrs)
diff --git a/syft/frameworks/torch/mpc/fss.py b/syft/frameworks/torch/mpc/fss.py
index 8b7a0d4f680..cbcfc809142 100644
--- a/syft/frameworks/torch/mpc/fss.py
+++ b/syft/frameworks/torch/mpc/fss.py
@@ -60,8 +60,8 @@ def initialize_crypto_plans(worker):
 
 
 def request_run_plan(worker, plan_tag, location, return_value, args=tuple(), kwargs=dict()):
-    response_ids = [sy.ID_PROVIDER.pop()]
-    args = [args, response_ids]
+    response_ids = (sy.ID_PROVIDER.pop(),)
+    args = (args, response_ids)
 
     command = ("run", plan_tag, args, kwargs)
 
diff --git a/syft/generic/object_storage.py b/syft/generic/object_storage.py
index e4b3e6138c4..11bc5112ab7 100644
--- a/syft/generic/object_storage.py
+++ b/syft/generic/object_storage.py
@@ -86,6 +86,7 @@ def set_obj(self, obj: Union[FrameworkTensorType, AbstractTensor]) -> None:
         Args:
             obj: A torch or syft tensor with an id.
         """
+        obj.owner = self
         self._objects[obj.id] = obj
         # Add entry in the tag index
         if obj.tags:
diff --git a/syft/generic/pointers/callable_pointer.py b/syft/generic/pointers/callable_pointer.py
index 44ad721ce3b..19de2f9788b 100644
--- a/syft/generic/pointers/callable_pointer.py
+++ b/syft/generic/pointers/callable_pointer.py
@@ -63,7 +63,7 @@ def __init__(
     def __call__(self, *args, **kwargs):
         command = ("__call__", self.id_at_location, args, kwargs)
 
-        return_ids = [sy.ID_PROVIDER.pop()]
+        return_ids = (sy.ID_PROVIDER.pop(),)
         response = self.owner.send_command(
             message=command, recipient=self.location, return_ids=return_ids
         )
diff --git a/syft/generic/pointers/pointer_dataset.py b/syft/generic/pointers/pointer_dataset.py
index 97cfe041e74..9a20d733579 100644
--- a/syft/generic/pointers/pointer_dataset.py
+++ b/syft/generic/pointers/pointer_dataset.py
@@ -32,13 +32,13 @@ def __init__(
 
     @property
     def data(self):
-        command = ("get_data", self.id_at_location, [], {})
+        command = ("get_data", self.id_at_location, (), {})
         ptr = self.owner.send_command(message=command, recipient=self.location).wrap()
         return ptr
 
     @property
     def targets(self):
-        command = ("get_targets", self.id_at_location, [], {})
+        command = ("get_targets", self.id_at_location, (), {})
         ptr = self.owner.send_command(message=command, recipient=self.location).wrap()
         return ptr
 
@@ -83,12 +83,12 @@ def __repr__(self):
         return out
 
     def __len__(self):
-        command = ("__len__", self.id_at_location, [], {})
+        command = ("__len__", self.id_at_location, (), {})
         len = self.owner.send_command(message=command, recipient=self.location)
         return len
 
     def __getitem__(self, index):
-        command = ("__getitem__", self.id_at_location, [index], {})
+        command = ("__getitem__", self.id_at_location, (index,), {})
         data_elem, target_elem = self.owner.send_command(message=command, recipient=self.location)
         return data_elem.wrap(), target_elem.wrap()
 
diff --git a/syft/generic/pointers/pointer_plan.py b/syft/generic/pointers/pointer_plan.py
index b85d66eee59..109bcd924de 100644
--- a/syft/generic/pointers/pointer_plan.py
+++ b/syft/generic/pointers/pointer_plan.py
@@ -115,7 +115,7 @@ def parameters(self) -> List:
         location = self._locations[0]
         id_at_location = self._ids_at_location[0]
 
-        command = ("parameters", id_at_location, [], {})
+        command = ("parameters", id_at_location, (), {})
 
         pointers = self.owner.send_command(message=command, recipient=location)
 
@@ -161,10 +161,10 @@ def request_run_plan(
                 id_at_location = id_at_loc
                 break
 
-        command = ("run", id_at_location, args, kwargs)
+        command = ("run", id_at_location, tuple(args), kwargs)
 
         response = self.owner.send_command(
-            message=command, recipient=location, return_ids=response_ids
+            message=command, recipient=location, return_ids=tuple(response_ids)
         )
         response = hook_args.hook_response(plan_name, response, wrap_type=FrameworkTensor[0])
         if isinstance(response, (list, tuple)):
diff --git a/syft/generic/pointers/pointer_tensor.py b/syft/generic/pointers/pointer_tensor.py
index ccf7539714c..0e969b1af5e 100644
--- a/syft/generic/pointers/pointer_tensor.py
+++ b/syft/generic/pointers/pointer_tensor.py
@@ -295,7 +295,7 @@ def remote_send(self, destination: AbstractWorker, requires_grad: bool = False):
         """
         kwargs_ = {"inplace": False, "requires_grad": requires_grad}
         message = TensorCommandMessage.communication(
-            self.id_at_location, "remote_send", self.location.id, [destination.id], kwargs_
+            "remote_send", self, (destination.id,), kwargs_, (self.id,)
         )
         self.owner.send_msg(message=message, location=self.location)
         return self
diff --git a/syft/messaging/message.py b/syft/messaging/message.py
index 4fa7ed76bc4..0b3ed29d89f 100644
--- a/syft/messaging/message.py
+++ b/syft/messaging/message.py
@@ -115,11 +115,11 @@ def computation(name, target, args_, kwargs_, return_ids, return_value=False):
         return TensorCommandMessage(action)
 
     @staticmethod
-    def communication(obj_id, name, source, destinations, kwargs_):
+    def communication(name, target, args_, kwargs_, return_ids):
         """ Helper function to build a TensorCommandMessage containing a CommunicationAction
         directly from the action arguments.
         """
-        action = CommunicationAction(obj_id, name, source, destinations, kwargs_)
+        action = CommunicationAction(name, target, args_, kwargs_, return_ids)
         return TensorCommandMessage(action)
 
     @staticmethod
diff --git a/syft/serde/protobuf/serde.py b/syft/serde/protobuf/serde.py
index 8de81bfaff3..d08a7f39f0b 100644
--- a/syft/serde/protobuf/serde.py
+++ b/syft/serde/protobuf/serde.py
@@ -401,22 +401,27 @@ def bufferize_arg(worker: AbstractWorker, arg: object) -> ArgPB:
         setattr(protobuf_arg, attr_name, arg)
     except:
         getattr(protobuf_arg, attr_name).CopyFrom(_bufferize(worker, arg))
+
     return protobuf_arg
 
 
 def unbufferize_args(worker: AbstractWorker, protobuf_args: list) -> list:
     args_ = []
     for protobuf_arg in protobuf_args:
-        args_.append(unbufferize_arg(worker, protobuf_arg))
-    return args_
+        unbufferized = unbufferize_arg(worker, protobuf_arg)
+        args_.append(unbufferized)
+    return tuple(args_)
 
 
 def unbufferize_arg(worker: AbstractWorker, protobuf_arg: ArgPB) -> object:
-    protobuf_arg_field = getattr(protobuf_arg, protobuf_arg.WhichOneof("arg"))
+    protobuf_field_name = protobuf_arg.WhichOneof("arg")
+
+    protobuf_arg_field = getattr(protobuf_arg, protobuf_field_name)
     try:
         arg = _unbufferize(worker, protobuf_arg_field)
     except:
         arg = protobuf_arg_field
+
     return arg
 
 
diff --git a/syft/workers/base.py b/syft/workers/base.py
index 2dbf6aa98ca..9b49c3e8083 100644
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -549,23 +549,22 @@ def execute_computation_action(self, action: ComputationAction) -> PointerTensor
                 raise ResponseSignatureError(new_ids)
 
     def execute_communication_action(self, action: CommunicationAction) -> PointerTensor:
-        obj_id = action.obj_id
-        source = action.source
-        destinations = action.destinations
+        owner = action.target.owner
+        destinations = [self.get_worker(id_) for id_ in action.args]
         kwargs_ = action.kwargs
-        source_worker = self.get_worker(source)
-        if source_worker != self:
+
+        if owner != self:
             return None
         else:
-            obj = self.get_obj(obj_id)
-            response = source_worker.send(obj, *destinations, **kwargs_)
+            obj = self.get_obj(action.target.id)
+            response = owner.send(obj, *destinations, **kwargs_)
             response.garbage_collect_data = False
             if kwargs_.get("requires_grad", False):
                 response = hook_args.register_response(
                     "send", response, [sy.ID_PROVIDER.pop()], self
                 )
             else:
-                self.rm_obj(obj_id)
+                self.rm_obj(action.target.id)
             return response
 
     def execute_worker_command(self, message: tuple):
diff --git a/test/execution/test_communication.py b/test/execution/test_communication.py
new file mode 100644
index 00000000000..3150d9783ae
--- /dev/null
+++ b/test/execution/test_communication.py
@@ -0,0 +1,16 @@
+import pytest
+
+import syft as sy
+
+from syft.execution.communication import CommunicationAction
+
+
+def test_communication_methods_accepted():
+    c = CommunicationAction("send", None, (), {}, ())
+
+    assert c.name == "send"
+
+
+def test_computation_methods_rejected():
+    with pytest.raises(ValueError):
+        CommunicationAction("__add__", None, (), {}, ())
diff --git a/test/execution/test_role.py b/test/execution/test_role.py
new file mode 100644
index 00000000000..cbc57179e63
--- /dev/null
+++ b/test/execution/test_role.py
@@ -0,0 +1,52 @@
+import pytest
+
+import torch
+
+import syft as sy
+from syft.execution.role import Role
+from syft.execution.placeholder import PlaceHolder
+from syft.execution.computation import ComputationAction
+from syft.execution.communication import CommunicationAction
+
+
+def test_register_computation_action():
+    role = Role()
+    placeholder = PlaceHolder()
+    target = torch.ones([1])
+
+    action = ("__add__", target, (), {})
+
+    role.register_action((action, placeholder), ComputationAction)
+
+    assert len(role.actions) == 1
+
+    registered = role.actions[0]
+
+    assert isinstance(registered, ComputationAction)
+    assert registered.name == "__add__"
+    assert registered.target == target
+    assert registered.args == ()
+    assert registered.kwargs == {}
+    assert registered.return_ids == (placeholder.id,)
+
+
+def test_register_communication_action():
+    role = Role()
+    placeholder = PlaceHolder()
+    target = torch.ones([1])
+
+    action = ("get", target, (), {})
+
+    role.register_action((action, placeholder), CommunicationAction)
+
+    assert len(role.actions) == 1
+
+    registered = role.actions[0]
+
+    assert isinstance(registered, CommunicationAction)
+    assert registered.name == "get"
+    assert registered.target == target
+    assert registered.args == ()
+    assert registered.kwargs == {}
+
+    assert registered.return_ids == (placeholder.id,)
diff --git a/test/generic/test_object_storage.py b/test/generic/test_object_storage.py
index 6bf623099db..5bb5a3adbbc 100644
--- a/test/generic/test_object_storage.py
+++ b/test/generic/test_object_storage.py
@@ -45,3 +45,19 @@ def test_clear_objects_return_None():
     objs = obj_storage.current_objects()
     assert len(objs) == 0
     assert ret_val is None
+
+
+def test_set_obj_takes_ownership(workers):
+    me = workers["me"]
+    bob = workers["bob"]
+
+    x = torch.tensor(1)
+
+    x.owner = bob
+
+    me.set_obj(x)
+
+    objs = me._objects
+
+    assert objs[x.id] == x
+    assert objs[x.id].owner == workers["me"]
diff --git a/test/serde/msgpack/test_msgpack_serde_full.py b/test/serde/msgpack/test_msgpack_serde_full.py
index dfcba4b4e41..d9f7898723a 100644
--- a/test/serde/msgpack/test_msgpack_serde_full.py
+++ b/test/serde/msgpack/test_msgpack_serde_full.py
@@ -85,7 +85,7 @@
 samples[syft.messaging.message.ObjectRequestMessage] = make_objectrequestmessage
 samples[syft.messaging.message.PlanCommandMessage] = make_plancommandmessage
 samples[syft.messaging.message.SearchMessage] = make_searchmessage
-samples[syft.messaging.message.TensorCommandMessage] = make_command_message
+samples[syft.messaging.message.TensorCommandMessage] = make_tensor_command_message
 samples[syft.messaging.message.WorkerCommandMessage] = make_workercommandmessage
 
 samples[syft.workers.base.BaseWorker] = make_baseworker
diff --git a/test/serde/protobuf/test_protobuf_serde_full.py b/test/serde/protobuf/test_protobuf_serde_full.py
index ee8f9d4f2f9..c256792e825 100644
--- a/test/serde/protobuf/test_protobuf_serde_full.py
+++ b/test/serde/protobuf/test_protobuf_serde_full.py
@@ -41,7 +41,7 @@
 
 # Syft Messages
 samples[syft.messaging.message.ObjectMessage] = make_objectmessage
-samples[syft.messaging.message.TensorCommandMessage] = make_command_message
+samples[syft.messaging.message.TensorCommandMessage] = make_tensor_command_message
 
 
 def test_serde_coverage():
diff --git a/test/serde/serde_helpers.py b/test/serde/serde_helpers.py
index b1fc1c22ada..88c5f1d5319 100644
--- a/test/serde/serde_helpers.py
+++ b/test/serde/serde_helpers.py
@@ -1378,26 +1378,34 @@ def make_communication_action(**kwargs):
     bob.log_msgs = False
 
     def compare(detailed, original):
+        assert type(detailed) == syft.messaging.message.CommunicationAction
+
         detailed_msg = (
-            detailed.obj_id,
             detailed.name,
-            detailed.source,
-            detailed.destinations,
+            detailed.target,
+            detailed.args,
             detailed.kwargs,
+            detailed.return_ids,
+            detailed.return_value,
         )
         original_msg = (
-            original.obj_id,
             original.name,
-            original.source,
-            original.destinations,
+            original.target,
+            original.args,
             original.kwargs,
+            original.return_ids,
+            original.return_value,
         )
-        assert type(detailed) == syft.messaging.message.CommunicationAction
+
         for i in range(len(original_msg)):
-            assert detailed_msg[i] == original_msg[i]
-        return True
+            if type(original_msg[i]) != torch.Tensor:
+                assert detailed_msg[i] == original_msg[i], f"{detailed_msg[i]} != {original_msg[i]}"
+            else:
+                assert detailed_msg[i].equal(
+                    original_msg[i]
+                ), f"{detailed_msg[i]} != {original_msg[i]}"
 
-    msg = (com.obj_id, com.name, com.source, com.destinations, com.kwargs)
+        return True
 
     return [
         {
@@ -1405,11 +1413,12 @@ def compare(detailed, original):
             "simplified": (
                 CODE[syft.execution.communication.CommunicationAction],
                 (
-                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.obj_id),
                     msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.name),
-                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.source),
-                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.destinations),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.target),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.args),
                     msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.kwargs),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.return_ids),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], com.return_value),
                 ),
             ),
             "cmp_detailed": compare,
@@ -1433,19 +1442,35 @@ def make_computation_action(**kwargs):
     bob.log_msgs = False
 
     def compare(detailed, original):
-        detailed_msg = (detailed.name, detailed.target, detailed.args, detailed.kwargs)
-        original_msg = (original.name, original.target, original.args, original.kwargs)
         assert type(detailed) == syft.execution.computation.ComputationAction
+
+        detailed_msg = (
+            detailed.name,
+            detailed.target,
+            detailed.args,
+            detailed.kwargs,
+            detailed.return_ids,
+            detailed.return_value,
+        )
+
+        original_msg = (
+            original.name,
+            original.target,
+            original.args,
+            original.kwargs,
+            original.return_ids,
+            original.return_value,
+        )
+
         for i in range(len(original_msg)):
             if type(original_msg[i]) != torch.Tensor:
-                assert detailed_msg[i] == original_msg[i]
+                assert detailed_msg[i] == original_msg[i], f"{detailed_msg[i]} != {original_msg[i]}"
             else:
-                assert detailed_msg[i].equal(original_msg[i])
-        assert detailed.return_ids == original.return_ids
-        return True
+                assert detailed_msg[i].equal(
+                    original_msg[i]
+                ), f"{detailed_msg[i]} != {original_msg[i]}"
 
-    message1 = (op1.name, op1.target, op1.args, op1.kwargs)
-    message2 = (op2.name, op2.target, op2.args, op2.kwargs)
+        return True
 
     return [
         {
@@ -1453,11 +1478,12 @@ def compare(detailed, original):
             "simplified": (
                 CODE[syft.execution.computation.ComputationAction],
                 (
-                    msgpack.serde._simplify(
-                        kwargs["workers"]["serde_worker"], message1
-                    ),  # (Any) message
-                    (CODE[tuple], (op1.return_ids[0],)),  # (tuple) return_ids
-                    False,  # return value
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.name),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.target),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.args),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.kwargs),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.return_ids),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op1.return_value),
                 ),
             ),
             "cmp_detailed": compare,
@@ -1467,11 +1493,12 @@ def compare(detailed, original):
             "simplified": (
                 CODE[syft.execution.computation.ComputationAction],
                 (
-                    msgpack.serde._simplify(
-                        kwargs["workers"]["serde_worker"], message2
-                    ),  # (Any) message
-                    (CODE[tuple], (op2.return_ids[0],)),  # (tuple) return_ids,
-                    False,  # return value
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.name),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.target),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.args),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.kwargs),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.return_ids),
+                    msgpack.serde._simplify(kwargs["workers"]["serde_worker"], op2.return_value),
                 ),
             ),
             "cmp_detailed": compare,
@@ -1480,7 +1507,7 @@ def compare(detailed, original):
 
 
 # syft.messaging.message.TensorCommandMessage
-def make_command_message(**kwargs):
+def make_tensor_command_message(**kwargs):
     bob = kwargs["workers"]["bob"]
     alice = kwargs["workers"]["alice"]
     bob.log_msgs = True
@@ -1500,23 +1527,34 @@ def make_command_message(**kwargs):
     bob.log_msgs = False
 
     def compare(detailed, original):
-        if isinstance(detailed.action, syft.execution.computation.ComputationAction):
-            detailed = detailed.action
-            original = original.action
-
-            detailed_msg = (detailed.name, detailed.target, detailed.args, detailed.kwargs)
-            original_msg = (original.name, original.target, original.args, original.kwargs)
-            for i in range(len(original_msg)):
-                if type(original_msg[i]) != torch.Tensor:
-                    assert detailed_msg[i] == original_msg[i]
-                else:
-                    assert detailed_msg[i].equal(original_msg[i])
-            assert detailed.return_ids == original.return_ids
-            return True
-
-        elif isinstance(detailed.action, syft.execution.communication.CommunicationAction):
-            assert detailed.action == original.action
-            return True
+        detailed_action = detailed.action
+        original_action = original.action
+
+        detailed_action = (
+            detailed.name,
+            detailed.target,
+            detailed.args,
+            detailed.kwargs,
+            detailed.return_ids,
+            detailed.return_value,
+        )
+
+        original_action = (
+            original.name,
+            original.target,
+            original.args,
+            original.kwargs,
+            original.return_ids,
+            original.return_value,
+        )
+
+        for i in range(len(original_action)):
+            if type(original_action[i]) != torch.Tensor:
+                assert detailed_action[i] == original_action[i]
+            else:
+                assert detailed_action[i].equal(original_action[i])
+
+        return True
 
     return [
         {


Code review comments:
diff_hunk: @@ -401,22 +401,27 @@ def bufferize_arg(worker: AbstractWorker, arg: object) -> ArgPB:
         setattr(protobuf_arg, attr_name, arg)
     except:
         getattr(protobuf_arg, attr_name).CopyFrom(_bufferize(worker, arg))
+
     return protobuf_arg
 
 
 def unbufferize_args(worker: AbstractWorker, protobuf_args: list) -> list:
     args_ = []
     for protobuf_arg in protobuf_args:
-        args_.append(unbufferize_arg(worker, protobuf_arg))
-    return args_
+        unbufferized = unbufferize_arg(worker, protobuf_arg)
path: syft/serde/protobuf/serde.py
text: @user1:
I would have done it with a list comprehension here, but it's a detail 

@author:
You're right, that's cleaner. Fixing.