Allow DataCollector to accept property names

examples/Schelling/model.py

@@ -56,7 +56,7 @@ def __init__(self, height, width, density, minority_pc, homophily):
 
         self.happy = 0
         self.datacollector = DataCollector(
-            {"happy": lambda m: m.happy},  # Model-level count of happy agents
+            {"happy": "happy"},  # Model-level count of happy agents
             # For testing purposes, agent's individual x and y
             {"x": lambda a: a.pos[0], "y": lambda a: a.pos[1]})
 

mesa/datacollection.py

@@ -104,6 +104,8 @@ def _new_model_reporter(self, reporter_name, reporter_function):
                                given a model instance.
 
         """
+        if type(reporter_function) is str:
+            reporter_function = self._make_attribute_collector(reporter_function)
         self.model_reporters[reporter_name] = reporter_function
         self.model_vars[reporter_name] = []
 
@@ -116,6 +118,8 @@ def _new_agent_reporter(self, reporter_name, reporter_function):
                                given an agent object.
 
         """
+        if type(reporter_function) is str:
+            reporter_function = self._make_attribute_collector(reporter_function)
         self.agent_reporters[reporter_name] = reporter_function
         self.agent_vars[reporter_name] = []
 
@@ -164,6 +168,19 @@ def add_table_row(self, table_name, row, ignore_missing=False):
             else:
                 raise Exception("Could not insert row with missing column")
 
+    @staticmethod
+    def _make_attribute_collector(attr):
+        '''
+        Create a function which collects the value of a named attribute
+        '''
+
+        def attr_collector(obj):
+            return getattr(obj, attr)
+
+        return attr_collector
+
+    @staticmethod
+
     def get_model_vars_dataframe(self):
         """ Create a pandas DataFrame from the model variables.