Add support for %load_ext kedro (#4406)

* initial test draft Signed-off-by: ravi_kumar_pilla <[email protected]> * initial test draft Signed-off-by: ravi_kumar_pilla <[email protected]> * fix lint Signed-off-by: ravi_kumar_pilla <[email protected]> * update release note Signed-off-by: ravi_kumar_pilla <[email protected]> * update type for ipython Signed-off-by: ravi_kumar_pilla <[email protected]> * update type for ipython Signed-off-by: ravi_kumar_pilla <[email protected]> * fix lint Signed-off-by: ravi_kumar_pilla <[email protected]> * mypy ignore Signed-off-by: ravi_kumar_pilla <[email protected]> --------- Signed-off-by: ravi_kumar_pilla <[email protected]>
kedro-org · Jan 15, 2025 · 1ed504c · 1ed504c
1 parent a565d66
commit 1ed504c
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 4 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -4,6 +4,7 @@
 * Implemented `KedroDataCatalog.to_config()` method that converts the catalog instance into a configuration format suitable for serialization.
 * Improve OmegaConfigLoader performance.
 * Replaced `trufflehog` with `detect-secrets` for detecting secrets within a code base.
+* Added support for `%load_ext kedro`.
 
 ## Bug fixes and other changes
 * Added validation to ensure dataset versions consistency across catalog.

diff --git a/kedro/__init__.py b/kedro/__init__.py
@@ -3,6 +3,13 @@
 configuration and pipeline assembly.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from IPython.core.interactiveshell import InteractiveShell
+
 import sys
 import warnings
 
@@ -29,3 +36,9 @@ class KedroPythonVersionWarning(UserWarning):
 or set the PYTHONWARNINGS environment variable accordingly.""",
         KedroPythonVersionWarning,
     )
+
+
+def load_ipython_extension(ipython: InteractiveShell) -> None:
+    import kedro.ipython
+
+    kedro.ipython.load_ipython_extension(ipython)
diff --git a/kedro/ipython/__init__.py b/kedro/ipython/__init__.py
@@ -19,6 +19,8 @@
 if TYPE_CHECKING:
     from collections import OrderedDict
 
+    from IPython.core.interactiveshell import InteractiveShell
+
 from IPython.core.getipython import get_ipython
 from IPython.core.magic import needs_local_scope, register_line_magic
 from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
@@ -50,16 +52,16 @@
 RICH_INSTALLED = True if importlib.util.find_spec("rich") is not None else False
 
 
-def load_ipython_extension(ipython: Any) -> None:
+def load_ipython_extension(ipython: InteractiveShell) -> None:
     """
     Main entry point when %load_ext kedro.ipython is executed, either manually or
     automatically through `kedro ipython` or `kedro jupyter lab/notebook`.
     IPython will look for this function specifically.
     See https://ipython.readthedocs.io/en/stable/config/extensions/index.html
     """
-    ipython.register_magic_function(magic_reload_kedro, magic_name="reload_kedro")
+    ipython.register_magic_function(func=magic_reload_kedro, magic_name="reload_kedro")  # type: ignore[call-arg]
     logger.info("Registered line magic '%reload_kedro'")
-    ipython.register_magic_function(magic_load_node, magic_name="load_node")
+    ipython.register_magic_function(func=magic_load_node, magic_name="load_node")  # type: ignore[call-arg]
     logger.info("Registered line magic '%load_node'")
 
     if _find_kedro_project(Path.cwd()) is None:

diff --git a/pyproject.toml b/pyproject.toml
@@ -178,7 +178,8 @@ layers = [
 ]
 ignore_imports = [
     "kedro.runner.task -> kedro.framework.project",
-    "kedro.framework.hooks.specs -> kedro.framework.context"
+    "kedro.framework.hooks.specs -> kedro.framework.context",
+    "kedro -> kedro.ipython"
 ]
 
 [[tool.importlinter.contracts]]
@@ -188,6 +189,9 @@ modules = [
     "kedro.pipeline",
     "kedro.io"
 ]
+ignore_imports = [
+    "kedro -> kedro.ipython"
+]
 
 [[tool.importlinter.contracts]]
 name = "Config cannot import Runner et al"

diff --git a/tests/ipython/test_ipython.py b/tests/ipython/test_ipython.py
@@ -222,6 +222,17 @@ def test_line_magic_with_invalid_arguments(self, mocker, ipython):
         ):
             ipython.magic("reload_kedro --invalid_arg=dummy")
 
+    def test_ipython_kedro_extension_alias(self, mocker, ipython):
+        mock_ipython_extension = mocker.patch(
+            "kedro.ipython.load_ipython_extension", autospec=True
+        )
+        # Ensure that `kedro` is not loaded initially
+        assert "kedro" not in ipython.extension_manager.loaded
+        ipython.magic("load_ext kedro")
+        mock_ipython_extension.assert_called_once_with(ipython)
+        # Ensure that `kedro` extension has been loaded
+        assert "kedro" in ipython.extension_manager.loaded
+
 
 class TestProjectPathResolution:
     def test_only_path_specified(self):