Repository: AgnostiqHQ/covalent (language: python)

Issue:
`Electron.wait_for` in a workflow causes the workflow function to no longer be callable
The new `Electron.wait_for()` method allows users to instruct Covalent that a task (task B) should wait for another task (task A) to return before being dispatched. This is needed in some instances because Covalent normally infers the dependency relations between electrons just from their inputs and outputs and therefore would not be able to handle scenarios where task B depends only on a side-effect of task A.

However, when users construct a function describing their workflow, invoking the `wait_for()` method on an electron causes the workflow function to no longer be callable:
```python

@ct.lattice
def workflow_function():
    res1 = task_1() # task_1 is an electron which returns an int at runtime
    res2 = task_2() # evaluates to an int at runtime
    res2.wait_for(res1) # int doesn't have a wait_for method
```
When `workflow_function` is called, `res1` and `res2` are the return values of `task_1` and `task_2`, not Electrons, so Python function cannot evaluate the third statement. 

This behavior is unintuitive because workflow functions are Python *functions*, and functions are expected to be callable. The `ct.lattice` decorator is not relevant here because:
* Users may construct their workflow function first as an ordinary function, and then construct a lattice explicitly using `ct.lattice()`, [as we demonstrate](https://covalent.readthedocs.io/en/latest/tutorials/machine_learning/mnist_classifier.html#Transform-and-run-workflow-with-Covalent) in the MNIST tutorial.
* Decorators are expected to merely [extend the behavior of a function](https://realpython.com/primer-on-python-decorators/). 



**Suggestion**: Instead of exposing the `Electron.wait_for` method directly to users, we can treat that a private method like the other methods that manipulate the transport graph (e.g. `Electron.connect_node_with_others()`), and expose a static wrapper function that invokes `Electron.wait_for` only in the context of an active lattice (during graph building and postprocessing). For example:

```python
def wait(child, parents):
    active_lattice = active_lattice_manager.get_active_lattice()

    if active_lattice:
        return child.wait_for(parents)
    else:
        return child
```
Then the following modified workflow function can both be called and dispatched:

```python
@ct.lattice
def workflow_function():
    res1 = task_1()
    res2 = task_2() 
    ct.wait(child=res2, parents=res1)
```



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a26e2c7e5..0df2688f5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Changed
+
+- Expose a public `wait()` function compatible with both calling and dispatching lattices
+
+### Docs
+
+- Updated the RTD on `wait_for()` to use the static `wait()` function
 ### Operations
 
 - pre-commit autoupdate
diff --git a/covalent/__init__.py b/covalent/__init__.py
index f3d7bb284..41a31c384 100644
--- a/covalent/__init__.py
+++ b/covalent/__init__.py
@@ -32,6 +32,7 @@
 from ._shared_files.config import get_config, reload_config, set_config
 from ._shared_files.util_classes import RESULT_STATUS as status
 from ._workflow import DepsBash, DepsCall, DepsPip, Lepton, TransportableObject, electron, lattice
+from ._workflow.electron import wait
 
 __all__ = [s for s in dir() if not s.startswith("_")]
 
diff --git a/covalent/_workflow/electron.py b/covalent/_workflow/electron.py
index 1513bf785..ce919d89f 100644
--- a/covalent/_workflow/electron.py
+++ b/covalent/_workflow/electron.py
@@ -611,3 +611,26 @@ def wrapper(*args, **kwargs):
         return decorator_electron
     else:  # decorator is called without arguments
         return decorator_electron(_func)
+
+
+def wait(child, parents):
+    """Instructs Covalent that an electron should wait for some other
+    tasks to complete before it is dispatched.
+
+    Args:
+        child: the dependent electron
+        parents: Electron(s) which must complete before `waiting_electron` starts
+
+    Returns:
+        waiting_electron
+
+    Useful when execution of an electron relies on a side-effect
+    from another one.
+
+    """
+    active_lattice = active_lattice_manager.get_active_lattice()
+
+    if active_lattice:
+        return child.wait_for(parents)
+    else:
+        return child
diff --git a/doc/source/concepts/concepts.rst b/doc/source/concepts/concepts.rst
index 632ee0083..015f54c02 100644
--- a/doc/source/concepts/concepts.rst
+++ b/doc/source/concepts/concepts.rst
@@ -183,7 +183,11 @@ This ensures that the independent experiments are performed in parallel rather t
 Waiting for other electrons
 ----------------------
 
-Sometimes the user might want to wait for a task's execution before executing another task even when the output of one is not the input of another. For those cases a `wait_for()` function is there which is called on the dependent electron and takes as input the depending electrons,
+Covalent normally infers the dependencies between electrons from
+their inputs and ouputs. Sometimes the user might want to wait for a
+task's execution before executing another task even when the output of
+one is not the input of another. The `wait()` function can be
+handy in those cases.
 
 .. code-block:: python
     :linenos:
@@ -214,14 +218,15 @@ Sometimes the user might want to wait for a task's execution before executing an
         res_1b = task_1b(2)
         res_1c = task_1c(2)
         res_2 = task_2(res_1a, 3)
-        res_3 = task_3(5).wait_for([res_1a, res_1b, res_1c])
+        res_3 = task_3(5)
+	ct.wait(child=res_3, parents=[res_1a, res_1b, res_1c])
 
         return task_2(res_2, res_3)
         ...
 
     res = ct.dispatch_sync(workflow)()
 
-This means that even though `task_3`'s input does not depend on anything, it's execution will still wait until all the tasks passed in `wait_for` have completed execution.
+The `wait()` statement instructs Covalent to wait for `task_1a`, `task_1b`, and `task_1c` to finish before dispatching `task_3`.
 
 
 Best Practices


Test patch:
diff --git a/tests/covalent_tests/workflow_stack_test.py b/tests/covalent_tests/workflow_stack_test.py
index 88ff2ca46..fc4c83bcc 100644
--- a/tests/covalent_tests/workflow_stack_test.py
+++ b/tests/covalent_tests/workflow_stack_test.py
@@ -709,7 +709,7 @@ def workflow():
         res_1a = task_1a(2)
         res_1b = task_1b(2)
         res_2 = task_2(res_1a, 3)
-        res_3 = task_3(5).wait_for([res_1a, res_1b])
+        res_3 = ct.wait(task_3(5), [res_1a, res_1b])
 
         return task_2(res_2, res_3)
 
@@ -729,6 +729,10 @@ def workflow():
     assert result.result == 1500
     rm._delete_result(dispatch_id)
 
+    # Check that workflow function returns the same result when called directly
+
+    assert workflow() == 1500
+
 
 def test_electron_getitem(db):
     """Test electron __getitem__, both with raw keys and with electron keys"""