From d2d01f9fe1551ae55545add00d16da048648d8d7 Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Fri, 9 Feb 2024 14:44:34 +0100
Subject: [PATCH] feat: Enhance `Pipeline.__repr__()` (#6963)

* Enhance Pipeline.draw() to show image directly in Jupyter notebook

* Add util method to check if we're in a Jupyter notebook

* Split Pipeline.draw() in two methods

* Update tests

* Update releasenotes

* Enhance Pipeline.__repr__

* Simplify Pipeline.__repr__

* Update release notes
---
 haystack/core/pipeline/pipeline.py            | 28 ++++++++++++
 .../notes/enhance-repr-0c5efa1e2ca6bafa.yaml  |  5 +++
 test/core/pipeline/test_pipeline.py           | 43 +++++++++++++++++++
 3 files changed, 76 insertions(+)
 create mode 100644 releasenotes/notes/enhance-repr-0c5efa1e2ca6bafa.yaml

diff --git a/haystack/core/pipeline/pipeline.py b/haystack/core/pipeline/pipeline.py
index 632cc73ef..98ba8df87 100644
--- a/haystack/core/pipeline/pipeline.py
+++ b/haystack/core/pipeline/pipeline.py
@@ -71,6 +71,34 @@ class Pipeline:
             return False
         return self.to_dict() == other.to_dict()
 
+    def __repr__(self) -> str:
+        """
+        Returns a text representation of the Pipeline.
+        If this runs in a Jupyter notebook, it will instead display the Pipeline image.
+        """
+        if is_in_jupyter():
+            # If we're in a Jupyter notebook we want to display the image instead of the text repr.
+            self.show()
+            return ""
+
+        res = f"{object.__repr__(self)}\n"
+        if self.metadata:
+            res += "🧱 Metadata\n"
+            for k, v in self.metadata.items():
+                res += f"  - {k}: {v}\n"
+
+        res += "🚅 Components\n"
+        for name, instance in self.graph.nodes(data="instance"):
+            res += f"  - {name}: {instance.__class__.__name__}\n"
+
+        res += "🛤️ Connections\n"
+        for sender, receiver, edge_data in self.graph.edges(data=True):
+            sender_socket = edge_data["from_socket"].name
+            receiver_socket = edge_data["to_socket"].name
+            res += f"  - {sender}.{sender_socket} -> {receiver}.{receiver_socket} ({edge_data['conn_type']})\n"
+
+        return res
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Returns this Pipeline instance as a dictionary.
diff --git a/releasenotes/notes/enhance-repr-0c5efa1e2ca6bafa.yaml b/releasenotes/notes/enhance-repr-0c5efa1e2ca6bafa.yaml
new file mode 100644
index 000000000..a9f1914ef
--- /dev/null
+++ b/releasenotes/notes/enhance-repr-0c5efa1e2ca6bafa.yaml
@@ -0,0 +1,5 @@
+---
+enhancements:
+  - |
+    Customize `Pipeline.__repr__()` to return a nice text representation of it.
+    If run on a Jupyter notebook it will instead have the same behaviour as `Pipeline.show()`.
diff --git a/test/core/pipeline/test_pipeline.py b/test/core/pipeline/test_pipeline.py
index c6dec1328..4e66f38f1 100644
--- a/test/core/pipeline/test_pipeline.py
+++ b/test/core/pipeline/test_pipeline.py
@@ -79,6 +79,49 @@ def test_get_component_name_not_added_to_pipeline():
     assert pipe.get_component_name(some_component) == ""
 
 
+@patch("haystack.core.pipeline.pipeline.is_in_jupyter")
+def test_repr(mock_is_in_jupyter):
+    pipe = Pipeline(metadata={"test": "test"}, max_loops_allowed=42)
+    pipe.add_component("add_two", AddFixedValue(add=2))
+    pipe.add_component("add_default", AddFixedValue())
+    pipe.add_component("double", Double())
+    pipe.connect("add_two", "double")
+    pipe.connect("double", "add_default")
+
+    expected_repr = (
+        f"{object.__repr__(pipe)}\n"
+        "🧱 Metadata\n"
+        "  - test: test\n"
+        "🚅 Components\n"
+        "  - add_two: AddFixedValue\n"
+        "  - add_default: AddFixedValue\n"
+        "  - double: Double\n"
+        "🛤️ Connections\n"
+        "  - add_two.result -> double.value (int)\n"
+        "  - double.value -> add_default.value (int)\n"
+    )
+    # Simulate not being in a notebook
+    mock_is_in_jupyter.return_value = False
+    assert repr(pipe) == expected_repr
+
+
+@patch("haystack.core.pipeline.pipeline.is_in_jupyter")
+def test_repr_in_notebook(mock_is_in_jupyter):
+    pipe = Pipeline(metadata={"test": "test"}, max_loops_allowed=42)
+    pipe.add_component("add_two", AddFixedValue(add=2))
+    pipe.add_component("add_default", AddFixedValue())
+    pipe.add_component("double", Double())
+    pipe.connect("add_two", "double")
+    pipe.connect("double", "add_default")
+
+    # Simulate being in a notebook
+    mock_is_in_jupyter.return_value = True
+
+    with patch.object(Pipeline, "show") as mock_show:
+        assert repr(pipe) == ""
+        mock_show.assert_called_once_with()
+
+
 def test_run_with_component_that_does_not_return_dict():
     BrokenComponent = component_class(
         "BrokenComponent", input_types={"a": int}, output_types={"b": int}, output=1  # type:ignore