Repository: kedro-org/kedro

Issue:
Revisit: Make KedroContext a dataclass and add `config_loader` as a property
# Description
As `KedroSession` now control the lifecycle of Kedro's run, `KedroContext` act like a container and it stores important attributes. Since we now dropped Python 3.6 support, we can make use of Python's `dataclass` to further tidy up `KedroContext`'s constructor code.

todo:
- [x] Make `KedroContext` a dataclass(https://github.com/kedro-org/kedro/pull/1465)
- [x] Add `config_loader` as @property instead of relying the private `context._get_config_loader()` https://github.com/kedro-org/kedro/pull/1505
- [x] Update corresponding document in Starter/Kedro



Discussion:
Turns out using the `dataclasss` doesn't seem to reduce any boilerplate code since we need to keep the attributes "private". You can find more details in this blog post https://florimond.dev/en/posts/2018/10/reconciling-dataclasses-and-properties-in-python/.

`attrs` support this but it's one more dependency to add. 
https://www.attrs.org/en/stable/examples.html



   
This was merged already, but as we approaching 0.19, I want to revise if it's a good idea to freeze `KedroContext`. 
Related Issues

We should also consider how easy for plugin developer to use `KedroContext`. In the past we were more negative about storing context and states in Hooks for following hooks. However in #2690, we may want to document this. If this is not a bad idea, do we need to re-visit the decision we made?
https://github.com/python/cpython/issues/93133#issuecomment-1660480431

Potentially helpful with the issues that using mock with dataclass or attr.
I cross-reference #3214 here which is likely a duplicate of this, sorry!

> # Description
> <!-- Is your feature request related to a problem? A clear and concise description of what the problem is: "I'm always frustrated when ..." -->
> 
> The ``KedroContext`` has been made a frozen dataclass in the develop branch since May 2022 and the PR #1465. 
> I think it should be unfrozen to keep te current behaviour and enable users to modify the context. 
> 
> The PR has triggered *a lot* of discussions (see in the issue), but it can be summarized as follows: 
> 1. if the class is frozen, it prevents users to modify the context with unintended side_effects, which has been made easier with the ``after_context_created``
> 2. if the class is *not* frozen, users can add attributes to the context, e.g. plugin configuration. The main advantage of this approach is that it makes the "interactive" workflow (i.e. within a notebook) consistent with the "command line" workflow (i.e. running with the kedro run command). 
> 
> Basically, we (@antonymilne, @noklam and myself) were not absolutely sure of what the best decision was, but @noklam pointed out that **we have a workaround** to force setting the attribute of a frozen dataclass by using ``context.__setattr__("attribute_name", attribute_value)`` **so we could have the benefits of both 1 & 2**. 
> 
> 
> ## Context
> <!-- Why is this change important to you? How would you use it? How can it benefit other users? -->
> 
> I am currently preparing the migration of ``kedro-mlflow`` to ``kedro=0.19``, and it turns out that **the workaround** identified back then **does _not_ work**. I can try do dark magic by digging deeply in python internals and kind of redefining ``__setattr__`` on the fly, but it looks very (very!) bad and I am not even sure I can make it work properly. 
> 
> Unless we find a workaround, **we are forced to choose between 1 & 2**. **My personal preference go to 2.** for the following reasons: 
> - Kedro has been open source for 4 years now, and I have never seen a bug due to a user involuntarily modifying the context, even during the last past year with the ``after_context_created`` hook. This obviously may happen, but we should acknowledge that the few users who modify the context do it "at their own risk" and know what they do.
> - I personally think (but I am biased as a plugin author :sweat_smile: ) that the consistency with the interactive workflow (not doing so will basically remove the mlflow configuration from context and make it hard to use in notebooks) and the simplified experience for plugin authors is worth letting them modifying the class. 
> 
> ## Possible Implementation
> <!-- (Optional) Suggest an idea for implementing the addition or change. -->
> 
> Replace ``@frozen`` decorator by ``@define(slots=False)``. The ``@slots=False`` is needed to enable adding new attributes. 
> 
> https://github.com/kedro-org/kedro/blob/0c31dd1560ce7426b89d3d06d76ee8073ba6084e/kedro/framework/context/context.py#L161-L164
> 
> ## Possible Alternatives
> 
> Keep it frozen. 
> 
Thanks for x-ref the issues. The issue has been 1.5 years old and my thinking has changed a little bit as well. I agree with you mostly because:
1. It has not been causing any issue since the hook is added 1.5 years ago.
2. Stateful hook is more common in Kedro now, so users are likely to bump into this problem often.

Basically, I think we can take a middle-ground which is Mutable class + immutable attributes. The key is to use `attrs.setters.frozen`, I mentioned that in https://github.com/kedro-org/kedro/pull/1465/ but honestly I don't remember why I did not use it but go for a full `frozen` class. Would this be sufficient?

```python
import attr
from attrs.setters import frozen

@attr.s
class Foo:
  # A partial frozen class
  bar: int = attr.field(default=0, on_setattr=frozen)

f = Foo()
f.abc = 3 # It's ok to mutate the class
f.bar = 3 # It's not ok to mutate existing attribtue
```


Is this possible with plain dataclass? (To avoid the extra dependency if possible)
## Choosing between ``attrs`` and ``dataclasses``

In the original issue https://github.com/kedro-org/kedro/pull/1465, there are detailed analysis of the pros and cons and justifications to choose over ``attrs``. The conclusion **seems to be that it is much easier to make the class "partially" frozen** (e.g. some attributes are frozen and other are not,namely the one prefixed with an underscore  like ``_hook_manager``) with ``attrs`` rather than with ``dataclasses`` (even if it was possible too with a lot of boilerplate code).  

Since we finally ended up making the entire class frozen, I am not sure why this was relevant in the end, but I may miss someething. This **needs to be reassesed once we figure out if we really want to make the class or some of its attributes frozen**. 

## Should we make some fields frozen? 

If I understand correctly, this was the original motivation of the PR, but I think the following arguments still stands: 

> It has not been causing any issue since the hook is added 1.5 years ago.

So I would be in favor of not freezing anything and use a pure dataclass. This would also helps to get rid of ``attrs`` (the motivation for thsi package was to make only some attributes frozen)  and one less dependency, many less future maintenance problems :)

