Repository: kedro-org/kedro
File: kedro/ipython/__init__.py (Python)

Code before:
    for module in to_remove:
        del sys.modules[module]


def _find_kedro_project(current_dir: Path) -> Any:  # pragma: no cover
    while current_dir != current_dir.parent:
        if _is_project(current_dir):
            return current_dir
        current_dir = current_dir.parent

    return None


@typing.no_type_check
@magic_arguments()
@argument(
    "node",
    type=str,
    help=("Name of the Node."),
    nargs="?",
    default=None,
)
def magic_load_node(node: str) -> None:
    """The line magic %load_node <node_name>
    Currently it only supports Jupyter Notebook (>7.0) and Jupyter Lab. This line magic
    will generate code in multiple cells to load datasets from `DataCatalog`, import
    relevant functions and modules, node function definition and a function call.
    """
    cells = _load_node(node, pipelines)
    from ipylab import JupyterFrontEnd

    app = JupyterFrontEnd()

    def _create_cell_with_text(text: str) -> None:
        # Noted this only works with Notebook >7.0 or Jupyter Lab. It doesn't work with
        # VS Code Notebook due to imcompatible backends.
        app.commands.execute("notebook:insert-cell-below")
        app.commands.execute("notebook:replace-selection", {"text": text})

    for cell in cells:
        _create_cell_with_text(cell)


def _load_node(node_name: str, pipelines: _ProjectPipelines) -> list[str]:
    """Prepare the code to load dataset from catalog, import statements and function body.

    Args:
        node_name (str): The name of the node.

    Returns:
        list[str]: A list of string which is the generated code, each string represent a

Review comment (suggestion):
Again, mypy complains a lot about untype calls. It would be great if we can set this to `kedro` related libraries only, because we have no control of 3-party libraries. If they don't use type hint we cannot do anything here.

Code after:

def _remove_cached_modules(package_name: str) -> None:  # pragma: no cover
    to_remove = [mod for mod in sys.modules if mod.startswith(package_name)]
    # `del` is used instead of `reload()` because: If the new version of a module does not
    # define a name that was defined by the old version, the old definition remains.
    for module in to_remove:
        del sys.modules[module]


def _find_kedro_project(current_dir: Path) -> Any:  # pragma: no cover
    while current_dir != current_dir.parent:
        if _is_project(current_dir):
            return current_dir
        current_dir = current_dir.parent

    return None


def _guess_run_environment() -> str:  # pragma: no cover
    """Best effort to guess the IPython/Jupyter environment"""
    # https://github.com/microsoft/vscode-jupyter/issues/7380
    if os.environ.get("VSCODE_PID") or os.environ.get("VSCODE_CWD"):
        return "vscode"
    elif _is_databricks():
        return "databricks"
    elif hasattr(get_ipython(), "kernel"):  # type: ignore[no-untyped-call]
        # IPython terminal does not have this attribute
        return "jupyter"
    else:
        return "ipython"


@typing.no_type_check
@magic_arguments()
@argument(
    "node",
    type=str,
    help=("Name of the Node."),
    nargs="?",
    default=None,
)
def magic_load_node(args: str) -> None:
    """The line magic %load_node <node_name>
    Currently this feature has better supports with Jupyter Notebook (>7.0) and Jupyter Lab
    and VSCode Notebook. This line magic will generate code in multiple cells to load
    datasets from `DataCatalog`, import relevant functions and modules, node function
    definition and a function call. If generating code is not possible, it will print
    the code instead.
    """

    parameters = parse_argstring(magic_load_node, args)