From 04c022e9d8ac4f36d484a35acfb9dce53635eba3 Mon Sep 17 00:00:00 2001
From: Ahdra Merali <ahdra.merali@quantumblack.com>
Date: Mon, 29 Jan 2024 14:48:10 +0000
Subject: [PATCH] Naive fix for return statements

Signed-off-by: Ahdra Merali <ahdra.merali@quantumblack.com>
---
 kedro/ipython/__init__.py     | 25 ++++++++++++++++++++++---
 tests/ipython/test_ipython.py |  9 ++++++---
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/kedro/ipython/__init__.py b/kedro/ipython/__init__.py
index 1dfcfc3381..2e3eb249f9 100644
--- a/kedro/ipython/__init__.py
+++ b/kedro/ipython/__init__.py
@@ -301,12 +301,31 @@ def _prepare_function_body(func: Callable) -> str:
     # https://stackoverflow.com/questions/38050649/getting-a-python-functions-source-code-without-the-definition-lines
     all_source_lines = inspect.getsourcelines(func)[0]
     # Remove any decorators
-    func_lines = dropwhile(lambda x: x.startswith("@"), all_source_lines)
-    line: str = next(func_lines).strip()
+    raw_func_lines = dropwhile(lambda x: x.startswith("@"), all_source_lines)
+    line: str = next(raw_func_lines).strip()
     if not line.startswith("def "):
         return line.rsplit(",")[0].strip()
     # Handle functions that are not one-liners
-    first_line = next(func_lines)
+    first_line = next(raw_func_lines)
+
+    # Convert return statements into print statements
+    func_lines = []
+    for line in raw_func_lines:
+        if line.lstrip().startswith("return"):
+            return_statement = line.lstrip()
+            line_indentation = len(line) - len(return_statement)
+            commented_return = line[:line_indentation] + "# " + return_statement
+            printed_return = (
+                line[:line_indentation]
+                + "print("
+                + return_statement[len("return") :].strip()
+                + ")"
+            )
+            func_lines.append(commented_return)
+            func_lines.append(printed_return)
+        else:
+            func_lines.append(line)
+
     # Find the indentation of the first line
     indentation = len(first_line) - len(first_line.lstrip())
     body = "".join(
diff --git a/tests/ipython/test_ipython.py b/tests/ipython/test_ipython.py
index 7f1f428863..5609ad74cb 100644
--- a/tests/ipython/test_ipython.py
+++ b/tests/ipython/test_ipython.py
@@ -420,7 +420,8 @@ def test_load_node(self, mocker, dummy_function_file_lines, dummy_pipelines):
 # this is an in-line comment in the body of the function
 random_assignment = "Added for a longer function"
 random_assignment += "make sure to modify variable"
-return not dummy_input"""
+# return not dummy_input
+print(not dummy_input)"""
 
         expected_cells = [
             node_inputs,
@@ -492,7 +493,8 @@ def test_prepare_function_body(self):
 # this is an in-line comment in the body of the function
 random_assignment = "Added for a longer function"
 random_assignment += "make sure to modify variable"
-return not dummy_input"""
+# return not dummy_input
+print(not dummy_input)"""
 
         result = _prepare_function_body(dummy_function)
         assert result == func_strings
@@ -515,7 +517,8 @@ def test_get_nested_function_body(self):
     def test_get_function_with_loop_body(self):
         func_strings = """for x in dummy_list:
     continue
-return len(dummy_list)"""
+# return len(dummy_list)
+print(len(dummy_list))"""
 
         result = _prepare_function_body(dummy_function_with_loop)
         assert result == func_strings