[Tidy] Remove component to data mapping from data manager (#451)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

vizro-core/changelog.d/20240502_175131_antony.milne_data_source_mapping.md

@@ -0,0 +1,48 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+<!--
+### Highlights ✨
+
+- A bullet item for the Highlights ✨ category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Removed
+
+- A bullet item for the Removed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Added
+
+- A bullet item for the Added category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Changed
+
+- A bullet item for the Changed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Deprecated
+
+- A bullet item for the Deprecated category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Fixed
+
+- A bullet item for the Fixed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Security
+
+- A bullet item for the Security category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->

vizro-core/src/vizro/actions/_actions_utils.py

@@ -173,8 +173,8 @@ def _get_filtered_data(
 ) -> Dict[ModelID, pd.DataFrame]:
     filtered_data = {}
     for target in targets:
-        data_frame = data_manager._get_component_data(target)
-
+        data_source_name = model_manager[target]["data_frame"]
+        data_frame = data_manager[data_source_name].load()
         data_frame = _apply_filters(data_frame=data_frame, ctds_filters=ctds_filters, target=target)
         data_frame = _apply_filter_interaction(
             data_frame=data_frame, ctds_filter_interaction=ctds_filter_interaction, target=target

vizro-core/src/vizro/managers/_data_manager.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import logging
-import os
 import warnings
 from typing import Any, Callable, Dict, Optional, Protocol, Union
 
@@ -15,10 +14,8 @@
 
 logger = logging.getLogger(__name__)
 
-# Really ComponentID and DataSourceName should be NewType and not just aliases but then for a user's code to type check
+# Really DataSourceName should be NewType and not just aliases but then for a user's code to type check
 # correctly they would need to cast all strings to these types.
-# TODO: remove these type aliases once have moved component to data mapping to models
-ComponentID = str
 DataSourceName = str
 pd_DataFrameCallable = Callable[[], pd.DataFrame]
 
@@ -186,7 +183,6 @@ class DataManager:
 
     def __init__(self):
         self.__data: Dict[DataSourceName, Union[_DynamicData, _StaticData]] = {}
-        self.__component_to_data: Dict[ComponentID, DataSourceName] = {}
         self._frozen_state = False
         self.cache = Cache(config={"CACHE_TYPE": "NullCache"})
         # In future, possibly we will accept just a config dict. Would need to work out whether to handle merging with
@@ -227,33 +223,6 @@ def __getitem__(self, name: DataSourceName) -> Union[_DynamicData, _StaticData]:
         except KeyError as exc:
             raise KeyError(f"Data source {name} does not exist.") from exc
 
-    @_state_modifier
-    def _add_component(self, component_id: ComponentID, name: DataSourceName):
-        """Adds a mapping from `component_id` to `name`."""
-        # TODO: once have removed self.__component_to_data, we shouldn't need this function any more.
-        #  Maybe always updated capturedcallable data_frame to  data source name string then.
-        if name not in self.__data:
-            raise KeyError(f"Data source {name} does not exist.")
-        if component_id in self.__component_to_data:
-            raise ValueError(
-                f"Component with id={component_id} already exists and is mapped to data "
-                f"{self.__component_to_data[component_id]}. Components must uniquely map to a data source across the "
-                f"whole dashboard. If you are working from a Jupyter Notebook, please either restart the kernel, or "
-                f"use 'from vizro import Vizro; Vizro._reset()`."
-            )
-        self.__component_to_data[component_id] = name
-
-    def _get_component_data(self, component_id: ComponentID) -> pd.DataFrame:
-        # TODO: once have removed self.__component_to_data, we shouldn't need this function any more. Calling
-        #  functions would just do data_manager[name].load().
-        """Returns the original data for `component_id`."""
-        if component_id not in self.__component_to_data:
-            raise KeyError(f"Component {component_id} does not exist. You need to call add_component first.")
-        name = self.__component_to_data[component_id]
-
-        logger.debug("Loading data %s on process %s", name, os.getpid())
-        return self[name].load()
-
     def _clear(self):
         # We do not actually call self.cache.clear() because (a) it would only work when self._cache_has_app is True,
         # which is not the case when e.g. Vizro._reset is called, and (b) because we do not want to accidentally

vizro-core/src/vizro/models/_components/_components_utils.py

@@ -24,28 +24,30 @@ def _callable_mode_validator_factory(mode: str):
     return validator("figure", allow_reuse=True)(check_callable_mode)
 
 
-def _process_callable_data_frame(captured_callable, values):
+def _process_callable_data_frame(captured_callable):
+    # Possibly all this validator's functionality should move into CapturedCallable (or a subclass of it) in the
+    # future. This would mean that data is added to the data manager outside the context of a dashboard though,
+    # which might not be desirable.
     data_frame = captured_callable["data_frame"]
 
     if isinstance(data_frame, str):
         # Named data source, which could be dynamic or static. This means px.scatter("iris") from the Python API and
-        # specification of "data_frame": "iris" through JSON. In these cases, data already exists in the data manager
-        # and just needs to be linked to the component.
-        data_source_name = data_frame
-    else:
-        # Unnamed data source, which must be a pd.DataFrame and hence static data. This means px.scatter(pd.DataFrame())
-        # and is only possible from the Python API. Extract dataframe from the captured function and put it into the
-        # data manager.
-        # Unlike with model_manager, it doesn't matter if the random seed is different across workers here. So long as
-        # we always fetch static data from the data manager by going through the appropriate Figure component, the right
-        # data source name will be fetched. It also doesn't matter if multiple Figures with the same underlying data
-        # each have their own entry in the data manager, since the underlying pd.DataFrame will still be the same and
-        # not copied into each one, so no memory is wasted.
-        logger.debug("Adding data to data manager for Figure with id %s", values["id"])
-        data_source_name = str(uuid.uuid4())
-        data_manager[data_source_name] = data_frame
-
-    data_manager._add_component(values["id"], data_source_name)
-    # No need to keep the data in the captured function any more so remove it to save memory.
-    del captured_callable["data_frame"]
+        # specification of "data_frame": "iris" through JSON. In these cases, data already exists in the data manager.
+        return captured_callable
+
+    # Unnamed data source, which must be a pd.DataFrame and hence static data. This means px.scatter(pd.DataFrame())
+    # and is only possible from the Python API. Extract dataframe from the captured function and put it into the
+    # data manager.
+    # Unlike with model_manager, it doesn't matter if the random seed is different across workers here. So long as
+    # we always fetch static data from the data manager by going through the appropriate Figure component, the right
+    # data source name will be fetched. It also doesn't matter if multiple Figures with the same underlying data
+    # each have their own entry in the data manager, since the underlying pd.DataFrame will still be the same and
+    # not copied into each one, so no memory is wasted.
+    # Replace the "data_frame" argument in the captured callable with the data_source_name for consistency with
+    # dynamic data and to save memory. This way we always access data via the same interface regardless of whether it's
+    # static or dynamic.
+    data_source_name = str(uuid.uuid4())
+    data_manager[data_source_name] = data_frame
+    captured_callable["data_frame"] = data_source_name
+
     return captured_callable

vizro-core/src/vizro/models/_components/ag_grid.py

@@ -51,7 +51,7 @@ class AgGrid(VizroBaseModel):
 
     # Convenience wrapper/syntactic sugar.
     def __call__(self, **kwargs):
-        kwargs.setdefault("data_frame", data_manager._get_component_data(self.id))
+        kwargs.setdefault("data_frame", data_manager[self["data_frame"]].load())
         figure = self.figure(**kwargs)
         figure.id = self._input_component_id
         return figure

vizro-core/src/vizro/models/_components/graph.py

@@ -50,7 +50,10 @@ class Graph(VizroBaseModel):
 
     # Convenience wrapper/syntactic sugar.
     def __call__(self, **kwargs):
-        kwargs.setdefault("data_frame", data_manager._get_component_data(str(self.id)))
+        # This default value is not actually used anywhere at the moment since __call__ is always used with data_frame
+        # specified. It's here to match Table and AgGrid and because we might want to use __call__ more in future.
+        # If the functionality of process_callable_data_frame moves to CapturedCallable then this would move there too.
+        kwargs.setdefault("data_frame", data_manager[self["data_frame"]].load())
         fig = self.figure(**kwargs)
 
         # Remove top margin if title is provided

vizro-core/src/vizro/models/_components/table.py

@@ -50,7 +50,7 @@ class Table(VizroBaseModel):
 
     # Convenience wrapper/syntactic sugar.
     def __call__(self, **kwargs):
-        kwargs.setdefault("data_frame", data_manager._get_component_data(self.id))
+        kwargs.setdefault("data_frame", data_manager[self["data_frame"]].load())
         figure = self.figure(**kwargs)
         figure.id = self._input_component_id
         return figure

vizro-core/src/vizro/models/_controls/filter.py

@@ -112,14 +112,20 @@ def _set_targets(self):
             for component_id in model_manager._get_page_model_ids_with_figure(
                 page_id=model_manager._get_model_page_id(model_id=ModelID(str(self.id)))
             ):
-                data_frame = data_manager._get_component_data(component_id)
+                # TODO: consider making a helper method in data_manager or elsewhere to reduce this operation being
+                #  duplicated across Filter so much, and/or consider storing the result to avoid repeating it.
+                #  Need to think about this in connection with how to update filters on the fly and duplicated calls
+                #  issue outlined in https://github.com/mckinsey/vizro/pull/398#discussion_r1559120849.
+                data_source_name = model_manager[component_id]["data_frame"]
+                data_frame = data_manager[data_source_name].load()
                 if self.column in data_frame.columns:
                     self.targets.append(component_id)
             if not self.targets:
                 raise ValueError(f"Selected column {self.column} not found in any dataframe on this page.")
 
     def _set_column_type(self):
-        data_frame = data_manager._get_component_data(self.targets[0])
+        data_source_name = model_manager[self.targets[0]]["data_frame"]
+        data_frame = data_manager[data_source_name].load()
 
         if is_numeric_dtype(data_frame[self.column]):
             self._column_type = "numerical"
@@ -146,7 +152,8 @@ def _set_numerical_and_temporal_selectors_values(self):
             min_values = []
             max_values = []
             for target_id in self.targets:
-                data_frame = data_manager._get_component_data(target_id)
+                data_source_name = model_manager[target_id]["data_frame"]
+                data_frame = data_manager[data_source_name].load()
                 min_values.append(data_frame[self.column].min())
                 max_values.append(data_frame[self.column].max())
 
@@ -173,7 +180,8 @@ def _set_categorical_selectors_options(self):
         if isinstance(self.selector, SELECTORS["categorical"]) and not self.selector.options:
             options = set()
             for target_id in self.targets:
-                data_frame = data_manager._get_component_data(target_id)
+                data_source_name = model_manager[target_id]["data_frame"]
+                data_frame = data_manager[data_source_name].load()
                 options |= set(data_frame[self.column])
 
             self.selector.options = sorted(options)

vizro-core/src/vizro/models/types.py

@@ -138,9 +138,9 @@ def __getitem__(self, arg_name: str):
         """Gets the value of a bound argument."""
         return self.__bound_arguments[arg_name]
 
-    def __delitem__(self, arg_name: str):
-        """Deletes a bound argument."""
-        del self.__bound_arguments[arg_name]
+    def __setitem__(self, arg_name: str, value):
+        """Sets the value of a bound argument."""
+        self.__bound_arguments[arg_name] = value
 
     @property
     def _arguments(self):

vizro-core/tests/unit/vizro/models/_components/test_ag_grid.py

@@ -91,15 +91,11 @@ class TestProcessAgGridDataFrame:
     def test_process_figure_data_frame_str_df(self, dash_ag_grid_with_str_dataframe, gapminder):
         data_manager["gapminder"] = gapminder
         ag_grid = vm.AgGrid(id="ag_grid", figure=dash_ag_grid_with_str_dataframe)
-        assert data_manager._get_component_data("ag_grid").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            ag_grid["data_frame"]
+        assert data_manager[ag_grid["data_frame"]].load().equals(gapminder)
 
     def test_process_figure_data_frame_df(self, standard_ag_grid, gapminder):
         ag_grid = vm.AgGrid(id="ag_grid", figure=standard_ag_grid)
-        assert data_manager._get_component_data("ag_grid").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            ag_grid["data_frame"]
+        assert data_manager[ag_grid["data_frame"]].load().equals(gapminder)
 
 
 class TestPreBuildAgGrid:

vizro-core/tests/unit/vizro/models/_components/test_graph.py

@@ -123,15 +123,11 @@ class TestProcessGraphDataFrame:
     def test_process_figure_data_frame_str_df(self, standard_px_chart_with_str_dataframe, gapminder):
         data_manager["gapminder"] = gapminder
         graph = vm.Graph(id="graph", figure=standard_px_chart_with_str_dataframe)
-        assert data_manager._get_component_data("graph").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            graph["data_frame"]
+        assert data_manager[graph["data_frame"]].load().equals(gapminder)
 
     def test_process_figure_data_frame_df(self, standard_px_chart, gapminder):
         graph = vm.Graph(id="graph", figure=standard_px_chart)
-        assert data_manager._get_component_data("graph").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            graph["data_frame"]
+        assert data_manager[graph["data_frame"]].load().equals(gapminder)
 
 
 class TestBuild:

vizro-core/tests/unit/vizro/models/_components/test_table.py

@@ -91,15 +91,11 @@ class TestProcessTableDataFrame:
     def test_process_figure_data_frame_str_df(self, dash_table_with_str_dataframe, gapminder):
         data_manager["gapminder"] = gapminder
         table = vm.Table(id="table", figure=dash_table_with_str_dataframe)
-        assert data_manager._get_component_data("table").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            table["data_frame"]
+        assert data_manager[table["data_frame"]].load().equals(gapminder)
 
     def test_process_figure_data_frame_df(self, standard_dash_table, gapminder):
         table = vm.Table(id="table", figure=standard_dash_table)
-        assert data_manager._get_component_data("table").equals(gapminder)
-        with pytest.raises(KeyError, match="'data_frame'"):
-            table["data_frame"]
+        assert data_manager[table["data_frame"]].load().equals(gapminder)
 
 
 class TestPreBuildTable:

vizro-core/tests/unit/vizro/models/test_types.py

@@ -76,11 +76,9 @@ def test_getitem_unknown_args(self, captured_callable):
         with pytest.raises(KeyError):
             captured_callable["c"]
 
-    def test_delitem(self, captured_callable):
-        del captured_callable["a"]
-
-        with pytest.raises(KeyError):
-            captured_callable["a"]
+    def test_setitem(self, captured_callable):
+        captured_callable["a"] = 2
+        assert captured_callable["a"] == 2
 
 
 @pytest.mark.parametrize(