## Workaround to add fields on a frozen dataclass

I've played a bit around this, and here are some food for thoughts. In a nutshell: 

|Package     |Override existing attribute| Add extra attribute|
|-------------|-----------------------------|---------------------|
|``dataclass``|:white_check_mark:            | :white_check_mark:|
|``attrs``       |:white_check_mark:             | :x:|

This implies that **if we switch to ``dataclass`` instead of attrs, we can keep the exact current behaviour**, but the price is to have a boilerplate code as described in the first paragraph.

#### frozen ``dataclass`` can be overriden and added

```python
from dataclasses import dataclass

@dataclass(frozen=True)
class PointDataClass:
    x: int
    y: int

point_data_class = PointDataClass(x=1, y=1)
point_data_class.z = 2 # # FrozenInstanceError: cannot assign to field 'z'
object.__setattr__(point_data_class, "y", 4)  # works: you can override an existing field
print(point_data_class.y)
point_data_class.__setattr__("z", 2)  # self.__setattr__ does not work: FrozenInstanceError: cannot assign to field 'z'
object.__setattr__(point_data_class, "z", 2)  # object.__setattr__ works! You can add an attribute
print(point_data_class.z) # 2
```

#### frozen ``attrs`` attributes can be overriden and but not added

```python
from attrs import frozen

@frozen
class PointAttrs:
    x: int
    y: int

point_attrs = PointAttrs(x=1, y=1)
point_attrs.z = 2
point_attrs.__setattr__("z", 2)  # self.__setattr__ does not work: FrozenInstanceError
object.__setattr__(point_attrs, "y", 4)  # object.__setattr__ works : you can override an existing field
print(point_attrs.y) # 4
object.__setattr__(point_attrs, "z", 2)  #  # object.__setattr__ does NOT work! You CANNOT add an attribute: AttributeError: 'PointAttrs' object has no attribute 'z'
```
@astrojuanlu `attrs` aren't really "extra" because it was always a secondary dependencies used by other library. Partial immutability is only possible with `attr`.

