From 0ca0d5aee1019913f9461c946dd430e223b189c1 Mon Sep 17 00:00:00 2001 From: Philippe Moussalli Date: Thu, 1 Feb 2024 12:08:06 +0100 Subject: [PATCH] implement PR feedback --- src/fondant/pipeline/lightweight_component.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/fondant/pipeline/lightweight_component.py b/src/fondant/pipeline/lightweight_component.py index a3d7cf97..1a39a708 100644 --- a/src/fondant/pipeline/lightweight_component.py +++ b/src/fondant/pipeline/lightweight_component.py @@ -285,13 +285,17 @@ def new_getfile(_object, _old_getfile=inspect.getfile): raise TypeError(msg) -def is_running_in_notebook(): - """Check if the code is running in a Jupyter notebook.""" +def is_running_interactively(): + """Check if the code is running in an interactive environment.""" try: from IPython import get_ipython shell = get_ipython().__class__.__name__ - return shell == "ZMQInteractiveShell" + return shell in [ + "ZMQInteractiveShell", + "TerminalInteractiveShell", + "PyDevTerminalInteractiveShell", + ] except ModuleNotFoundError: return False @@ -313,18 +317,20 @@ def build_python_script(component_cls: t.Type[Component]) -> str: """, ) - if is_running_in_notebook(): + if is_running_interactively(): from IPython.core.magics.code import extract_symbols - inspect.getfile = new_getfile component_source = "".join( inspect.linecache.getlines( # type: ignore[attr-defined] new_getfile(component_cls), ), ) - component_source = extract_symbols(component_source, component_cls.__name__)[0][ - 0 - ] + component_source = extract_symbols( + component_source, + component_cls.__name__, + ) + component_source = component_source[0][0] + else: component_source = inspect.getsource(component_cls)