Fix: Add unique_id and agent_type to agent reports

Author: AgrimGupta04

mesa_frames/concrete/datacollector.py

@@ -72,7 +72,7 @@ def __init__(
         self,
         model: Model,
         model_reporters: dict[str, Callable] | None = None,
-        agent_reporters: dict[str, str] | None = None,
+        agent_reporters: dict[str, str | Callable] | None = None,  # <-- ALLOWS CALLABLE
         trigger: Callable[[Any], bool] | None = None,
         reset_memory: bool = True,
         storage: Literal[
@@ -92,7 +92,10 @@ def __init__(
         model_reporters : dict[str, Callable] | None
             Functions to collect data at the model level.
         agent_reporters : dict[str, str | Callable] | None
-            Attributes or functions to collect data at the agent level.
+            (MODIFIED) A dictionary mapping new column names to existing 
+            column names (str) or callables. Callables are not currently
+            processed by the agent data collector but are allowed for API compatibility.
+            Example: {"agent_wealth": "wealth", "age_in_years": "age"}
         trigger : Callable[[Any], bool] | None
             A function(model) -> bool that determines whether to collect data.
         reset_memory : bool
@@ -108,10 +111,14 @@ def __init__(
         """
         if agent_reporters:
             for key, value in agent_reporters.items():
-                if not isinstance(value, str):
+                if not isinstance(key, str):
                     raise TypeError(
-                        f"Agent reporter for '{key}' must be a string (the column name), "
-                        f"not a {type(value)}. Callable reporters are not supported for agents."
+                        f"Agent reporter keys must be strings (the final column name), not a {type(key)}."
+                    )
+                if not (isinstance(value, str) or callable(value)):
+                    raise TypeError(
+                        f"Agent reporter for '{key}' must be either a string (the source column name) "
+                        f"or a callable (function taking an agent and returning a value), not a {type(value)}."
                     )
 
         super().__init__(

tests/test_datacollector.py

@@ -171,13 +171,28 @@ def test__init__(self, fix1_model, postgres_uri):
         #         model=model, model_reporters={"total_agents": "sum"}
         #     )
 
+        model.test_dc = DataCollector(
+            model=model, 
+            agent_reporters={"wealth": lambda m: 1}
+        )
+        assert model.test_dc is not None
+
         with pytest.raises(
             TypeError,
-            match="Agent reporter for 'wealth' must be a string",
+            match="Agent reporter keys must be strings",
         ):
             model.test_dc = DataCollector(
-                model=model,
-                agent_reporters={"wealth": lambda m: 1},  # This is now illegal
+                model=model, 
+                agent_reporters={123: "wealth"}
+            )
+
+        with pytest.raises(
+            TypeError,
+            match="must be either a string .* or a callable",
+        ):
+            model.test_dc = DataCollector(
+                model=model, 
+                agent_reporters={"wealth": 123} # This is illegal
             )
 
         with pytest.raises(
@@ -729,3 +744,47 @@ def test_batch_save(self, fix2_model):
             assert sorted(agent_df_step2_batch1["wealth"].to_list()) == sorted(
                 expected_wealth_s2b1
             )
+
+    def test_collect_no_agentsets_list(self, fix1_model, caplog):
+        """Tests that the collector logs an error and exits gracefully if _agentsets is missing."""
+        model = fix1_model
+        del model.sets._agentsets
+
+        dc = DataCollector(model=model, agent_reporters={"wealth": "wealth"})
+        dc.collect()
+        
+        assert "could not find '_agentsets'" in caplog.text
+        assert dc.data["agent"].shape == (0, 0)
+
+    def test_collect_agent_set_no_df(self, fix1_model, caplog):
+        """Tests that the collector logs a warning and skips a set if it has no .df attribute."""
+
+        class NoDfSet:
+            def __init__(self):
+                self.__class__ = type("NoDfSet", (object,), {})
+        
+        fix1_model.sets._agentsets.append(NoDfSet())
+        
+        fix1_model.dc = DataCollector(fix1_model, agent_reporters={"wealth": "wealth", "age": "age"})
+        fix1_model.dc.collect()
+        
+        assert "has no 'df' attribute" in caplog.text
+        assert fix1_model.dc.data["agent"].shape == (12, 7) 
+
+    def test_collect_df_no_unique_id(self, fix1_model, caplog):
+        """Tests that the collector logs a warning and skips a set if its df has no unique_id."""
+        bad_set = fix1_model.sets._agentsets[0] 
+        bad_set._df = bad_set._df.drop("unique_id") 
+            
+        fix1_model.dc = DataCollector(fix1_model, agent_reporters={"wealth": "wealth", "age": "age"})
+        fix1_model.dc.collect()
+        
+        assert "has no 'unique_id' column" in caplog.text
+        assert fix1_model.dc.data["agent"].shape == (8, 7)
+
+    def test_collect_no_matching_reporters(self, fix1_model):
+        """Tests that the collector returns an empty frame if no reporters match any columns."""
+        fix1_model.dc = DataCollector(fix1_model, agent_reporters={"baz": "foo", "qux": "bar"})
+        fix1_model.dc.collect()
+
+        assert fix1_model.dc.data["agent"].shape == (0, 0)