From eba334fe64d39059d318eeb23b2cfe1619831366 Mon Sep 17 00:00:00 2001
From: Casey Jao <casey@agnostiq.ai>
Date: Mon, 5 Feb 2024 08:36:08 -0500
Subject: [PATCH 1/2] Parse function string in Electron constructor

Waiting until `build_graph()` won't work if the graph is being built
in a remote executor without the original source code.
---
 covalent/_workflow/electron.py                | 31 ++-----------------
 .../covalent_tests/workflow/electron_test.py  | 21 +------------
 2 files changed, 4 insertions(+), 48 deletions(-)

diff --git a/covalent/_workflow/electron.py b/covalent/_workflow/electron.py
index 12f18cbf5..a8a9055be 100644
--- a/covalent/_workflow/electron.py
+++ b/covalent/_workflow/electron.py
@@ -96,6 +96,7 @@ def __init__(
         self.metadata = metadata
         self.task_group_id = task_group_id
         self._packing_tasks = packing_tasks
+        self._function_string = get_serialized_function_str(function)
 
     @property
     def packing_tasks(self) -> bool:
@@ -442,7 +443,7 @@ def __call__(self, *args, **kwargs) -> Union[Any, "Electron"]:
             )
 
             name = sublattice_prefix + self.function.__name__
-            function_string = get_serialized_function_str(self.function)
+            function_string = self._function_string
             bound_electron = sub_electron(
                 self.function, json.dumps(parent_metadata), *args, **kwargs
             )
@@ -463,7 +464,7 @@ def __call__(self, *args, **kwargs) -> Union[Any, "Electron"]:
             name=self.function.__name__,
             function=self.function,
             metadata=self.metadata.copy(),
-            function_string=get_serialized_function_str(self.function),
+            function_string=self._function_string,
             task_group_id=self.task_group_id if self.packing_tasks else None,
         )
         self.task_group_id = self.task_group_id if self.packing_tasks else self.node_id
@@ -608,32 +609,6 @@ def _auto_dict_node(*args, **kwargs):
                 arg_index=arg_index,
             )
 
-    def add_collection_node_to_graph(self, graph: "_TransportGraph", prefix: str) -> int:
-        """
-        Adds the node to lattice's transport graph in the case
-        where a collection of electrons is passed as an argument
-        to another electron.
-
-        Args:
-            graph: Transport graph of the lattice
-            prefix: Prefix of the node
-
-        Returns:
-            node_id: Node id of the added node
-        """
-
-        new_metadata = encode_metadata(DEFAULT_METADATA_VALUES.copy())
-        if "executor" in self.metadata:
-            new_metadata["executor"] = self.metadata["executor"]
-            new_metadata["executor_data"] = self.metadata["executor_data"]
-
-        node_id = graph.add_node(
-            name=prefix,
-            function=to_decoded_electron_collection,
-            metadata=new_metadata,
-            function_string=get_serialized_function_str(to_decoded_electron_collection),
-        )
-
         return node_id
 
     def wait_for(self, electrons: Union["Electron", Iterable["Electron"]]):
diff --git a/tests/covalent_tests/workflow/electron_test.py b/tests/covalent_tests/workflow/electron_test.py
index a3db76bb7..d7a3e3192 100644
--- a/tests/covalent_tests/workflow/electron_test.py
+++ b/tests/covalent_tests/workflow/electron_test.py
@@ -36,7 +36,7 @@
     to_decoded_electron_collection,
 )
 from covalent._workflow.lattice import Lattice
-from covalent._workflow.transport import TransportableObject, _TransportGraph, encode_metadata
+from covalent._workflow.transport import TransportableObject, encode_metadata
 from covalent.executor.executor_plugins.local import LocalExecutor
 
 
@@ -252,25 +252,6 @@ def test_collection_node_helper_electron():
     assert to_decoded_electron_collection(x=dict_collection) == {"a": 1, "b": 2}
 
 
-def test_electron_add_collection_node():
-    """Test `to_decoded_electron_collection` in `Electron.add_collection_node`"""
-
-    def f(x):
-        return x
-
-    e = Electron(f)
-    tg = _TransportGraph()
-    node_id = e.add_collection_node_to_graph(tg, prefix=":")
-    collection_fn = tg.get_node_value(node_id, "function").get_deserialized()
-
-    collection = [
-        TransportableObject.make_transportable(1),
-        TransportableObject.make_transportable(2),
-    ]
-
-    assert collection_fn(x=collection) == [1, 2]
-
-
 def test_injected_inputs_are_not_in_tg():
     """Test that arguments to electrons injected by calldeps aren't
     added to the transport graph"""

From fe046a91db1d31df326aea2d4a14922f5b8e6c6b Mon Sep 17 00:00:00 2001
From: Casey Jao <casey@agnostiq.ai>
Date: Tue, 6 Feb 2024 17:56:14 -0500
Subject: [PATCH 2/2] Changelog

---
 CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce526d4da..2c0f77891 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Fixed
+
+- Sublattice electron function strings are now parsed correctly
+
 ### Operations
 
 - Added qelectron tests to the `tests` workflow