From 28e52d5c0327dd7bb7f6b4211c7908ade6150104 Mon Sep 17 00:00:00 2001
From: Ewout ter Hoeven <E.M.terHoeven@student.tudelft.nl>
Date: Thu, 26 Sep 2024 09:49:29 +0200
Subject: [PATCH] Deprecate `initialize_data_collector` (#2327)

Deprecate the Model's `initialize_data_collector`.

Back when we had a schedulers it was logical, but now you can't forget to add Agents to a scheduler since that happens automatically.

The other reason was the `batch_run` assuming a Model attribute called `datacollector` to be present. An error in `batch_run` itself is far more robust to enforce that behavior, so that's added..

It also removes the behavior of directly collecting after init. Now you can do it when you want. It makes using the DataCollector explicit.

The migration guide was updated. Basically, replace:
```python
self.initialize_data_collector(...)
```

With:
```python
self.datacollector = DataCollector(...)
```
---
 docs/migration_guide.md     | 16 ++++++++++++++++
 mesa/batchrunner.py         |  7 +++++++
 mesa/model.py               | 15 +++++++--------
 tests/test_datacollector.py | 27 ++-------------------------
 4 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/docs/migration_guide.md b/docs/migration_guide.md
index 51789cabcdd..4741b8fb28b 100644
--- a/docs/migration_guide.md
+++ b/docs/migration_guide.md
@@ -295,3 +295,19 @@ def show_steps(model):
 
 SolaraViz(model, components=[show_steps])
 ```
+
+### Other changes
+#### Removal of Model.initialize_data_collector
+The `initialize_data_collector` in the Model class is removed. In the Model class, replace:
+
+Replace:
+```python
+self.initialize_data_collector(...)
+```
+
+With:
+```python
+self.datacollector = DataCollector(...)
+```
+
+- Ref: [PR #2327](https://github.com/projectmesa/mesa/pull/2327), Mesa-examples [PR #208](https://github.com/projectmesa/mesa-examples/pull/208))
diff --git a/mesa/batchrunner.py b/mesa/batchrunner.py
index d50476c6f52..2a6d50a3b77 100644
--- a/mesa/batchrunner.py
+++ b/mesa/batchrunner.py
@@ -38,6 +38,9 @@ def batch_run(
     Returns:
         List[Dict[str, Any]]
 
+    Notes:
+        batch_run assumes the model has a `datacollector` attribute that has a DataCollector object initialized.
+
     """
     runs_list = []
     run_id = 0
@@ -173,6 +176,10 @@ def _collect_data(
     step: int,
 ) -> tuple[dict[str, Any], list[dict[str, Any]]]:
     """Collect model and agent data from a model using mesas datacollector."""
+    if not hasattr(model, "datacollector"):
+        raise AttributeError(
+            "The model does not have a datacollector attribute. Please add a DataCollector to your model."
+        )
     dc = model.datacollector
 
     model_data = {param: values[step] for param, values in dc.model_vars.items()}
diff --git a/mesa/model.py b/mesa/model.py
index a3227c10cdc..d9566f44f3c 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -210,14 +210,13 @@ def initialize_data_collector(
             tables: tables to collect
 
         """
-        if not hasattr(self, "schedule") or self.schedule is None:
-            raise RuntimeError(
-                "You must initialize the scheduler (self.schedule) before initializing the data collector."
-            )
-        if self.schedule.get_agent_count() == 0:
-            raise RuntimeError(
-                "You must add agents to the scheduler before initializing the data collector."
-            )
+        warnings.warn(
+            "initialize_data_collector() is deprecated. Please use the DataCollector class directly. "
+            "by using `self.datacollector = DataCollector(...)`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         self.datacollector = DataCollector(
             model_reporters=model_reporters,
             agent_reporters=agent_reporters,
diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
index b2760fc1b44..5c2575df31b 100644
--- a/tests/test_datacollector.py
+++ b/tests/test_datacollector.py
@@ -70,7 +70,7 @@ def __init__(self):  # noqa: D107
         self.n = 10
         for i in range(1, self.n + 1):
             self.schedule.add(MockAgent(self, val=i))
-        self.initialize_data_collector(
+        self.datacollector = DataCollector(
             model_reporters={
                 "total_agents": lambda m: m.schedule.get_agent_count(),
                 "model_value": "model_val",
@@ -132,6 +132,7 @@ class TestDataCollector(unittest.TestCase):
     def setUp(self):
         """Create the model and run it a set number of steps."""
         self.model = MockModel()
+        self.model.datacollector.collect(self.model)
         for i in range(7):
             if i == 4:
                 self.model.schedule.remove(self.model.schedule._agents[3])
@@ -234,30 +235,6 @@ def test_exports(self):
             table_df = data_collector.get_table_dataframe("not a real table")
 
 
-class TestDataCollectorInitialization(unittest.TestCase):
-    """Tests for DataCollector initialization."""
-
-    def setUp(self):  # noqa: D102
-        self.model = Model()
-
-    def test_initialize_before_scheduler(self):  # noqa: D102
-        with self.assertRaises(RuntimeError) as cm:
-            self.model.initialize_data_collector()
-        self.assertEqual(
-            str(cm.exception),
-            "You must initialize the scheduler (self.schedule) before initializing the data collector.",
-        )
-
-    def test_initialize_before_agents_added_to_scheduler(self):  # noqa: D102
-        with self.assertRaises(RuntimeError) as cm:
-            self.model.schedule = BaseScheduler(self)
-            self.model.initialize_data_collector()
-        self.assertEqual(
-            str(cm.exception),
-            "You must add agents to the scheduler before initializing the data collector.",
-        )
-
-
 class TestDataCollectorWithAgentTypes(unittest.TestCase):
     """Tests for DataCollector with agent-type-specific reporters."""