Repository: PrefectHQ/prefect

Issue:
Support `.pipe` for `EdgeAnnotation`
## Current behavior

`EdgeAnnotation` doesn't support the newly introduced `.pipe`


## Proposed behavior
```python
mapped = task.map([1,2,3,4]).pipe(flatten).pipe(task2)
```



Discussion:

Resolving patch:
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -30,6 +30,7 @@
 from prefect.utilities import logging
 from prefect.utilities.notifications import callback_factory
 from prefect.utilities.edges import EdgeAnnotation
+from prefect.utilities.tasks import as_task
 
 if TYPE_CHECKING:
     from prefect.core.flow import Flow
@@ -935,7 +936,9 @@ def outputs(self) -> Any:
             return_annotation = Any
         return return_annotation
 
-    def pipe(_prefect_self, _prefect_task: "Task", **kwargs: Any) -> "Task":
+    def pipe(
+        _prefect_self, _prefect_task: Union[Type[EdgeAnnotation], "Task"], **kwargs: Any
+    ) -> "Task":
         """
         "Pipes" the result of this task through another task. ``some_task().pipe(other_task)`` is
         equivalent to ``other_task(some_task())``, but can result in more readable code when used in a
@@ -950,7 +953,12 @@ def pipe(_prefect_self, _prefect_task: "Task", **kwargs: Any) -> "Task":
         """
         if "self" in kwargs:
             raise ValueError('You cannot use the keyword argument "self" in .pipe.')
-        return _prefect_task(_prefect_self, **kwargs)
+
+        if inspect.isclass(_prefect_task) and issubclass(_prefect_task, EdgeAnnotation):  # type: ignore
+            return as_task(_prefect_task(_prefect_self))
+        else:
+            # mypy<0.900 doesn't infer type as Task due to Union type
+            return _prefect_task(_prefect_self, **kwargs)  # type: ignore
 
     # Serialization ------------------------------------------------------------
 


Test patch:
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -47,6 +47,7 @@
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.debug import raise_on_exception
 from prefect.exceptions import TaskTimeoutSignal
+from prefect.utilities.edges import flatten, mapped, unmapped
 from prefect.utilities.tasks import pause_task
 
 
@@ -2553,3 +2554,30 @@ def accepts_anything(arg, **kwargs):
                     accepts_anything, self="some kwarg"
                 )
         state = flow.run()
+
+    def test_task_pipeline_edge_annotation(self):
+        """
+        Verify that passing EdgeAnnotation to pipe returns Task objects
+        that can chain further.
+        """
+        from prefect import Flow
+
+        @prefect.task
+        def add_one(number) -> int:
+            return number + 1
+
+        @prefect.task
+        def aggregate(numbers) -> int:
+            return sum(numbers)
+
+        with Flow(name="test") as flow:
+            total = (
+                add_one.map(range(5))
+                .pipe(flatten)
+                .pipe(mapped)
+                .pipe(unmapped)
+                .pipe(aggregate)
+            )
+
+        state = flow.run()
+        assert state.result[total].result == 15