Repository: kedro-org/kedro
Pull request: Generic typings abstract dataset
Signed-off-by: Niels Drost <codingdutchman@gmail.com>

## Description
Resolves #1564 

## Development notes
Only generic types have been added for `AbstractDataSet` and `AbstractVersionedDataSet`, to define load and save types.
In all classes in extras that use before mentioned classes, the correct types have been added. 
Checked using mypy.

## Checklist

- [x] Read the [contributing](https://github.com/kedro-org/kedro/blob/main/CONTRIBUTING.md) guidelines
- [x] Updated the documentation to reflect the code changes
- [x] Added a description of this change in the [`RELEASE.md`](https://github.com/kedro-org/kedro/blob/main/RELEASE.md) file
❌  Dit not add tests, because it has no run time influence.


<a href="https://gitpod.io/#https://github.com/kedro-org/kedro/pull/1612"><img src="https://gitpod.io/button/open-in-gitpod.svg"/></a>



Problem statement:
Feature: make AbstractDataset take data type as generic parameter.
## Description

Abstract dataset is currently incompletely typed, as _load and _save use "Any".

## Context

I would like to write typed code.

## Possible Implementation

Use:
```
class AbstractDataset(Generic[DS]): ...
```
If necessary for deprecation, could have something like:
```
class BaseDataset(Generic[DS]): ...

class AbstractDataset(BaseDataset[Any]): ...
```
and deprecate `AbstractDataset` in favor of `BaseDataset`.

Hints from issue discussion:
Hi @shaunc, thanks very much for the suggestion! This sounds like a very good idea, but I just want to understand it a bit better since I'm not quite sure how you would intend to use it in practice?

Are you just suggesting that we add the `Generic` to `AbstractDataset` inheritance + `_load` + `_save` methods or are there other changes involved here also?

`AbstractDataSet` currently inherits from `abc.ABC`. As far as I can tell there's no reason we shouldn't just do this?
```
class AbstractDataSet(abc.ABC, Generic[DS]):
```

Also, as far as I can see, there's nothing at all breaking about this change and all datasets will continue to function exactly as they currently do. So I don't see any need to do it in two stages and have a new `BaseDataSet`. Or am I missing something here? Can adding `Generic[DS]` to the inheritance actually break anything for anyone?
Everything "runtime" should work exactly the same. Users who type check, however, will have to use `AbstractDataSet[Foo]` or `AbstractDataSet[Any]` in code where they currently just use `AbstractDataSet` ... from mypy perspective (but not that of your test suite, presumably) it would be a "breaking change".

If I were them (that is, such a user), I would be happy, of course.... I suggested that "just in case".

NB -- e.g. your various derived classes will have to be changed as well. I guess in a few cases you'll want to use `Union[ ... ]`s or pass through the generic for types that accept various formats for saving. One developer should be able to make the changes quickly just by tracking down the places that mypy points at....
Cool, thanks for the explanation. I think this is fine to regard as non-breaking since it's only type-checking rather than actual execution. It will only affect a tiny fraction of users anyway, so I will take you as a representative when you say that it would make such users happy 😀  So long as we add something to the release notes about the change I think we're good here.

So please do go ahead and raise a PR for this, since in all honesty we're not going to do it ourselves 😅  Our type hints could do with improvements all round but it's not a priority really.

Two things I don't quite understand here I hope you can explain.

1. "NB -- e.g. your various derived classes will have to be changed as well". I just tried `class AbstractDataSet(abc.ABC, Generic[DS])` and running mypy and it didn't trigger any messages about our datasets that inherit from `AbstractDataSet`. Would you expect it to?
2. The only way we use mypy currently is [this pre-commit hook](https://github.com/kedro-org/kedro/blob/main/.pre-commit-config.yaml#L31-L41). When I run it on all files there's no errors returned and it also runs on CI without showing anything. I'm pretty sure there are lots of problems though, and indeed when I run `mypy` from the command line with the same options it shows them. Do you know why the pre-commit hook isn't flagging them?

Commit message: fixed tests

Signed-off-by: Niels Drost <codingdutchman@gmail.com>

Patch to review:
diff --git a/RELEASE.md b/RELEASE.md
index 1bf296a84f..1fe5f2efa2 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -25,6 +25,7 @@
 * Reduced number of log lines by changing the logging level from `INFO` to `DEBUG` for low priority messages.
 * Kedro's framework-side logging configuration no longer performs file-based logging. Hence superfluous `info.log`/`errors.log` files are no longer created in your project root, and running Kedro on read-only file systems such as Databricks Repos is now possible.
 * The `root` logger is now set to the Python default level of `WARNING` rather than `INFO`. Kedro's logger is still set to emit `INFO` level messages.
+* Generic typing for `AbstractDataSet` and `AbstractVersionedDataSet`.
 
 ## Upcoming deprecations for Kedro 0.19.0
 * `kedro.extras.ColorHandler` will be removed in 0.19.0.
diff --git a/docs/source/extend_kedro/custom_datasets.md b/docs/source/extend_kedro/custom_datasets.md
index 43ac551940..116b26ed44 100644
--- a/docs/source/extend_kedro/custom_datasets.md
+++ b/docs/source/extend_kedro/custom_datasets.md
@@ -30,6 +30,9 @@ At the minimum, a valid Kedro dataset needs to subclass the base [AbstractDataSe
 * `_save`
 * `_describe`
 
+`AbstractDataSet` is generically typed with an input data type for saving data, and an output data type for loading data.
+This typing is optional however, and defaults to `Any` type.
+
 Here is an example skeleton for `ImageDataSet`:
 
 <details>
@@ -43,7 +46,7 @@ import numpy as np
 from kedro.io import AbstractDataSet
 
 
-class ImageDataSet(AbstractDataSet):
+class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
     """``ImageDataSet`` loads / save image data from a given filepath as `numpy` array using Pillow.
 
     Example:
@@ -109,7 +112,7 @@ from kedro.io import AbstractDataSet
 from kedro.io.core import get_filepath_str, get_protocol_and_path
 
 
-class ImageDataSet(AbstractDataSet):
+class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
     def __init__(self, filepath: str):
         """Creates a new instance of ImageDataSet to load / save image data for given filepath.
 
@@ -166,7 +169,7 @@ Similarly, we can implement the `_save` method as follows:
 
 
 ```python
-class ImageDataSet(AbstractDataSet):
+class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
     def _save(self, data: np.ndarray) -> None:
         """Saves image data to the specified filepath."""
         # using get_filepath_str ensures that the protocol and path are appended correctly for different filesystems
@@ -190,7 +193,7 @@ You can open the file to verify that the data was written back correctly.
 The `_describe` method is used for printing purposes. The convention in Kedro is for the method to return a dictionary describing the attributes of the dataset.
 
 ```python
-class ImageDataSet(AbstractDataSet):
+class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
     def _describe(self) -> Dict[str, Any]:
         """Returns a dict that describes the attributes of the dataset."""
         return dict(filepath=self._filepath, protocol=self._protocol)
@@ -215,7 +218,7 @@ from kedro.io import AbstractDataSet
 from kedro.io.core import get_filepath_str, get_protocol_and_path
 
 
-class ImageDataSet(AbstractDataSet):
+class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
     """``ImageDataSet`` loads / save image data from a given filepath as `numpy` array using Pillow.
 
     Example:
@@ -321,7 +324,7 @@ from kedro.io import AbstractVersionedDataSet
 from kedro.io.core import get_filepath_str, get_protocol_and_path, Version
 
 
-class ImageDataSet(AbstractVersionedDataSet):
+class ImageDataSet(AbstractVersionedDataSet[np.ndarray, np.ndarray]):
     """``ImageDataSet`` loads / save image data from a given filepath as `numpy` array using Pillow.
 
     Example:
@@ -394,8 +397,8 @@ The difference between the original `ImageDataSet` and the versioned `ImageDataS
 +from kedro.io.core import get_filepath_str, get_protocol_and_path, Version
 
 
--class ImageDataSet(AbstractDataSet):
-+class ImageDataSet(AbstractVersionedDataSet):
+-class ImageDataSet(AbstractDataSet[np.ndarray, np.ndarray]):
++class ImageDataSet(AbstractVersionedDataSet[np.ndarray, np.ndarray]):
      """``ImageDataSet`` loads / save image data from a given filepath as `numpy` array using Pillow.
 
      Example:
diff --git a/kedro/extras/datasets/api/api_dataset.py b/kedro/extras/datasets/api/api_dataset.py
index 58c635dee1..09d0dbcab3 100644
--- a/kedro/extras/datasets/api/api_dataset.py
+++ b/kedro/extras/datasets/api/api_dataset.py
@@ -9,7 +9,7 @@
 from kedro.io.core import AbstractDataSet, DataSetError
 
 
-class APIDataSet(AbstractDataSet):
+class APIDataSet(AbstractDataSet[Any, requests.Response]):
     """``APIDataSet`` loads the data from HTTP(S) APIs.
     It uses the python requests library: https://requests.readthedocs.io/en/master/
 
diff --git a/kedro/extras/datasets/biosequence/biosequence_dataset.py b/kedro/extras/datasets/biosequence/biosequence_dataset.py
index 80906abca6..669270f4b9 100644
--- a/kedro/extras/datasets/biosequence/biosequence_dataset.py
+++ b/kedro/extras/datasets/biosequence/biosequence_dataset.py
@@ -11,7 +11,7 @@
 from kedro.io.core import AbstractDataSet, get_filepath_str, get_protocol_and_path
 
 
-class BioSequenceDataSet(AbstractDataSet):
+class BioSequenceDataSet(AbstractDataSet[List, List]):
     r"""``BioSequenceDataSet`` loads and saves data to a sequence file.
 
     Example:
diff --git a/kedro/extras/datasets/dask/parquet_dataset.py b/kedro/extras/datasets/dask/parquet_dataset.py
index 24c1e7a585..b4bbe05bdf 100644
--- a/kedro/extras/datasets/dask/parquet_dataset.py
+++ b/kedro/extras/datasets/dask/parquet_dataset.py
@@ -10,7 +10,7 @@
 from kedro.io.core import AbstractDataSet, get_protocol_and_path
 
 
-class ParquetDataSet(AbstractDataSet):
+class ParquetDataSet(AbstractDataSet[dd.DataFrame, dd.DataFrame]):
     """``ParquetDataSet`` loads and saves data to parquet file(s). It uses Dask
     remote data services to handle the corresponding load and save operations:
     https://docs.dask.org/en/latest/how-to/connect-to-remote-data.html
diff --git a/kedro/extras/datasets/email/message_dataset.py b/kedro/extras/datasets/email/message_dataset.py
index 7b0ed58407..08dc157fc6 100644
--- a/kedro/extras/datasets/email/message_dataset.py
+++ b/kedro/extras/datasets/email/message_dataset.py
@@ -22,7 +22,7 @@
 
 
 class EmailMessageDataSet(
-    AbstractVersionedDataSet
+    AbstractVersionedDataSet[Message, Message]
 ):  # pylint: disable=too-many-instance-attributes
     """``EmailMessageDataSet`` loads/saves an email message from/to a file
     using an underlying filesystem (e.g.: local, S3, GCS). It uses the
diff --git a/kedro/extras/datasets/geopandas/geojson_dataset.py b/kedro/extras/datasets/geopandas/geojson_dataset.py
index 0ef3812f04..6e82eae7de 100644
--- a/kedro/extras/datasets/geopandas/geojson_dataset.py
+++ b/kedro/extras/datasets/geopandas/geojson_dataset.py
@@ -18,7 +18,11 @@
 )
 
 
-class GeoJSONDataSet(AbstractVersionedDataSet):
+class GeoJSONDataSet(
+    AbstractVersionedDataSet[
+        gpd.GeoDataFrame, Union[gpd.GeoDataFrame, Dict[str, gpd.GeoDataFrame]]
+    ]
+):
     """``GeoJSONDataSet`` loads/saves data to a GeoJSON file using an underlying filesystem
     (eg: local, S3, GCS).
     The underlying functionality is supported by geopandas, so it supports all
diff --git a/kedro/extras/datasets/holoviews/holoviews_writer.py b/kedro/extras/datasets/holoviews/holoviews_writer.py
index 60ad02d840..aac1024d8d 100644
--- a/kedro/extras/datasets/holoviews/holoviews_writer.py
+++ b/kedro/extras/datasets/holoviews/holoviews_writer.py
@@ -21,7 +21,7 @@
 HoloViews = TypeVar("HoloViews")
 
 
-class HoloviewsWriter(AbstractVersionedDataSet):
+class HoloviewsWriter(AbstractVersionedDataSet[HoloViews, str]):
     """``HoloviewsWriter`` saves Holoviews objects to image file(s) in an underlying
     filesystem (e.g. local, S3, GCS).
 
diff --git a/kedro/extras/datasets/json/json_dataset.py b/kedro/extras/datasets/json/json_dataset.py
index 067e4071f9..0cca0f1794 100644
--- a/kedro/extras/datasets/json/json_dataset.py
+++ b/kedro/extras/datasets/json/json_dataset.py
@@ -17,7 +17,7 @@
 )
 
 
-class JSONDataSet(AbstractVersionedDataSet):
+class JSONDataSet(AbstractVersionedDataSet[Dict, Dict]):
     """``JSONDataSet`` loads/saves data from/to a JSON file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses native json to handle the JSON file.
 
diff --git a/kedro/extras/datasets/matplotlib/matplotlib_writer.py b/kedro/extras/datasets/matplotlib/matplotlib_writer.py
index 9f9edef489..0043314d20 100644
--- a/kedro/extras/datasets/matplotlib/matplotlib_writer.py
+++ b/kedro/extras/datasets/matplotlib/matplotlib_writer.py
@@ -19,7 +19,11 @@
 )
 
 
-class MatplotlibWriter(AbstractVersionedDataSet):
+class MatplotlibWriter(
+    AbstractVersionedDataSet[
+        Union[plt.figure, List[plt.figure], Dict[str, plt.figure]], None
+    ]
+):
     """``MatplotlibWriter`` saves one or more Matplotlib objects as
     image files to an underlying filesystem (e.g. local, S3, GCS).
 
diff --git a/kedro/extras/datasets/networkx/gml_dataset.py b/kedro/extras/datasets/networkx/gml_dataset.py
index 350878aeb3..5d615ed609 100644
--- a/kedro/extras/datasets/networkx/gml_dataset.py
+++ b/kedro/extras/datasets/networkx/gml_dataset.py
@@ -18,7 +18,7 @@
 )
 
 
-class GMLDataSet(AbstractVersionedDataSet):
+class GMLDataSet(AbstractVersionedDataSet[networkx.Graph, networkx.Graph]):
     """``GMLDataSet`` loads and saves graphs to a GML file using an
     underlying filesystem (e.g.: local, S3, GCS). ``NetworkX`` is used to
     create GML data.
diff --git a/kedro/extras/datasets/networkx/graphml_dataset.py b/kedro/extras/datasets/networkx/graphml_dataset.py
index e913e9b52b..e599902117 100644
--- a/kedro/extras/datasets/networkx/graphml_dataset.py
+++ b/kedro/extras/datasets/networkx/graphml_dataset.py
@@ -17,7 +17,7 @@
 )
 
 
-class GraphMLDataSet(AbstractVersionedDataSet):
+class GraphMLDataSet(AbstractVersionedDataSet[networkx.Graph, networkx.Graph]):
     """``GraphMLDataSet`` loads and saves graphs to a GraphML file using an
     underlying filesystem (e.g.: local, S3, GCS). ``NetworkX`` is used to
     create GraphML data.
diff --git a/kedro/extras/datasets/networkx/json_dataset.py b/kedro/extras/datasets/networkx/json_dataset.py
index fae22b4128..519c2a424c 100644
--- a/kedro/extras/datasets/networkx/json_dataset.py
+++ b/kedro/extras/datasets/networkx/json_dataset.py
@@ -18,7 +18,7 @@
 )
 
 
-class JSONDataSet(AbstractVersionedDataSet):
+class JSONDataSet(AbstractVersionedDataSet[networkx.Graph, networkx.Graph]):
     """NetworkX ``JSONDataSet`` loads and saves graphs to a JSON file using an
     underlying filesystem (e.g.: local, S3, GCS). ``NetworkX`` is used to
     create JSON data.
diff --git a/kedro/extras/datasets/pandas/csv_dataset.py b/kedro/extras/datasets/pandas/csv_dataset.py
index 0c63f87bc3..336e10e4b8 100644
--- a/kedro/extras/datasets/pandas/csv_dataset.py
+++ b/kedro/extras/datasets/pandas/csv_dataset.py
@@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class CSVDataSet(AbstractVersionedDataSet):
+class CSVDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
     """``CSVDataSet`` loads/saves data from/to a CSV file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to handle the CSV file.
 
diff --git a/kedro/extras/datasets/pandas/excel_dataset.py b/kedro/extras/datasets/pandas/excel_dataset.py
index 0dc8dad5ee..9ace17ada0 100644
--- a/kedro/extras/datasets/pandas/excel_dataset.py
+++ b/kedro/extras/datasets/pandas/excel_dataset.py
@@ -22,7 +22,12 @@
 logger = logging.getLogger(__name__)
 
 
-class ExcelDataSet(AbstractVersionedDataSet):
+class ExcelDataSet(
+    AbstractVersionedDataSet[
+        Union[pd.DataFrame, Dict[str, pd.DataFrame]],
+        Union[pd.DataFrame, Dict[str, pd.DataFrame]],
+    ]
+):
     """``ExcelDataSet`` loads/saves data from/to a Excel file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to handle the Excel file.
 
diff --git a/kedro/extras/datasets/pandas/feather_dataset.py b/kedro/extras/datasets/pandas/feather_dataset.py
index 52ea920fd6..9adba3c9b3 100644
--- a/kedro/extras/datasets/pandas/feather_dataset.py
+++ b/kedro/extras/datasets/pandas/feather_dataset.py
@@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class FeatherDataSet(AbstractVersionedDataSet):
+class FeatherDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
     """``FeatherDataSet`` loads and saves data to a feather file using an
     underlying filesystem (e.g.: local, S3, GCS). The underlying functionality
     is supported by pandas, so it supports all allowed pandas options
diff --git a/kedro/extras/datasets/pandas/gbq_dataset.py b/kedro/extras/datasets/pandas/gbq_dataset.py
index 9bd04e9141..5a5dc99406 100644
--- a/kedro/extras/datasets/pandas/gbq_dataset.py
+++ b/kedro/extras/datasets/pandas/gbq_dataset.py
@@ -21,7 +21,7 @@
 )
 
 
-class GBQTableDataSet(AbstractDataSet):
+class GBQTableDataSet(AbstractDataSet[pd.DataFrame, pd.DataFrame]):
     """``GBQTableDataSet`` loads and saves data from/to Google BigQuery.
     It uses pandas-gbq to read and write from/to BigQuery table.
 
diff --git a/kedro/extras/datasets/pandas/generic_dataset.py b/kedro/extras/datasets/pandas/generic_dataset.py
index 0d74e0f45a..60bd4c9193 100644
--- a/kedro/extras/datasets/pandas/generic_dataset.py
+++ b/kedro/extras/datasets/pandas/generic_dataset.py
@@ -29,7 +29,7 @@
 ]
 
 
-class GenericDataSet(AbstractVersionedDataSet):
+class GenericDataSet(AbstractVersionedDataSet[pd.DataFrame, Any]):
     """`pandas.GenericDataSet` loads/saves data from/to a data file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to dynamically select the
     appropriate type of read/write target on a best effort basis.
diff --git a/kedro/extras/datasets/pandas/hdf_dataset.py b/kedro/extras/datasets/pandas/hdf_dataset.py
index 6e12350e27..5eeb38ff93 100644
--- a/kedro/extras/datasets/pandas/hdf_dataset.py
+++ b/kedro/extras/datasets/pandas/hdf_dataset.py
@@ -20,7 +20,7 @@
 HDFSTORE_DRIVER = "H5FD_CORE"
 
 
-class HDFDataSet(AbstractVersionedDataSet):
+class HDFDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
     """``HDFDataSet`` loads/saves data from/to a hdf file using an underlying
     filesystem (e.g. local, S3, GCS). It uses pandas.HDFStore to handle the hdf file.
 
diff --git a/kedro/extras/datasets/pandas/json_dataset.py b/kedro/extras/datasets/pandas/json_dataset.py
index 38b8199375..72d77e972d 100644
--- a/kedro/extras/datasets/pandas/json_dataset.py
+++ b/kedro/extras/datasets/pandas/json_dataset.py
@@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class JSONDataSet(AbstractVersionedDataSet):
+class JSONDataSet(AbstractVersionedDataSet[pd.DataFrame, Any]):
     """``JSONDataSet`` loads/saves data from/to a JSON file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to handle the json file.
 
diff --git a/kedro/extras/datasets/pandas/parquet_dataset.py b/kedro/extras/datasets/pandas/parquet_dataset.py
index 9a01c17b34..664c470fa7 100644
--- a/kedro/extras/datasets/pandas/parquet_dataset.py
+++ b/kedro/extras/datasets/pandas/parquet_dataset.py
@@ -23,7 +23,7 @@
 logger = logging.getLogger(__name__)
 
 
-class ParquetDataSet(AbstractVersionedDataSet):
+class ParquetDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
     """``ParquetDataSet`` loads/saves data from/to a Parquet file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to handle the Parquet file.
 
diff --git a/kedro/extras/datasets/pandas/sql_dataset.py b/kedro/extras/datasets/pandas/sql_dataset.py
index 88b60c90a4..12ce0f964d 100644
--- a/kedro/extras/datasets/pandas/sql_dataset.py
+++ b/kedro/extras/datasets/pandas/sql_dataset.py
@@ -87,7 +87,7 @@ def _get_sql_alchemy_missing_error() -> DataSetError:
     )
 
 
-class SQLTableDataSet(AbstractDataSet):
+class SQLTableDataSet(AbstractDataSet[pd.DataFrame, pd.DataFrame]):
     """``SQLTableDataSet`` loads data from a SQL table and saves a pandas
     dataframe to a table. It uses ``pandas.DataFrame`` internally,
     so it supports all allowed pandas options on ``read_sql_table`` and
@@ -257,7 +257,7 @@ def _exists(self) -> bool:
         return exists
 
 
-class SQLQueryDataSet(AbstractDataSet):
+class SQLQueryDataSet(AbstractDataSet[pd.DataFrame, pd.DataFrame]):
     """``SQLQueryDataSet`` loads data from a provided SQL query. It
     uses ``pandas.DataFrame`` internally, so it supports all allowed
     pandas options on ``read_sql_query``. Since Pandas uses SQLAlchemy behind
diff --git a/kedro/extras/datasets/pandas/xml_dataset.py b/kedro/extras/datasets/pandas/xml_dataset.py
index ebfc42797d..3e5d59989b 100644
--- a/kedro/extras/datasets/pandas/xml_dataset.py
+++ b/kedro/extras/datasets/pandas/xml_dataset.py
@@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class XMLDataSet(AbstractVersionedDataSet):
+class XMLDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
     """``XMLDataSet`` loads/saves data from/to a XML file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses pandas to handle the XML file.
 
diff --git a/kedro/extras/datasets/pickle/pickle_dataset.py b/kedro/extras/datasets/pickle/pickle_dataset.py
index bf8027750b..bdddf815bf 100644
--- a/kedro/extras/datasets/pickle/pickle_dataset.py
+++ b/kedro/extras/datasets/pickle/pickle_dataset.py
@@ -19,7 +19,7 @@
 )
 
 
-class PickleDataSet(AbstractVersionedDataSet):
+class PickleDataSet(AbstractVersionedDataSet[Any, Any]):
     """``PickleDataSet`` loads/saves data from/to a Pickle file using an underlying
     filesystem (e.g.: local, S3, GCS). The underlying functionality is supported by
     the specified backend library passed in (defaults to the ``pickle`` library), so it
diff --git a/kedro/extras/datasets/pillow/image_dataset.py b/kedro/extras/datasets/pillow/image_dataset.py
index c85776282d..c498361815 100644
--- a/kedro/extras/datasets/pillow/image_dataset.py
+++ b/kedro/extras/datasets/pillow/image_dataset.py
@@ -17,7 +17,7 @@
 )
 
 
-class ImageDataSet(AbstractVersionedDataSet):
+class ImageDataSet(AbstractVersionedDataSet[Image.Image, Image.Image]):
     """``ImageDataSet`` loads/saves image data as `numpy` from an underlying
     filesystem (e.g.: local, S3, GCS). It uses Pillow to handle image file.
 
@@ -107,13 +107,13 @@ def _describe(self) -> Dict[str, Any]:
             version=self._version,
         )
 
-    def _load(self) -> Image:
+    def _load(self) -> Image.Image:
         load_path = get_filepath_str(self._get_load_path(), self._protocol)
 
         with self._fs.open(load_path, **self._fs_open_args_load) as fs_file:
             return Image.open(fs_file).copy()
 
-    def _save(self, data: Image) -> None:
+    def _save(self, data: Image.Image) -> None:
         save_path = get_filepath_str(self._get_save_path(), self._protocol)
 
         with self._fs.open(save_path, **self._fs_open_args_save) as fs_file:
diff --git a/kedro/extras/datasets/plotly/json_dataset.py b/kedro/extras/datasets/plotly/json_dataset.py
index f1ec1ea73c..77da180394 100644
--- a/kedro/extras/datasets/plotly/json_dataset.py
+++ b/kedro/extras/datasets/plotly/json_dataset.py
@@ -17,7 +17,9 @@
 )
 
 
-class JSONDataSet(AbstractVersionedDataSet):
+class JSONDataSet(
+    AbstractVersionedDataSet[go.Figure, Union[go.Figure, go.FigureWidget]]
+):
     """``JSONDataSet`` loads/saves a plotly figure from/to a JSON file using an
     underlying filesystem (e.g.: local, S3, GCS).
 
diff --git a/kedro/extras/datasets/redis/redis_dataset.py b/kedro/extras/datasets/redis/redis_dataset.py
index 1e8e2d49e7..b00226ea18 100644
--- a/kedro/extras/datasets/redis/redis_dataset.py
+++ b/kedro/extras/datasets/redis/redis_dataset.py
@@ -12,7 +12,7 @@
 from kedro.io.core import AbstractDataSet, DataSetError
 
 
-class PickleDataSet(AbstractDataSet):
+class PickleDataSet(AbstractDataSet[Any, Any]):
     """``PickleDataSet`` loads/saves data from/to a Redis database. The
     underlying functionality is supported by the redis library, so it supports
     all allowed options for instantiating the redis app ``from_url`` and setting
diff --git a/kedro/extras/datasets/spark/deltatable_dataset.py b/kedro/extras/datasets/spark/deltatable_dataset.py
index da4ddae39b..ec00d5fc45 100644
--- a/kedro/extras/datasets/spark/deltatable_dataset.py
+++ b/kedro/extras/datasets/spark/deltatable_dataset.py
@@ -15,7 +15,7 @@
 from kedro.io.core import AbstractDataSet, DataSetError
 
 
-class DeltaTableDataSet(AbstractDataSet):
+class DeltaTableDataSet(AbstractDataSet[Any, DeltaTable]):
     """``DeltaTableDataSet`` loads data into DeltaTable objects.
 
         Example adding a catalog entry with
diff --git a/kedro/extras/datasets/spark/spark_dataset.py b/kedro/extras/datasets/spark/spark_dataset.py
index 85b4e52589..250504433c 100644
--- a/kedro/extras/datasets/spark/spark_dataset.py
+++ b/kedro/extras/datasets/spark/spark_dataset.py
@@ -157,7 +157,7 @@ def hdfs_glob(self, pattern: str) -> List[str]:
         return sorted(matched)
 
 
-class SparkDataSet(AbstractVersionedDataSet):
+class SparkDataSet(AbstractVersionedDataSet[DataFrame, DataFrame]):
     """``SparkDataSet`` loads and saves Spark dataframes.
 
     Example adding a catalog entry with
diff --git a/kedro/extras/datasets/spark/spark_hive_dataset.py b/kedro/extras/datasets/spark/spark_hive_dataset.py
index 17b1020799..fd743e64c5 100644
--- a/kedro/extras/datasets/spark/spark_hive_dataset.py
+++ b/kedro/extras/datasets/spark/spark_hive_dataset.py
@@ -12,7 +12,7 @@
 
 
 # pylint:disable=too-many-instance-attributes
-class SparkHiveDataSet(AbstractDataSet):
+class SparkHiveDataSet(AbstractDataSet[DataFrame, DataFrame]):
     """``SparkHiveDataSet`` loads and saves Spark dataframes stored on Hive.
     This data set also handles some incompatible file types such as using partitioned parquet on
     hive which will not normally allow upserts to existing data without a complete replacement
diff --git a/kedro/extras/datasets/spark/spark_jdbc_dataset.py b/kedro/extras/datasets/spark/spark_jdbc_dataset.py
index df9f236bd2..e4c6ccdb63 100644
--- a/kedro/extras/datasets/spark/spark_jdbc_dataset.py
+++ b/kedro/extras/datasets/spark/spark_jdbc_dataset.py
@@ -10,7 +10,7 @@
 __all__ = ["SparkJDBCDataSet"]
 
 
-class SparkJDBCDataSet(AbstractDataSet):
+class SparkJDBCDataSet(AbstractDataSet[DataFrame, DataFrame]):
     """``SparkJDBCDataSet`` loads data from a database table accessible
     via JDBC URL url and connection properties and saves the content of
     a PySpark DataFrame to an external database table via JDBC.  It uses
diff --git a/kedro/extras/datasets/tensorflow/tensorflow_model_dataset.py b/kedro/extras/datasets/tensorflow/tensorflow_model_dataset.py
index 354aab5823..441ffd455f 100644
--- a/kedro/extras/datasets/tensorflow/tensorflow_model_dataset.py
+++ b/kedro/extras/datasets/tensorflow/tensorflow_model_dataset.py
@@ -20,7 +20,7 @@
 TEMPORARY_H5_FILE = "tmp_tensorflow_model.h5"
 
 
-class TensorFlowModelDataset(AbstractVersionedDataSet):
+class TensorFlowModelDataset(AbstractVersionedDataSet[tf.keras.Model, tf.keras.Model]):
     """``TensorflowModelDataset`` loads and saves TensorFlow models.
     The underlying functionality is supported by, and passes input arguments through to,
     TensorFlow 2.X load_model and save_model methods.
diff --git a/kedro/extras/datasets/text/text_dataset.py b/kedro/extras/datasets/text/text_dataset.py
index 274742f7bd..a2e01e6a93 100644
--- a/kedro/extras/datasets/text/text_dataset.py
+++ b/kedro/extras/datasets/text/text_dataset.py
@@ -16,7 +16,7 @@
 )
 
 
-class TextDataSet(AbstractVersionedDataSet):
+class TextDataSet(AbstractVersionedDataSet[str, str]):
     """``TextDataSet`` loads/saves data from/to a text file using an underlying
     filesystem (e.g.: local, S3, GCS)
 
diff --git a/kedro/extras/datasets/yaml/yaml_dataset.py b/kedro/extras/datasets/yaml/yaml_dataset.py
index b8d043eb5d..affa38afcf 100644
--- a/kedro/extras/datasets/yaml/yaml_dataset.py
+++ b/kedro/extras/datasets/yaml/yaml_dataset.py
@@ -17,7 +17,7 @@
 )
 
 
-class YAMLDataSet(AbstractVersionedDataSet):
+class YAMLDataSet(AbstractVersionedDataSet[Dict, Dict]):
     """``YAMLDataSet`` loads/saves data from/to a YAML file using an underlying
     filesystem (e.g.: local, S3, GCS). It uses PyYAML to handle the YAML file.
 
diff --git a/kedro/io/core.py b/kedro/io/core.py
index c0f13bf9a1..f475ef02d3 100644
--- a/kedro/io/core.py
+++ b/kedro/io/core.py
@@ -13,7 +13,7 @@
 from glob import iglob
 from operator import attrgetter
 from pathlib import Path, PurePath, PurePosixPath
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+from typing import Any, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
 from urllib.parse import urlsplit
 
 from cachetools import Cache, cachedmethod
@@ -66,7 +66,11 @@ class VersionNotFoundError(DataSetError):
     pass
 
 
-class AbstractDataSet(abc.ABC):
+_DI = TypeVar("_DI")
+_DO = TypeVar("_DO")
+
+
+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.
@@ -81,7 +85,7 @@ class AbstractDataSet(abc.ABC):
         >>> from kedro.io import AbstractDataSet
         >>>
         >>>
-        >>> class MyOwnDataSet(AbstractDataSet):
+        >>> class MyOwnDataSet(AbstractDataSet[pd.DataFrame, pd.DataFrame]):
         >>>     def __init__(self, filepath, param1, param2=True):
         >>>         self._filepath = PurePosixPath(filepath)
         >>>         self._param1 = param1
@@ -165,7 +169,7 @@ def from_config(
     def _logger(self) -> logging.Logger:
         return logging.getLogger(__name__)
 
-    def load(self) -> Any:
+    def load(self) -> _DO:
         """Loads data by delegation to the provided load method.
 
         Returns:
@@ -190,7 +194,7 @@ def load(self) -> Any:
             )
             raise DataSetError(message) from exc
 
-    def save(self, data: Any) -> None:
+    def save(self, data: _DI) -> None:
         """Saves data by delegation to the provided save method.
 
         Args:
@@ -244,14 +248,14 @@ def _to_str(obj, is_root=False):
         return f"{type(self).__name__}({_to_str(self._describe(), True)})"
 
     @abc.abstractmethod
-    def _load(self) -> Any:
+    def _load(self) -> _DO:
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataSet and "
             f"it must implement the '_load' method"
         )
 
     @abc.abstractmethod
-    def _save(self, data: Any) -> None:
+    def _save(self, data: _DI) -> None:
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataSet and "
             f"it must implement the '_save' method"
@@ -450,7 +454,7 @@ def _local_exists(filepath: str) -> bool:  # SKIP_IF_NO_SPARK
     return filepath.exists() or any(par.is_file() for par in filepath.parents)
 
 
-class AbstractVersionedDataSet(AbstractDataSet, abc.ABC):
+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
@@ -590,11 +594,11 @@ def _get_save_path(self) -> PurePosixPath:
     def _get_versioned_path(self, version: str) -> PurePosixPath:
         return self._filepath / version / self._filepath.name
 
-    def load(self) -> Any:
+    def load(self) -> _DO:
         self.resolve_load_version()  # Make sure last load version is set
         return super().load()
 
-    def save(self, data: Any) -> None:
+    def save(self, data: _DI) -> None:
         self._version_cache.clear()
         save_version = self.resolve_save_version()  # Make sure last save version is set
         try:


Code review comments:
diff_hunk: @@ -21,7 +21,7 @@
 HoloViews = TypeVar("HoloViews")
 
 
-class HoloviewsWriter(AbstractVersionedDataSet):
+class HoloviewsWriter(AbstractVersionedDataSet[HoloViews, str]):
path: kedro/extras/datasets/holoviews/holoviews_writer.py
text: @user1:
```suggestion
class HoloviewsWriter(AbstractVersionedDataSet[HoloViews, None]):
```

And change the type signature for `_load` 🙂

diff_hunk: @@ -22,7 +22,7 @@
 logger = logging.getLogger(__name__)
 
 
-class JSONDataSet(AbstractVersionedDataSet):
+class JSONDataSet(AbstractVersionedDataSet[pd.DataFrame, Any]):
path: kedro/extras/datasets/pandas/json_dataset.py
text: @user1:
```suggestion
class JSONDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
```
And change the type signature for `_load` 🙂

diff_hunk: @@ -29,7 +29,7 @@
 ]
 
 
-class GenericDataSet(AbstractVersionedDataSet):
+class GenericDataSet(AbstractVersionedDataSet[pd.DataFrame, Any]):
path: kedro/extras/datasets/pandas/generic_dataset.py
text: @user1:
```suggestion
class GenericDataSet(AbstractVersionedDataSet[pd.DataFrame, pd.DataFrame]):
```
And change the type signature for `_load` 🙂