From 98ba2695c21ed70e81f3a3a568d176e732190d52 Mon Sep 17 00:00:00 2001
From: Madhur Tandon <madhur@outerbounds.co>
Date: Tue, 28 May 2024 01:53:45 +0530
Subject: [PATCH] fix leaked message

---
 metaflow/runner/nbrun.py | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/metaflow/runner/nbrun.py b/metaflow/runner/nbrun.py
index 13f981658d5..21d4076e15a 100644
--- a/metaflow/runner/nbrun.py
+++ b/metaflow/runner/nbrun.py
@@ -1,22 +1,14 @@
 import ast
 import os
-import shutil
 import tempfile
 from typing import Dict, Optional
 
 from metaflow import Runner
 
-try:
-    from IPython import get_ipython
-
-    ipython = get_ipython()
-except ModuleNotFoundError:
-    print("'nbrun' requires an interactive python environment (such as Jupyter)")
-
 DEFAULT_DIR = tempfile.gettempdir()
 
 
-def get_current_cell():
+def get_current_cell(ipython):
     if ipython:
         return ipython.history_manager.input_hist_raw[-1]
     return None
@@ -56,7 +48,17 @@ def __init__(
         base_dir: str = DEFAULT_DIR,
         **kwargs,
     ):
-        self.cell = get_current_cell()
+        try:
+            from IPython import get_ipython
+
+            ipython = get_ipython()
+        except ModuleNotFoundError:
+            print(
+                "'nbrun' requires an interactive python environment (such as Jupyter)"
+            )
+            return
+
+        self.cell = get_current_cell(ipython)
         self.flow = flow
         self.show_output = show_output