Repository: kedro-org/kedro

Issue:
Render concrete parameter/return types for dataset `save`/`load`
## Description
The parameter type for a dataset's `save` is `_DI`, and the return type for `load` is `_DO`. This doesn't make any sense to a user (and barely makes sense to me, as a Kedro developer, because I saw the typing PR go in and have that context). I'd rather it not display any types, if getting the correct concrete types is too difficult.

<img width="1436" alt="image" src="https://user-images.githubusercontent.com/14007150/211978326-bd1cced3-03d3-4f4b-bc88-4c48269bf159.png">


Discussion:
This would indeed be great to fix but I think only someone very brave should attempt it as I can see it being very frustrating and difficult to achieve... Do you have any idea how you could actually get Sphinx to do that without us having to manually override every single page of automatically generated docs for datasets? Possibly we can edit a higher-level rst template to do it?
To do: Try updating to the newest version of `sphinx` and see if it has improved. 
We are using a later version of Sphinx and this is still not resolved.

It may be something that @astrojuanlu has familiarity with but if it's a non-trivial fix I suggest we put it as a very low priority. Could we potentially have an explanation of what is going on somewhere so if people want to know what it means, they can search and find the reasoning?
Upgrading sphinx or sphinx-autodoc-typehints proved to create lots of headaches in the past. I agree the current result is less than ideal but I don't think this is trivial to fix at all. It might be easier to actually redesign the datasets...
It doesn't look much better on `latest` with the spun-off `kedro_datasets` https://docs.kedro.org/en/stable/kedro_datasets.pandas.CSVDataSet.html