So I think the key here is:
- Should we make it partially immutable or just keep everything mutable?

[The example above ](https://github.com/kedro-org/kedro/issues/1459#issuecomment-1783944593)shows that we can make it works for partial immutability. In my mind `attr` is a superset of `dataclass` so I don't expect things that only work with `dataclasses`

Whether or not we should use `attr` is an implementation decision, it's not necessary an extra dependencies because it is already used by other core dependencies (maybe we can double check),

Resolving patch:
diff --git a/kedro/framework/context/context.py b/kedro/framework/context/context.py
--- a/kedro/framework/context/context.py
+++ b/kedro/framework/context/context.py
@@ -8,7 +8,7 @@
 from urllib.parse import urlparse
 from warnings import warn
 
-from attrs import field, frozen
+from attrs import define, field
 from omegaconf import OmegaConf
 from pluggy import PluginManager
 
@@ -142,18 +142,38 @@ def _expand_full_path(project_path: str | Path) -> Path:
     return Path(project_path).expanduser().resolve()
 
 
-@frozen
+@define(slots=False)  # Enable setting new attributes to `KedroContext`
 class KedroContext:
     """``KedroContext`` is the base class which holds the configuration and
     Kedro's main functionality.
+
+    Create a context object by providing the root of a Kedro project and
+    the environment configuration subfolders (see ``kedro.config.OmegaConfigLoader``)
+    Raises:
+        KedroContextError: If there is a mismatch
+            between Kedro project version and package version.
+    Args:
+        project_path: Project path to define the context for.
+        config_loader: Kedro's ``OmegaConfigLoader`` for loading the configuration files.
+        env: Optional argument for configuration default environment to be used
+            for running the pipeline. If not specified, it defaults to "local".
+        package_name: Package name for the Kedro project the context is
+            created for.
+        hook_manager: The ``PluginManager`` to activate hooks, supplied by the session.
+        extra_params: Optional dictionary containing extra project parameters.
+            If specified, will update (and therefore take precedence over)
+            the parameters retrieved from the project configuration.
+
     """
 
-    _package_name: str
-    project_path: Path = field(converter=_expand_full_path)
-    config_loader: AbstractConfigLoader
-    _hook_manager: PluginManager
-    env: str | None = None
-    _extra_params: dict[str, Any] | None = field(default=None, converter=deepcopy)
+    project_path: Path = field(init=True, converter=_expand_full_path)
+    config_loader: AbstractConfigLoader = field(init=True)
+    env: str | None = field(init=True)
+    _package_name: str = field(init=True)
+    _hook_manager: PluginManager = field(init=True)
+    _extra_params: dict[str, Any] | None = field(
+        init=True, default=None, converter=deepcopy
+    )
 
     @property
     def catalog(self) -> DataCatalog:


Test patch:
diff --git a/tests/framework/context/test_context.py b/tests/framework/context/test_context.py
--- a/tests/framework/context/test_context.py
+++ b/tests/framework/context/test_context.py
@@ -12,7 +12,6 @@
 import pytest
 import toml
 import yaml
-from attrs.exceptions import FrozenInstanceError
 from pandas.testing import assert_frame_equal
 
 from kedro import __version__ as kedro_version
@@ -201,11 +200,11 @@ def dummy_context(tmp_path, prepare_project_dir, env, extra_params):
     context = session.load_context()
     config_loader = context.config_loader
     context = KedroContext(
-        MOCK_PACKAGE_NAME,
-        str(tmp_path),
+        project_path=str(tmp_path),
         config_loader=config_loader,
-        hook_manager=_create_hook_manager(),
         env=env,
+        package_name=MOCK_PACKAGE_NAME,
+        hook_manager=_create_hook_manager(),
         extra_params=extra_params,
     )
 
@@ -218,9 +217,9 @@ def test_attributes(self, tmp_path, dummy_context):
         assert isinstance(dummy_context.project_path, Path)
         assert dummy_context.project_path == tmp_path.resolve()
 
-    def test_immutable_instance(self, dummy_context):
-        with pytest.raises(FrozenInstanceError):
-            dummy_context.catalog = 1
+    def test_set_new_attribute(self, dummy_context):
+        dummy_context.mlflow = 1
+        assert dummy_context.mlflow == 1
 
     def test_get_catalog_always_using_absolute_path(self, dummy_context):
         config_loader = dummy_context.config_loader
diff --git a/tests/framework/session/test_session.py b/tests/framework/session/test_session.py
--- a/tests/framework/session/test_session.py
+++ b/tests/framework/session/test_session.py
@@ -1,9 +1,12 @@
 import logging
 import re
 import subprocess
+import sys
 import textwrap
 from collections.abc import Mapping
 from pathlib import Path
+from typing import Any, Type
+from unittest.mock import create_autospec
 
 import pytest
 import toml
@@ -13,7 +16,6 @@
 from kedro import __version__ as kedro_version
 from kedro.config import AbstractConfigLoader, OmegaConfigLoader
 from kedro.framework.cli.utils import _split_params
-from kedro.framework.context import KedroContext
 from kedro.framework.project import (
     LOGGING,
     ValidationError,
@@ -23,7 +25,7 @@
     _ProjectSettings,
 )
 from kedro.framework.session import KedroSession
-from kedro.framework.session.session import KedroSessionError
+from kedro.framework.session.session import KedroContext, KedroSessionError
 from kedro.framework.session.shelvestore import ShelveStore
 from kedro.framework.session.store import BaseSessionStore
 
@@ -43,6 +45,36 @@ class BadConfigLoader:
     """
 
 
+ATTRS_ATTRIBUTE = "__attrs_attrs__"
+
+NEW_TYPING = sys.version_info[:3] >= (3, 7, 0)  # PEP 560
+
+
+def create_attrs_autospec(spec: Type, spec_set: bool = True) -> Any:
+    """Creates a mock of an attr class (creates mocks recursively on all attributes).
+    https://github.com/python-attrs/attrs/issues/462#issuecomment-1134656377
+
+    :param spec: the spec to mock
+    :param spec_set: if True, AttributeError will be raised if an attribute that is not in the spec is set.
+    """
+
+    if not hasattr(spec, ATTRS_ATTRIBUTE):
+        raise TypeError(f"{spec!r} is not an attrs class")
+    mock = create_autospec(spec, spec_set=spec_set)
+    for attribute in getattr(spec, ATTRS_ATTRIBUTE):
+        attribute_type = attribute.type
+        if NEW_TYPING:
+            # A[T] does not get a copy of __dict__ from A(Generic[T]) anymore, use __origin__ to get it
+            while hasattr(attribute_type, "__origin__"):
+                attribute_type = attribute_type.__origin__
+        if hasattr(attribute_type, ATTRS_ATTRIBUTE):
+            mock_attribute = create_attrs_autospec(attribute_type, spec_set)
+        else:
+            mock_attribute = create_autospec(attribute_type, spec_set=spec_set)
+        object.__setattr__(mock, attribute.name, mock_attribute)
+    return mock
+
+
 @pytest.fixture
 def mock_runner(mocker):
     mock_runner = mocker.patch(
@@ -55,7 +87,12 @@ def mock_runner(mocker):
 
 @pytest.fixture
 def mock_context_class(mocker):
-    return mocker.patch("kedro.framework.session.session.KedroContext", autospec=True)
+    mock_cls = create_attrs_autospec(KedroContext)
+    return mocker.patch(
+        "kedro.framework.session.session.KedroContext",
+        autospec=True,
+        return_value=mock_cls,
+    )
 
 
 def _mock_imported_settings_paths(mocker, mock_settings):
@@ -75,6 +112,7 @@ def mock_settings(mocker):
 @pytest.fixture
 def mock_settings_context_class(mocker, mock_context_class):
     class MockSettings(_ProjectSettings):
+        # dynaconf automatically deleted some attribute when the class is MagicMock
         _CONTEXT_CLASS = Validator(
             "CONTEXT_CLASS", default=lambda *_: mock_context_class
         )