Make Viz compatible with kedro-datasets (#1214)

* Make kedro viz compatible with kedro-datasets

Signed-off-by: Merel Theisen <merel.theisen@quantumblack.com>

* Fix linting

Signed-off-by: Merel Theisen <merel.theisen@quantumblack.com>

* Fix lint

Signed-off-by: Merel Theisen <merel.theisen@quantumblack.com>

* Make implementation consistent across flowchart and ET

Signed-off-by: Merel Theisen <merel.theisen@quantumblack.com>

* fe changes

* fix tests to work with kedro-datasets

* fix tests

* fix lint

* update release.md

* Update package/kedro_viz/models/experiment_tracking.py

Co-authored-by: Antony Milne <49395058+AntonyMilneQB@users.noreply.github.com>

* replaced dataset_module_class with dataset_type and remove the original dataset_type

* fix tests

* fix lint

* small fix on fe

* fix tests

* fix lint

Signed-off-by: Merel Theisen <merel.theisen@quantumblack.com>
Co-authored-by: Rashida Kanchwala <rashida_kanchwala@mckinsey.com>
Co-authored-by: rashidakanchwala <37628668+rashidakanchwala@users.noreply.github.com>
Co-authored-by: Antony Milne <49395058+AntonyMilneQB@users.noreply.github.com>

Makefile

@@ -6,7 +6,7 @@ package:
 	cd package && python setup.py clean --all
 	cd package && python setup.py sdist bdist_wheel
 
-build: 
+build:
 	rm -rf build package/build package/dist package/kedro_viz/html pip-wheel-metadata package/kedro_viz.egg-info
 	npm run build
 	cp -R build package/kedro_viz/html
@@ -16,10 +16,10 @@ PROJECT_PATH ?= demo-project
 run:
 	PYTHONPATH=$(shell pwd)/package python3 package/kedro_viz/server.py $(PROJECT_PATH)
 
-pytest: 
+pytest:
 	cd package && pytest --cov-fail-under=100
 
-e2e-tests: 
+e2e-tests:
 	cd package && behave
 
 lint: format-fix lint-check
@@ -34,7 +34,7 @@ format-check:
 
 lint-check:
 	pylint --rcfile=package/.pylintrc -j 0 package/kedro_viz
-	pylint --rcfile=package/.pylintrc -j 0 --disable=protected-access,missing-docstring,redefined-outer-name,no-self-use,invalid-name,too-few-public-methods,no-member,unused-argument,duplicate-code package/tests
+	pylint --rcfile=package/.pylintrc -j 0 --disable=protected-access,missing-docstring,redefined-outer-name,invalid-name,too-few-public-methods,no-member,unused-argument,duplicate-code package/tests
 	pylint --rcfile=package/.pylintrc -j 0 --disable=missing-docstring,no-name-in-module,unused-argument package/features
 	flake8 --config=package/.flake8 package
 	mypy --config-file=package/mypy.ini package

RELEASE.md

@@ -13,6 +13,7 @@ Please follow the established format:
 - Allow users to hide modular pipelines on the flowchart. (#1046)
 - Create URL parameters for each element/section in the flowchart. (#1138)
 - Improve CLI loading time. (#1196)
+- Make Kedro-Viz compatible with kedro-datasets. (#1214)
 
 ## Bug fixes and other changes
 

demo-project/data/08_reporting/feature_importance_plot.json/2023-01-16T15.55.27.850Z/feature_importance_plot.json

@@ -0,0 +1 @@
+{"data":[{"alignmentgroup":"True","hovertemplate":"Score=%{marker.color}<br>Feature=%{y}<extra></extra>","legendgroup":"","marker":{"color":[0.07220573068512559,0.09630163347191056,0.10757956840752514,0.16754043714468225,0.2331547228058376,0.3277623812767221,0.3342728074803427,0.36552145755102194,0.3759890244807811,0.3884506496930368,0.40808494303300524,0.6149014984739918,0.8792522745524257,0.9056186845540457,0.9251081185297177],"coloraxis":"coloraxis","pattern":{"shape":""}},"name":"","offsetgroup":"","orientation":"h","showlegend":false,"textposition":"auto","x":[0.07220573068512559,0.09630163347191056,0.10757956840752514,0.16754043714468225,0.2331547228058376,0.3277623812767221,0.3342728074803427,0.36552145755102194,0.3759890244807811,0.3884506496930368,0.40808494303300524,0.6149014984739918,0.8792522745524257,0.9056186845540457,0.9251081185297177],"xaxis":"x","y":["feature_5","feature_1","feature_2","feature_11","feature_4","feature_8","feature_3","feature_10","feature_7","feature_0","feature_9","feature_12","feature_6","feature_14","feature_13"],"yaxis":"y","type":"bar"}],"layout":{"template":{"data":{"histogram2dcontour":[{"type":"histogram2dcontour","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"choropleth":[{"type":"choropleth","colorbar":{"outlinewidth":0,"ticks":""}}],"histogram2d":[{"type":"histogram2d","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"heatmap":[{"type":"heatmap","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"heatmapgl":[{"type":"heatmapgl","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"contourcarpet":[{"type":"contourcarpet","colorbar":{"outlinewidth":0,"ticks":""}}],"contour":[{"type":"contour","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"surface":[{"type":"surface","colorbar":{"outlinewidth":0,"ticks":""},"colorscale":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]]}],"mesh3d":[{"type":"mesh3d","colorbar":{"outlinewidth":0,"ticks":""}}],"scatter":[{"fillpattern":{"fillmode":"overlay","size":10,"solidity":0.2},"type":"scatter"}],"parcoords":[{"type":"parcoords","line":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterpolargl":[{"type":"scatterpolargl","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"bar":[{"error_x":{"color":"#2a3f5f"},"error_y":{"color":"#2a3f5f"},"marker":{"line":{"color":"#E5ECF6","width":0.5},"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"bar"}],"scattergeo":[{"type":"scattergeo","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterpolar":[{"type":"scatterpolar","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"histogram":[{"marker":{"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"histogram"}],"scattergl":[{"type":"scattergl","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatter3d":[{"type":"scatter3d","line":{"colorbar":{"outlinewidth":0,"ticks":""}},"marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scattermapbox":[{"type":"scattermapbox","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scatterternary":[{"type":"scatterternary","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"scattercarpet":[{"type":"scattercarpet","marker":{"colorbar":{"outlinewidth":0,"ticks":""}}}],"carpet":[{"aaxis":{"endlinecolor":"#2a3f5f","gridcolor":"white","linecolor":"white","minorgridcolor":"white","startlinecolor":"#2a3f5f"},"baxis":{"endlinecolor":"#2a3f5f","gridcolor":"white","linecolor":"white","minorgridcolor":"white","startlinecolor":"#2a3f5f"},"type":"carpet"}],"table":[{"cells":{"fill":{"color":"#EBF0F8"},"line":{"color":"white"}},"header":{"fill":{"color":"#C8D4E3"},"line":{"color":"white"}},"type":"table"}],"barpolar":[{"marker":{"line":{"color":"#E5ECF6","width":0.5},"pattern":{"fillmode":"overlay","size":10,"solidity":0.2}},"type":"barpolar"}],"pie":[{"automargin":true,"type":"pie"}]},"layout":{"autotypenumbers":"strict","colorway":["#636efa","#EF553B","#00cc96","#ab63fa","#FFA15A","#19d3f3","#FF6692","#B6E880","#FF97FF","#FECB52"],"font":{"color":"#2a3f5f"},"hovermode":"closest","hoverlabel":{"align":"left"},"paper_bgcolor":"white","plot_bgcolor":"#E5ECF6","polar":{"bgcolor":"#E5ECF6","angularaxis":{"gridcolor":"white","linecolor":"white","ticks":""},"radialaxis":{"gridcolor":"white","linecolor":"white","ticks":""}},"ternary":{"bgcolor":"#E5ECF6","aaxis":{"gridcolor":"white","linecolor":"white","ticks":""},"baxis":{"gridcolor":"white","linecolor":"white","ticks":""},"caxis":{"gridcolor":"white","linecolor":"white","ticks":""}},"coloraxis":{"colorbar":{"outlinewidth":0,"ticks":""}},"colorscale":{"sequential":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]],"sequentialminus":[[0.0,"#0d0887"],[0.1111111111111111,"#46039f"],[0.2222222222222222,"#7201a8"],[0.3333333333333333,"#9c179e"],[0.4444444444444444,"#bd3786"],[0.5555555555555556,"#d8576b"],[0.6666666666666666,"#ed7953"],[0.7777777777777778,"#fb9f3a"],[0.8888888888888888,"#fdca26"],[1.0,"#f0f921"]],"diverging":[[0,"#8e0152"],[0.1,"#c51b7d"],[0.2,"#de77ae"],[0.3,"#f1b6da"],[0.4,"#fde0ef"],[0.5,"#f7f7f7"],[0.6,"#e6f5d0"],[0.7,"#b8e186"],[0.8,"#7fbc41"],[0.9,"#4d9221"],[1,"#276419"]]},"xaxis":{"gridcolor":"white","linecolor":"white","ticks":"","title":{"standoff":15},"zerolinecolor":"white","automargin":true,"zerolinewidth":2},"yaxis":{"gridcolor":"white","linecolor":"white","ticks":"","title":{"standoff":15},"zerolinecolor":"white","automargin":true,"zerolinewidth":2},"scene":{"xaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2},"yaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2},"zaxis":{"backgroundcolor":"#E5ECF6","gridcolor":"white","linecolor":"white","showbackground":true,"ticks":"","zerolinecolor":"white","gridwidth":2}},"shapedefaults":{"line":{"color":"#2a3f5f"}},"annotationdefaults":{"arrowcolor":"#2a3f5f","arrowhead":0,"arrowwidth":1},"geo":{"bgcolor":"white","landcolor":"#E5ECF6","subunitcolor":"white","showland":true,"showlakes":true,"lakecolor":"white"},"title":{"x":0.05},"mapbox":{"style":"light"}}},"xaxis":{"anchor":"y","domain":[0.0,1.0],"title":{"text":"Score"}},"yaxis":{"anchor":"x","domain":[0.0,1.0],"title":{"text":"Feature"}},"coloraxis":{"colorbar":{"title":{"text":"Score"}},"colorscale":[[0.0,"rgb(103,0,31)"],[0.1,"rgb(178,24,43)"],[0.2,"rgb(214,96,77)"],[0.3,"rgb(244,165,130)"],[0.4,"rgb(253,219,199)"],[0.5,"rgb(247,247,247)"],[0.6,"rgb(209,229,240)"],[0.7,"rgb(146,197,222)"],[0.8,"rgb(67,147,195)"],[0.9,"rgb(33,102,172)"],[1.0,"rgb(5,48,97)"]]},"legend":{"tracegroupgap":0},"margin":{"t":60},"barmode":"relative"}}

demo-project/data/09_tracking/linear_params.json/2023-01-16T15.55.27.850Z/linear_params.json

@@ -0,0 +1,6 @@
+{
+  "model_type": "LinearRegression",
+  "fit_intercept": true,
+  "copy_X": true,
+  "positive": false
+}

demo-project/data/09_tracking/linear_score.json/2023-01-16T15.55.27.850Z/linear_score.json

@@ -0,0 +1,3 @@
+{
+  "r2_score": 0.40296896595214116
+}

demo-project/data/09_tracking/rf_params.json/2023-01-16T15.55.27.850Z/rf_params.json

@@ -0,0 +1,15 @@
+{
+  "model_type": "RandomForestRegressor",
+  "n_estimators": 100,
+  "criterion": "squared_error",
+  "min_samples_split": 2,
+  "min_samples_leaf": 1,
+  "min_weight_fraction_leaf": 0,
+  "max_features": "auto",
+  "min_impurity_decrease": 0,
+  "bootstrap": true,
+  "oob_score": false,
+  "verbose": 0,
+  "warm_start": false,
+  "ccp_alpha": 0
+}

demo-project/data/09_tracking/rf_score.json/2023-01-16T15.55.27.850Z/rf_score.json

@@ -0,0 +1,3 @@
+{
+  "r2_score": 0.4113266331219384
+}

package/.pylintrc

@@ -50,7 +50,7 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=ungrouped-imports,bad-continuation,attribute-defined-outside-init,too-many-arguments,duplicate-code,fixme
+disable=ungrouped-imports,attribute-defined-outside-init,too-many-arguments,duplicate-code,fixme
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option
@@ -92,35 +92,22 @@ max-nested-blocks=5
 
 [BASIC]
 
-# Naming hint for argument names
-argument-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
-
 # Regular expression matching correct argument names
 argument-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
 
-# Naming hint for attribute names
-attr-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
 
 # Regular expression matching correct attribute names
 attr-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
 
 # Bad variable names which should always be refused, separated by a comma
 bad-names=foo,bar,baz,toto,tutu,tata
 
-# Naming hint for class attribute names
-class-attribute-name-hint=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$
-
 # Regular expression matching correct class attribute names
 class-attribute-rgx=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$
 
-# Naming hint for class names
-class-name-hint=[A-Z_][a-zA-Z0-9]+$
-
 # Regular expression matching correct class names
 class-rgx=[A-Z_][a-zA-Z0-9]+$
 
-# Naming hint for constant names
-const-name-hint=(([A-Z_][A-Z0-9_]*)|(__.*__))$
 
 # Regular expression matching correct constant names
 const-rgx=(([A-Z_][A-Z0-9_]*)|(__.*__))$
@@ -129,9 +116,6 @@ const-rgx=(([A-Z_][A-Z0-9_]*)|(__.*__))$
 # ones are exempt.
 docstring-min-length=-1
 
-# Naming hint for function names
-function-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
-
 # Regular expression matching correct function names
 function-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
 
@@ -141,21 +125,12 @@ good-names=i,j,k,ex,Run,_,id
 # Include a hint for the correct naming format with invalid-name
 include-naming-hint=no
 
-# Naming hint for inline iteration names
-inlinevar-name-hint=[A-Za-z_][A-Za-z0-9_]*$
-
 # Regular expression matching correct inline iteration names
 inlinevar-rgx=[A-Za-z_][A-Za-z0-9_]*$
 
-# Naming hint for method names
-method-name-hint=(([a-z][a-z0-9_]{2,60})|(_[a-z0-9_]*))$
-
 # Regular expression matching correct method names
 method-rgx=(([a-z][a-z0-9_]{2,60})|(_[a-z0-9_]*))$
 
-# Naming hint for module names
-module-name-hint=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$
-
 # Regular expression matching correct module names
 module-rgx=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$
 
@@ -171,9 +146,6 @@ no-docstring-rgx=^_
 # to this list to register other decorators that produce valid properties.
 property-classes=abc.abstractproperty
 
-# Naming hint for variable names
-variable-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
-
 # Regular expression matching correct variable names
 variable-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$
 
@@ -199,12 +171,6 @@ max-line-length=100
 # Maximum number of lines in a module
 max-module-lines=1000
 
-# List of optional constructs for which whitespace checking is disabled. `dict-
-# separator` is used to allow tabulation in dicts, etc.: {1  : 1,\n222: 2}.
-# `trailing-comma` allows a space between comma and closing bracket: (a, ).
-# `empty-line` allows space-only lines.
-no-space-check=trailing-comma,dict-separator
-
 # Allow the body of a class to be on the same line as the declaration if body
 # contains single statement.
 single-line-class-stmt=no

package/kedro_viz/api/graphql/serializers.py

@@ -85,7 +85,7 @@ def format_run_tracking_data(
         Dictionary with formatted tracking data for selected runs
 
     Example:
-        >>> from kedro.extras.datasets.tracking import MetricsDataSet
+        >>> from kedro_datasets.tracking import MetricsDataSet
         >>> tracking_data = {
         >>>     'My Favorite Sprint': {
         >>>         'bootstrap':0.8

package/kedro_viz/api/rest/responses.py

@@ -121,7 +121,7 @@ class Config:
         schema_extra = {
             "example": {
                 "filepath": "/my-kedro-project/data/03_primary/master_table.csv",
-                "type": "kedro.extras.datasets.pandas.csv_dataset.CSVDataSet",
+                "type": "pandas.csv_dataset.CSVDataSet",
                 "run_command": 'kedro run --to-outputs="master_table"',
             }
         }

package/kedro_viz/integrations/kedro/data_loader.py

@@ -10,12 +10,21 @@
 from typing import Any, Dict, Optional, Tuple
 
 from kedro import __version__
-from kedro.extras.datasets import (  # Safe since ImportErrors are suppressed within kedro.
-    json,
-    matplotlib,
-    plotly,
-    tracking,
-)
+
+try:
+    from kedro_datasets import (  # isort:skip
+        json,
+        matplotlib,
+        plotly,
+        tracking,
+    )
+except ImportError:
+    from kedro.extras.datasets import (  # Safe since ImportErrors are suppressed within kedro.
+        json,
+        matplotlib,
+        plotly,
+        tracking,
+    )
 from kedro.io import DataCatalog
 from kedro.io.core import get_filepath_str
 from kedro.pipeline import Pipeline
@@ -150,16 +159,16 @@ def matplotlib_writer_load(dataset: matplotlib.MatplotlibWriter) -> str:
             base64_bytes = base64.b64encode(img_file.read())
         return base64_bytes.decode("utf-8")
 
-    matplotlib.MatplotlibWriter._load = matplotlib_writer_load  # type:ignore
+    matplotlib.MatplotlibWriter._load = matplotlib_writer_load
 
 if hasattr(plotly, "JSONDataSet"):
-    plotly.JSONDataSet._load = json.JSONDataSet._load  # type:ignore
+    plotly.JSONDataSet._load = json.JSONDataSet._load
 
 if hasattr(plotly, "PlotlyDataSet"):
-    plotly.PlotlyDataSet._load = json.JSONDataSet._load  # type:ignore
+    plotly.PlotlyDataSet._load = json.JSONDataSet._load
 
 if hasattr(tracking, "JSONDataSet"):
-    tracking.JSONDataSet._load = json.JSONDataSet._load  # type:ignore
+    tracking.JSONDataSet._load = json.JSONDataSet._load
 
 if hasattr(tracking, "MetricsDataSet"):
-    tracking.MetricsDataSet._load = json.JSONDataSet._load  # type:ignore
+    tracking.MetricsDataSet._load = json.JSONDataSet._load

package/kedro_viz/models/experiment_tracking.py

@@ -51,13 +51,17 @@ class TrackingDatasetGroup(str, Enum):
     JSON = "json"
 
 
-# pylint: disable=line-too-long
+# The keys will match a dataset type with any prefix, e.g.
+# plotly.plotly_dataset.PlotlyDataSet would include the following:
+# plotly.plotly_dataset.PlotlyDataSet
+# plotly.plotly_dataset.PlotlyDataSet
+# my.custom.path.to.plotly.plotly_dataset.PlotlyDataSet
 TRACKING_DATASET_GROUPS = {
-    "kedro.extras.datasets.plotly.plotly_dataset.PlotlyDataSet": TrackingDatasetGroup.PLOT,
-    "kedro.extras.datasets.plotly.json_dataset.JSONDataSet": TrackingDatasetGroup.PLOT,
-    "kedro.extras.datasets.matplotlib.matplotlib_writer.MatplotlibWriter": TrackingDatasetGroup.PLOT,
-    "kedro.extras.datasets.tracking.metrics_dataset.MetricsDataSet": TrackingDatasetGroup.METRIC,
-    "kedro.extras.datasets.tracking.json_dataset.JSONDataSet": TrackingDatasetGroup.JSON,
+    "plotly.plotly_dataset.PlotlyDataSet": TrackingDatasetGroup.PLOT,
+    "plotly.json_dataset.JSONDataSet": TrackingDatasetGroup.PLOT,
+    "matplotlib.matplotlib_writer.MatplotlibWriter": TrackingDatasetGroup.PLOT,
+    "tracking.metrics_dataset.MetricsDataSet": TrackingDatasetGroup.METRIC,
+    "tracking.json_dataset.JSONDataSet": TrackingDatasetGroup.JSON,
 }
 
 
@@ -67,7 +71,7 @@ class TrackingDatasetModel:
 
     dataset_name: str
     # dataset is the actual dataset instance, whereas dataset_type is a string.
-    # e.g. "kedro.extras.datasets.tracking.metrics_dataset.MetricsDataSet"
+    # e.g. "tracking.metrics_dataset.MetricsDataSet"
     dataset: AbstractVersionedDataSet
     dataset_type: str = field(init=False)
     # runs is a mapping from run_id to loaded data.
@@ -110,4 +114,6 @@ def load_tracking_data(self, run_id: str):
 
 
 def get_dataset_type(dataset: AbstractVersionedDataSet) -> str:
-    return f"{dataset.__class__.__module__}.{dataset.__class__.__qualname__}"
+    class_name = f"{dataset.__class__.__qualname__}"
+    _, dataset_type, dataset_file = f"{dataset.__class__.__module__}".rsplit(".", 2)
+    return f"{dataset_type}.{dataset_file}.{class_name}"