Refactor KEDRO_ENV (#447)

kedro/context/__init__.py

@@ -30,9 +30,4 @@
 project context.
 """
 
-from .context import (  # NOQA
-    KEDRO_ENV_VAR,
-    KedroContext,
-    KedroContextError,
-    load_context,
-)
+from .context import KedroContext, KedroContextError, load_context  # NOQA

kedro/context/context.py

@@ -48,8 +48,6 @@
 from kedro.utils import load_obj
 from kedro.versioning import Journal
 
-KEDRO_ENV_VAR = "KEDRO_ENV"
-
 
 def _version_mismatch_error(context_version) -> str:
     return (
@@ -562,6 +560,10 @@ def load_context(project_path: Union[str, Path], **kwargs) -> KedroContext:
         )
         os.chdir(str(project_path))  # Move to project root
 
+    # update kwargs with env from the environment variable (defaults to None if not set)
+    # need to do this because some CLI command (e.g `kedro run`) defaults to passing in `env=None`
+    kwargs["env"] = kwargs.get("env") or os.getenv("KEDRO_ENV")
+
     # Instantiate the context after changing the cwd for logging to be properly configured.
     context = context_class(project_path, **kwargs)
     return context

kedro/template/{{ cookiecutter.repo_name }}/.ipython/profile_default/startup/00-kedro-init.py

@@ -1,5 +1,4 @@
 import logging.config
-import os
 import sys
 from pathlib import Path
 
@@ -19,7 +18,7 @@ def reload_kedro(path, line=None):
 
     try:
         import kedro.config.default_logger
-        from kedro.context import KEDRO_ENV_VAR, load_context
+        from kedro.context import load_context
         from kedro.cli.jupyter import collect_line_magic
     except ImportError:
         logging.error(
@@ -31,8 +30,7 @@ def reload_kedro(path, line=None):
     try:
         path = path or project_path
         logging.debug("Loading the context from %s", str(path))
-
-        context = load_context(path, env=os.getenv(KEDRO_ENV_VAR))
+        context = load_context(path)
         catalog = context.catalog
 
         # remove cached user modules

kedro/template/{{ cookiecutter.repo_name }}/kedro_cli.py

@@ -51,7 +51,7 @@
     forward_command,
     python_call,
 )
-from kedro.context import KEDRO_ENV_VAR, load_context
+from kedro.context import load_context
 from kedro.runner import SequentialRunner
 from kedro.utils import load_obj
 
@@ -216,7 +216,6 @@ def cli():
     type=str,
     default=None,
     multiple=False,
-    envvar=KEDRO_ENV_VAR,
     help=ENV_ARG_HELP,
 )
 @click.option("--tag", "-t", type=str, multiple=True, help=TAG_ARG_HELP)
@@ -456,12 +455,12 @@ def _build_jupyter_command(
 def _build_jupyter_env(kedro_env: str) -> Dict[str, Any]:
     """Build the environment dictionary that gets injected into the subprocess running
     Jupyter. Since the subprocess has access only to the environment variables passed
-    in, we need to copy the current environment and add ``KEDRO_ENV_VAR``.
+    in, we need to copy the current environment and add ``KEDRO_ENV``.
     """
     if not kedro_env:
         return {}
     jupyter_env = os.environ.copy()
-    jupyter_env[KEDRO_ENV_VAR] = kedro_env
+    jupyter_env["KEDRO_ENV"] = kedro_env
     return {"env": jupyter_env}
 
 
@@ -484,7 +483,6 @@ def jupyter():
     type=str,
     default=None,
     multiple=False,
-    envvar=KEDRO_ENV_VAR,
     help=ENV_ARG_HELP,
 )
 def jupyter_notebook(ip, all_kernels, env, idle_timeout, args):
@@ -512,7 +510,6 @@ def jupyter_notebook(ip, all_kernels, env, idle_timeout, args):
     type=str,
     default=None,
     multiple=False,
-    envvar=KEDRO_ENV_VAR,
     help=ENV_ARG_HELP,
 )
 def jupyter_lab(ip, all_kernels, env, idle_timeout, args):

tests/template/test_kedro_cli.py

@@ -37,7 +37,6 @@
 import pytest
 from click.testing import CliRunner
 
-from kedro.context import KEDRO_ENV_VAR
 from kedro.runner import ParallelRunner, SequentialRunner
 
 
@@ -238,7 +237,10 @@ def test_run_with_config(
             ({}, {}),
             ({"params": {"foo": "baz"}}, {"foo": "baz"}),
             ({"params": "foo:baz"}, {"foo": "baz"}),
-            ({"params": {"foo": "123.45", "baz": "678", "bar": 9}}, {"foo": "123.45", "baz": "678", "bar": 9}),
+            (
+                {"params": {"foo": "123.45", "baz": "678", "bar": 9}},
+                {"foo": "123.45", "baz": "678", "bar": 9},
+            ),
         ],
         indirect=["fake_run_config_with_params"],
     )
@@ -268,17 +270,6 @@ def test_run_with_params_in_config(
             Path.cwd(), env=mocker.ANY, extra_params=expected
         )
 
-    def test_run_env_environment_var(
-        self, fake_kedro_cli, fake_load_context, fake_repo_path, monkeypatch, mocker
-    ):
-        monkeypatch.setenv("KEDRO_ENV", "my_special_env")
-        result = CliRunner().invoke(fake_kedro_cli.cli, ["run"])
-        assert not result.exit_code
-
-        fake_load_context.assert_called_once_with(
-            Path.cwd(), env="my_special_env", extra_params=mocker.ANY
-        )
-
     @pytest.mark.parametrize(
         "cli_arg,expected_extra_params",
         [
@@ -614,6 +605,7 @@ def test_help(self, help_flag, fake_kedro_cli, fake_ipython_message):
     def test_env(
         self, env_flag, fake_kedro_cli, python_call_mock, default_jupyter_options
     ):
+        """This tests passing an environment variable to the jupyter subprocess."""
         result = CliRunner().invoke(
             fake_kedro_cli.cli, ["jupyter", "notebook", env_flag, "my_special_env"]
         )
@@ -622,21 +614,8 @@ def test_env(
         args, kwargs = python_call_mock.call_args
         assert args == default_jupyter_options
         assert "env" in kwargs
-        assert KEDRO_ENV_VAR in kwargs["env"]
-        assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env"
-
-    def test_env_environment_variable(
-        self, fake_kedro_cli, python_call_mock, monkeypatch, default_jupyter_options
-    ):
-        monkeypatch.setenv("KEDRO_ENV", "my_special_env")
-        result = CliRunner().invoke(fake_kedro_cli.cli, ["jupyter", "notebook"])
-        assert not result.exit_code
-
-        args, kwargs = python_call_mock.call_args
-        assert args == default_jupyter_options
-        assert "env" in kwargs
-        assert KEDRO_ENV_VAR in kwargs["env"]
-        assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env"
+        assert "KEDRO_ENV" in kwargs["env"]
+        assert kwargs["env"]["KEDRO_ENV"] == "my_special_env"
 
 
 class TestJupyterLabCommand:
@@ -707,6 +686,7 @@ def test_help(self, help_flag, fake_kedro_cli, fake_ipython_message):
     def test_env(
         self, env_flag, fake_kedro_cli, python_call_mock, default_jupyter_options
     ):
+        """This tests passing an environment variable to the jupyter subprocess."""
         result = CliRunner().invoke(
             fake_kedro_cli.cli, ["jupyter", "lab", env_flag, "my_special_env"]
         )
@@ -715,21 +695,8 @@ def test_env(
         args, kwargs = python_call_mock.call_args
         assert args == default_jupyter_options
         assert "env" in kwargs
-        assert KEDRO_ENV_VAR in kwargs["env"]
-        assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env"
-
-    def test_env_environment_variable(
-        self, fake_kedro_cli, python_call_mock, monkeypatch, default_jupyter_options
-    ):
-        monkeypatch.setenv("KEDRO_ENV", "my_special_env")
-        result = CliRunner().invoke(fake_kedro_cli.cli, ["jupyter", "lab"])
-        assert not result.exit_code
-
-        args, kwargs = python_call_mock.call_args
-        assert args == default_jupyter_options
-        assert "env" in kwargs
-        assert KEDRO_ENV_VAR in kwargs["env"]
-        assert kwargs["env"][KEDRO_ENV_VAR] == "my_special_env"
+        assert "KEDRO_ENV" in kwargs["env"]
+        assert kwargs["env"]["KEDRO_ENV"] == "my_special_env"
 
 
 class TestConvertNotebookCommand:

tests/template/test_load_context.py

@@ -47,6 +47,16 @@ def test_valid_context(self, mocker, fake_repo_path):
         assert str(fake_repo_path.resolve() / "src") in sys.path
         assert os.getcwd() == str(fake_repo_path.resolve())
 
+    def test_valid_context_with_env(self, mocker, monkeypatch, fake_repo_path):
+        """Test getting project context when Kedro config environment is specified in the environment variable."""
+        # Disable logging.config.dictConfig in KedroContext._setup_logging as
+        # it changes logging.config and affects other unit tests
+        mocker.patch("logging.config.dictConfig")
+        mocker.patch("kedro.config.config.ConfigLoader.get")
+        monkeypatch.setenv("KEDRO_ENV", "my_fake_env")
+        result = load_context(str(fake_repo_path))
+        assert result.env == "my_fake_env"
+
     def test_invalid_path(self, tmp_path):
         """Test for loading context from an invalid path. """
         other_path = tmp_path / "other"