![image](https://github.com/kedro-org/kedro/assets/316517/a76473f3-e52e-47ab-bd5f-d2c5ded69b47)

The root cause is that the heavy lifting is done by `_load`, which is a private method. Conversation about this is scattered around different issues (#1778, #1936, #2409 to name a few) but essentially I think the only sane way to fix this is to come up with a design that favours composition over inheritance, something that looks more similar to the proposal in https://github.com/kedro-org/kedro/issues/1936#issue-1410411315.
I think it's easier to just wait on a redesign of our catalog and datasets #1778.
# Desired output

We can redefine `load` and `save` in each class to produce our desired result:

```python
from pandas.util._decorators import doc  # We could roll our own, simpler version


class LambdaDataset(AbstractDataset):
    ...

    @doc(AbstractDataset.save.__doc__)
    def save(self, data: int) -> None:
        super().save(data)
 
    @doc(AbstractDataset.load.__doc__)
    def load(self) -> pd.DataFrame:
        return super().load()
```

<img width="759" alt="image" src="https://github.com/kedro-org/kedro/assets/14007150/4e533946-d9e2-4cdc-9409-de12bc30566b">

Note that clicking through to the code on the docs will show that generated code, so maybe we should make it explicit that this is a generated function in the comments:

<img width="910" alt="image" src="https://github.com/kedro-org/kedro/assets/14007150/f67c901b-4796-425d-a4c5-1c46afcd87cd">

# Possible implementations

Now, the question is—how do we accomplish this for all datasets? It is possible to either:

1. Actually add the `load` and `save` methods to each dataset implementation. There is probably a negligible performance hit from the added function call (since I/O is almost certainly the dominant part of any reasonable dataset). It's not the prettiest from a code perspective, though, where you'd rather rely on inheritance.
2. You can generate the additional methods just for docs generation.
3. Is there another way to modify classes parsed by Sphinx before docs generation?

# Notes

* I tried adding a `.pyi` file, but that didn't seem to have any effect. Can't figure out for certain whether `sphinx_autodoc_typehints` looks at them (I don't think so).
* https://github.com/tox-dev/sphinx-autodoc-typehints/issues/296 fixed `typing.overload` support, but is only available from `sphinx_autodoc_typehints==1.21.2`. While we probably shouldn't hard-pin `sphinx_autodoc_typehints` in `kedro-sphinx-theme`, and I personally am not sure what regression was the issue/if it may have been fixed in a future version (`sphinx_autodoc_typehints` is on 2.x now), the specific issue resolved re `typing.overload` (duplicate entries) does not seem to be the problem here.
* Upon further investigation, maybe `typing.overload` doesn't have any value here to begin with, as you must specify the implementation after the `typing.overload`. At that point, we may as well just provide the implementation only in our case.
We discussed this on tech design today (2024-05-08). While there weren't clear arguments against this solution, we discussed alternative approaches:

- Given that this would mainly be the presentation layer, maybe it would be worth exploring a Sphinx extension instead.
  - However, creating Sphinx extensions can be an _interesting_ exercise.
  - @deepyaman wasn't particularly motivated to pursue this approach but didn't close the door on somebody else doing it.
  - This is probably something that wouldn't be useful to the broader community, given that this is a very specific situation caused by the design.
- We discussed about the datasets design (xref #1778)
  - Of course we acknowledged that it will take us ~months to get there, so it would be nice to have an interim solution
  - We started discussing whether we could get rid of the custom error handling in `AbstractDataSet` (see below) given that it's directly affecting this issue _and_ it's making it difficult for users to debug their dataset errors https://github.com/kedro-org/kedro/issues/1936#issuecomment-1727172650 but we didn't have enough time to discuss

https://github.com/kedro-org/kedro/blob/c8a0e22464f2d229e26ec7fc138436c8863e3dcc/kedro/io/core.py#L190-L202
> We discussed this on tech design today (2024-05-08). While there weren't clear arguments against this solution, we discussed alternative approaches:
> 
> * Given that this would mainly be the presentation layer, maybe it would be worth exploring a Sphinx extension instead.
>   
>   * However, creating Sphinx extensions can be an _interesting_ exercise.
>   * @deepyaman wasn't particularly motivated to pursue this approach but didn't close the door on somebody else doing it.
>   * This is probably something that wouldn't be useful to the broader community, given that this is a very specific situation caused by the design.
> * We discussed about the datasets design (xref [Re-design io.core and io.data_catalog #1778](https://github.com/kedro-org/kedro/issues/1778))
>   
>   * Of course we acknowledged that it will take us ~months to get there, so it would be nice to have an interim solution
>   * We started discussing whether we could get rid of the custom error handling in `AbstractDataSet` (see below) given that it's directly affecting this issue _and_ it's making it difficult for users to debug their dataset errors [Easier CustomDataset Creation #1936 (comment)](https://github.com/kedro-org/kedro/issues/1936#issuecomment-1727172650) but we didn't have enough time to discuss
> 
> https://github.com/kedro-org/kedro/blob/c8a0e22464f2d229e26ec7fc138436c8863e3dcc/kedro/io/core.py#L190-L202

Thanks @astrojuanlu! Just copying in my summary, but I think you covered it:

> It sounds like overriding `load`/`save` may be workable, there’s a potential weekend project if somebody wants to write their own Sphinx plugin, but maybe the first followup is to see whether we can get rid of `_load` and `_save` altogether?

I have some ideas on the latter, and I will investigate this.



Patch:
diff --git a/RELEASE.md b/RELEASE.md
index f3eeec104d..32a47371f6 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,6 +1,7 @@
 # Upcoming Release 0.19.7
 
 ## Major features and improvements
+* Exposed `load` and `save` publicly for each dataset in the core `kedro` library, and enabled other datasets to do the same. If a dataset doesn't expose `load` or `save` publicly, Kedro will fall back to using `_load` or `_save`, respectively.
 * Kedro commands are now lazily loaded to add performance gains when running Kedro commands.
 
 ## Bug fixes and other changes
diff --git a/docs/source/data/how_to_create_a_custom_dataset.md b/docs/source/data/how_to_create_a_custom_dataset.md
index 5ebe50a8c5..01ad199f55 100644
--- a/docs/source/data/how_to_create_a_custom_dataset.md
+++ b/docs/source/data/how_to_create_a_custom_dataset.md
@@ -312,7 +312,7 @@ To add versioning support to the new dataset we need to extend the
  {py:class}`~kedro.io.AbstractVersionedDataset` to:
 
 * Accept a `version` keyword argument as part of the constructor
-* Adapt the `_save` and `_load` method to use the versioned data path obtained from `_get_save_path` and `_get_load_path` respectively
+* Adapt the `_load` and `_save` method to use the versioned data path obtained from `_get_load_path` and `_get_save_path` respectively
 
 The following amends the full implementation of our basic `ImageDataset`. It now loads and saves data to and from a versioned subfolder (`data/01_raw/pokemon-images-and-types/images/images/pikachu.png/<version>/pikachu.png` with `version` being a datetime-formatted string `YYYY-MM-DDThh.mm.ss.sssZ` by default):
 
diff --git a/kedro/io/cached_dataset.py b/kedro/io/cached_dataset.py
index 2ca0bac28b..c4424c34ef 100644
--- a/kedro/io/cached_dataset.py
+++ b/kedro/io/cached_dataset.py
@@ -71,7 +71,7 @@ def __init__(
                 "The argument type of 'dataset' should be either a dict/YAML "
                 "representation of the dataset, or the actual dataset object."
             )
-        self._cache = MemoryDataset(copy_mode=copy_mode)
+        self._cache = MemoryDataset(copy_mode=copy_mode)  # type: ignore[abstract]
         self.metadata = metadata
 
     def _release(self) -> None:
diff --git a/kedro/io/core.py b/kedro/io/core.py
index 53c2000e35..fc48e536e3 100644
--- a/kedro/io/core.py
+++ b/kedro/io/core.py
@@ -13,7 +13,7 @@
 import warnings
 from collections import namedtuple
 from datetime import datetime, timezone
-from functools import partial
+from functools import partial, wraps
 from glob import iglob
 from operator import attrgetter
 from pathlib import Path, PurePath, PurePosixPath
@@ -22,6 +22,7 @@
 
 from cachetools import Cache, cachedmethod
 from cachetools.keys import hashkey
+from typing_extensions import Self
 
 from kedro.utils import load_obj
 
@@ -74,6 +75,7 @@ class VersionNotFoundError(DatasetError):
 
 class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
     """``AbstractDataset`` is the base class for all data set implementations.
+
     All data set implementations should extend this abstract class
     and implement the methods marked as abstract.
     If a specific dataset implementation cannot be used in conjunction with
@@ -93,10 +95,10 @@ class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
         >>>         self._param1 = param1
         >>>         self._param2 = param2
         >>>
-        >>>     def _load(self) -> pd.DataFrame:
+        >>>     def load(self) -> pd.DataFrame:
         >>>         return pd.read_csv(self._filepath)
         >>>
-        >>>     def _save(self, df: pd.DataFrame) -> None:
+        >>>     def save(self, df: pd.DataFrame) -> None:
         >>>         df.to_csv(str(self._filepath))
         >>>
         >>>     def _exists(self) -> bool:
@@ -178,57 +180,6 @@ def from_config(
     def _logger(self) -> logging.Logger:
         return logging.getLogger(__name__)
 
-    def load(self) -> _DO:
-        """Loads data by delegation to the provided load method.
-
-        Returns:
-            Data returned by the provided load method.
-
-        Raises:
-            DatasetError: When underlying load method raises error.
-
-        """
-
-        self._logger.debug("Loading %s", str(self))
-
-        try:
-            return self._load()
-        except DatasetError:
-            raise
-        except Exception as exc:
-            # This exception handling is by design as the composed data sets
-            # can throw any type of exception.
-            message = (
-                f"Failed while loading data from data set {str(self)}.\n{str(exc)}"
-            )
-            raise DatasetError(message) from exc
-
-    def save(self, data: _DI) -> None:
-        """Saves data by delegation to the provided save method.
-
-        Args:
-            data: the value to be saved by provided save method.
-
-        Raises:
-            DatasetError: when underlying save method raises error.
-            FileNotFoundError: when save method got file instead of dir, on Windows.
-            NotADirectoryError: when save method got file instead of dir, on Unix.
-        """
-
-        if data is None:
-            raise DatasetError("Saving 'None' to a 'Dataset' is not allowed")
-
-        try:
-            self._logger.debug("Saving %s", str(self))
-            self._save(data)
-        except DatasetError:
-            raise
-        except (FileNotFoundError, NotADirectoryError):
-            raise
-        except Exception as exc:
-            message = f"Failed while saving data to data set {str(self)}.\n{str(exc)}"
-            raise DatasetError(message) from exc
-
     def __str__(self) -> str:
         # TODO: Replace with __repr__ implementation in 0.20.0 release.
         def _to_str(obj: Any, is_root: bool = False) -> str:
@@ -257,6 +208,85 @@ def _to_str(obj: Any, is_root: bool = False) -> str:
 
         return f"{type(self).__name__}({_to_str(self._describe(), True)})"
 
+    @classmethod
+    def _load_wrapper(cls, load_func: Callable[[Self], _DO]) -> Callable[[Self], _DO]:
+        """Decorate `load_func` with logging and error handling code."""
+
+        @wraps(load_func)
+        def load(self: Self) -> _DO:
+            self._logger.debug("Loading %s", str(self))
+
+            try:
+                return load_func(self)
+            except DatasetError:
+                raise
+            except Exception as exc:
+                # This exception handling is by design as the composed data sets
+                # can throw any type of exception.
+                message = (
+                    f"Failed while loading data from data set {str(self)}.\n{str(exc)}"
+                )
+                raise DatasetError(message) from exc
+
+        load.__annotations__["return"] = load_func.__annotations__.get("return")
+        load.__loadwrapped__ = True  # type: ignore[attr-defined]
+        return load
+
+    @classmethod
+    def _save_wrapper(
+        cls, save_func: Callable[[Self, _DI], None]
+    ) -> Callable[[Self, _DI], None]:
+        """Decorate `save_func` with logging and error handling code."""
+
+        @wraps(save_func)
+        def save(self: Self, data: _DI) -> None:
+            if data is None:
+                raise DatasetError("Saving 'None' to a 'Dataset' is not allowed")
+
+            try:
+                self._logger.debug("Saving %s", str(self))
+                save_func(self, data)
+            except (DatasetError, FileNotFoundError, NotADirectoryError):
+                raise
+            except Exception as exc:
+                message = (
+                    f"Failed while saving data to data set {str(self)}.\n{str(exc)}"
+                )
+                raise DatasetError(message) from exc
+
+        save.__annotations__["data"] = save_func.__annotations__.get("data", Any)
+        save.__annotations__["return"] = save_func.__annotations__.get("return")
+        save.__savewrapped__ = True  # type: ignore[attr-defined]
+        return save
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Decorate the `load` and `save` methods provided by the class.
+
+        If `_load` or `_save` are defined, alias them as a prerequisite.
+
+        """
+        super().__init_subclass__(**kwargs)
+
+        if hasattr(cls, "_load") and not cls._load.__qualname__.startswith("Abstract"):
+            cls.load = cls._load  # type: ignore[method-assign]
+
+        if hasattr(cls, "_save") and not cls._save.__qualname__.startswith("Abstract"):
+            cls.save = cls._save  # type: ignore[method-assign]
+
+        if hasattr(cls, "load") and not cls.load.__qualname__.startswith("Abstract"):
+            cls.load = cls._load_wrapper(  # type: ignore[assignment]
+                cls.load
+                if not getattr(cls.load, "__loadwrapped__", False)
+                else cls.load.__wrapped__  # type: ignore[attr-defined]
+            )
+
+        if hasattr(cls, "save") and not cls.save.__qualname__.startswith("Abstract"):
+            cls.save = cls._save_wrapper(  # type: ignore[assignment]
+                cls.save
+                if not getattr(cls.save, "__savewrapped__", False)
+                else cls.save.__wrapped__  # type: ignore[attr-defined]
+            )
+
     def _pretty_repr(self, object_description: dict[str, Any]) -> str:
         str_keys = []
         for arg_name, arg_descr in object_description.items():
@@ -276,17 +306,37 @@ def __repr__(self) -> str:
         return self._pretty_repr(self._describe())
 
     @abc.abstractmethod
-    def _load(self) -> _DO:
+    def load(self) -> _DO:
+        """Loads data by delegation to the provided load method.
+
+        Returns:
+            Data returned by the provided load method.
+
+        Raises:
+            DatasetError: When underlying load method raises error.
+
+        """
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataset and "
-            f"it must implement the '_load' method"
+            f"it must implement the 'load' method"
         )
 
     @abc.abstractmethod
-    def _save(self, data: _DI) -> None:
+    def save(self, data: _DI) -> None:
+        """Saves data by delegation to the provided save method.
+
+        Args:
+            data: the value to be saved by provided save method.
+
+        Raises:
+            DatasetError: when underlying save method raises error.
+            FileNotFoundError: when save method got file instead of dir, on Windows.
+            NotADirectoryError: when save method got file instead of dir, on Unix.
+
+        """
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataset and "
-            f"it must implement the '_save' method"
+            f"it must implement the 'save' method"
         )
 
     @abc.abstractmethod
@@ -509,7 +559,9 @@ def _local_exists(local_filepath: str) -> bool:  # SKIP_IF_NO_SPARK
 class AbstractVersionedDataset(AbstractDataset[_DI, _DO], abc.ABC):
     """
     ``AbstractVersionedDataset`` is the base class for all versioned data set
-    implementations. All data sets that implement versioning should extend this
+    implementations.
+
+    All data sets that implement versioning should extend this
     abstract class and implement the methods marked as abstract.
 
     Example:
@@ -526,11 +578,11 @@ class AbstractVersionedDataset(AbstractDataset[_DI, _DO], abc.ABC):
         >>>         self._param1 = param1
         >>>         self._param2 = param2
         >>>
-        >>>     def _load(self) -> pd.DataFrame:
+        >>>     def load(self) -> pd.DataFrame:
         >>>         load_path = self._get_load_path()
         >>>         return pd.read_csv(load_path)
         >>>
-        >>>     def _save(self, df: pd.DataFrame) -> None:
+        >>>     def save(self, df: pd.DataFrame) -> None:
         >>>         save_path = self._get_save_path()
         >>>         df.to_csv(str(save_path))
         >>>
@@ -652,34 +704,43 @@ def _get_save_path(self) -> PurePosixPath:
     def _get_versioned_path(self, version: str) -> PurePosixPath:
         return self._filepath / version / self._filepath.name
 
-    def load(self) -> _DO:
-        return super().load()
-
-    def save(self, data: _DI) -> None:
-        self._version_cache.clear()
-        save_version = self.resolve_save_version()  # Make sure last save version is set
-        try:
-            super().save(data)
-        except (FileNotFoundError, NotADirectoryError) as err:
-            # FileNotFoundError raised in Win, NotADirectoryError raised in Unix
-            _default_version = "YYYY-MM-DDThh.mm.ss.sssZ"
-            raise DatasetError(
-                f"Cannot save versioned dataset '{self._filepath.name}' to "
-                f"'{self._filepath.parent.as_posix()}' because a file with the same "
-                f"name already exists in the directory. This is likely because "
-                f"versioning was enabled on a dataset already saved previously. Either "
-                f"remove '{self._filepath.name}' from the directory or manually "
-                f"convert it into a versioned dataset by placing it in a versioned "
-                f"directory (e.g. with default versioning format "
-                f"'{self._filepath.as_posix()}/{_default_version}/{self._filepath.name}"
-                f"')."
-            ) from err
+    @classmethod
+    def _save_wrapper(
+        cls, save_func: Callable[[Self, _DI], None]
+    ) -> Callable[[Self, _DI], None]:
+        """Decorate `save_func` with logging and error handling code."""
+
+        @wraps(save_func)
+        def save(self: Self, data: _DI) -> None:
+            self._version_cache.clear()
+            save_version = (
+                self.resolve_save_version()
+            )  # Make sure last save version is set
+            try:
+                super()._save_wrapper(save_func)(self, data)
+            except (FileNotFoundError, NotADirectoryError) as err:
+                # FileNotFoundError raised in Win, NotADirectoryError raised in Unix
+                _default_version = "YYYY-MM-DDThh.mm.ss.sssZ"
+                raise DatasetError(
+                    f"Cannot save versioned dataset '{self._filepath.name}' to "
+                    f"'{self._filepath.parent.as_posix()}' because a file with the same "
+                    f"name already exists in the directory. This is likely because "
+                    f"versioning was enabled on a dataset already saved previously. Either "
+                    f"remove '{self._filepath.name}' from the directory or manually "
+                    f"convert it into a versioned dataset by placing it in a versioned "
+                    f"directory (e.g. with default versioning format "
+                    f"'{self._filepath.as_posix()}/{_default_version}/{self._filepath.name}"
+                    f"')."
+                ) from err
+
+            load_version = self.resolve_load_version()
+            if load_version != save_version:
+                warnings.warn(
+                    _CONSISTENCY_WARNING.format(save_version, load_version, str(self))
+                )
+                self._version_cache.clear()
 
-        load_version = self.resolve_load_version()
-        if load_version != save_version:
-            warnings.warn(
-                _CONSISTENCY_WARNING.format(save_version, load_version, str(self))
-            )
+        return save
 
     def exists(self) -> bool:
         """Checks whether a data set's output already exists by calling
diff --git a/kedro/io/data_catalog.py b/kedro/io/data_catalog.py
index 5b95c12478..6c70cd038a 100644
--- a/kedro/io/data_catalog.py
+++ b/kedro/io/data_catalog.py
@@ -732,7 +732,7 @@ def add_feed_dict(self, feed_dict: dict[str, Any], replace: bool = False) -> Non
             if isinstance(feed_dict[dataset_name], AbstractDataset):
                 dataset = feed_dict[dataset_name]
             else:
-                dataset = MemoryDataset(data=feed_dict[dataset_name])
+                dataset = MemoryDataset(data=feed_dict[dataset_name])  # type: ignore[abstract]
 
             self.add(dataset_name, dataset, replace)
 
diff --git a/kedro/io/lambda_dataset.py b/kedro/io/lambda_dataset.py
index 4ed3a7a248..126bef6632 100644
--- a/kedro/io/lambda_dataset.py
+++ b/kedro/io/lambda_dataset.py
@@ -49,21 +49,21 @@ def _to_str(func: Any) -> str | None:
 
         return descr
 
-    def _save(self, data: Any) -> None:
-        if not self.__save:
+    def load(self) -> Any:
+        if not self.__load:
             raise DatasetError(
-                "Cannot save to data set. No 'save' function "
+                "Cannot load data set. No 'load' function "
                 "provided when LambdaDataset was created."
             )
-        self.__save(data)
+        return self.__load()
 
-    def _load(self) -> Any:
-        if not self.__load:
+    def save(self, data: Any) -> None:
+        if not self.__save:
             raise DatasetError(
-                "Cannot load data set. No 'load' function "
+                "Cannot save to data set. No 'save' function "
                 "provided when LambdaDataset was created."
             )
-        return self.__load()
+        self.__save(data)
 
     def _exists(self) -> bool:
         if not self.__exists:
diff --git a/kedro/io/memory_dataset.py b/kedro/io/memory_dataset.py
index b3b08e1fae..30b6e19fa8 100644
--- a/kedro/io/memory_dataset.py
+++ b/kedro/io/memory_dataset.py
@@ -57,7 +57,7 @@ def __init__(
         self.metadata = metadata
         self._EPHEMERAL = True
         if data is not _EMPTY:
-            self._save(data)
+            self.save.__wrapped__(self, data)  # type: ignore[attr-defined]
 
     def _load(self) -> Any:
         if self._data is _EMPTY:
diff --git a/pyproject.toml b/pyproject.toml
index 383b7aa34e..0aef9e1cc2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
     "rich>=12.0,<14.0",
     "rope>=0.21,<2.0",  # subject to LGPLv3 license
     "toml>=0.10.0",
+    "typing_extensions>=4.0",
     "graphlib_backport>=1.0.0; python_version < '3.9'",
 ]
 keywords = [


Test patch:
diff --git a/tests/io/test_core.py b/tests/io/test_core.py
index c2f07a106c..ebff20dec2 100644
--- a/tests/io/test_core.py
+++ b/tests/io/test_core.py
@@ -309,10 +309,18 @@ def test_parse_dataset_definition(self):
         dataset, _ = parse_dataset_definition(config)
         assert dataset is LambdaDataset
 
-    def test_test_parse_dataset_definition_with_python_class_type(self):
+    def test_parse_dataset_definition_with_python_class_type(self):
         config = {"type": MyDataset}
         parse_dataset_definition(config)
 
+    def test_load_and_save_are_wrapped_once(self):
+        assert not getattr(
+            MyOtherVersionedDataset.load.__wrapped__, "__loadwrapped__", False
+        )
+        assert not getattr(
+            MyOtherVersionedDataset.save.__wrapped__, "__savewrapped__", False
+        )
+
 
 class TestAbstractVersionedDataset:
     def test_version_str_repr(self, load_version, save_version):
@@ -431,3 +439,135 @@ def test_cache_release(self, my_versioned_dataset):
 
         my_versioned_dataset._release()
         assert my_versioned_dataset._version_cache.currsize == 0
+
+
+class MyLegacyDataset(AbstractDataset):
+    def __init__(self, filepath="", save_args=None, fs_args=None, var=None):
+        self._filepath = PurePosixPath(filepath)
+        self.save_args = save_args
+        self.fs_args = fs_args
+        self.var = var
+
+    def _describe(self):
+        return {"filepath": self._filepath, "var": self.var}
+
+    def _exists(self) -> bool:
+        return Path(self._filepath.as_posix()).exists()
+
+    def _load(self):
+        return pd.read_csv(self._filepath)
+
+    def _save(self, data: str) -> None:
+        with open(self._filepath, mode="w") as file:
+            file.write(data)
+
+
+class MyLegacyVersionedDataset(AbstractVersionedDataset[str, str]):
+    def __init__(  # noqa: PLR0913
+        self,
+        filepath: str,
+        version: Version = None,
+    ) -> None:
+        _fs_args: dict[Any, Any] = {}
+        _fs_args.setdefault("auto_mkdir", True)
+        protocol, path = get_protocol_and_path(filepath, version)
+
+        self._protocol = protocol
+        self._fs = fsspec.filesystem(self._protocol, **_fs_args)
+
+        super().__init__(
+            filepath=PurePosixPath(path),
+            version=version,
+            exists_function=self._fs.exists,
+            glob_function=self._fs.glob,
+        )
+
+    def _describe(self) -> dict[str, Any]:
+        return dict(filepath=self._filepath, version=self._version)
+
+    def _load(self) -> str:
+        load_path = get_filepath_str(self._get_load_path(), self._protocol)
+
+        with self._fs.open(load_path, mode="r") as fs_file:
+            return fs_file.read()
+
+    def _save(self, data: str) -> None:
+        save_path = get_filepath_str(self._get_save_path(), self._protocol)
+
+        with self._fs.open(save_path, mode="w") as fs_file:
+            fs_file.write(data)
+
+    def _exists(self) -> bool:
+        try:
+            load_path = get_filepath_str(self._get_load_path(), self._protocol)
+        except DatasetError:
+            return False
+
+        return self._fs.exists(load_path)
+
+
+@pytest.fixture
+def my_legacy_dataset(filepath_versioned, save_args, fs_args):
+    return MyLegacyDataset(
+        filepath=filepath_versioned, save_args=save_args, fs_args=fs_args
+    )
+
+
+@pytest.fixture
+def my_legacy_versioned_dataset(filepath_versioned, load_version, save_version):
+    return MyLegacyVersionedDataset(
+        filepath=filepath_versioned, version=Version(load_version, save_version)
+    )
+
+
+class TestLegacyLoadAndSave:
+    def test_saving_none(self, my_legacy_dataset):
+        """Check the error when attempting to save the dataset without
+        providing the data"""
+        pattern = r"Saving 'None' to a 'Dataset' is not allowed"
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_dataset.save(None)
+
+    def test_saving_invalid_data(self, my_legacy_dataset, dummy_data):
+        pattern = r"Failed while saving data to data set"
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_dataset.save(pd.DataFrame())
+
+    @pytest.mark.parametrize(
+        "load_version", ["2019-01-01T23.59.59.999Z"], indirect=True
+    )
+    @pytest.mark.parametrize(
+        "save_version", ["2019-01-02T00.00.00.000Z"], indirect=True
+    )
+    def test_save_version_warning(
+        self, my_legacy_versioned_dataset, load_version, save_version, dummy_data
+    ):
+        """Check the warning when saving to the path that differs from
+        the subsequent load path."""
+        pattern = (
+            f"Save version '{save_version}' did not match "
+            f"load version '{load_version}' for "
+            r"MyLegacyVersionedDataset\(.+\)"
+        )
+        with pytest.warns(UserWarning, match=pattern):
+            my_legacy_versioned_dataset.save(dummy_data)
+
+    def test_versioning_existing_dataset(
+        self, my_legacy_dataset, my_legacy_versioned_dataset, dummy_data
+    ):
+        """Check the error when attempting to save a versioned dataset on top of an
+        already existing (non-versioned) dataset."""
+        my_legacy_dataset.save(dummy_data)
+        assert my_legacy_dataset.exists()
+        assert my_legacy_dataset._filepath == my_legacy_versioned_dataset._filepath
+        pattern = (
+            f"(?=.*file with the same name already exists in the directory)"
+            f"(?=.*{my_legacy_versioned_dataset._filepath.parent.as_posix()})"
+        )
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_versioned_dataset.save(dummy_data)
+
+        # Remove non-versioned dataset and try again
+        Path(my_legacy_dataset._filepath.as_posix()).unlink()
+        my_legacy_versioned_dataset.save(dummy_data)
+        assert my_legacy_versioned_dataset.exists()