edsl 0.1.47__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/agents/{AgentList.py → agent_list.py}

@@ -4,26 +4,33 @@
 from __future__ import annotations
 import csv
 import sys
+import random
+import logging
+from collections import defaultdict
+from itertools import product
+
+
 from collections import UserList
 from collections.abc import Iterable
-
 from typing import Any, List, Optional, Union, TYPE_CHECKING
 
 from simpleeval import EvalWithCompoundTypes, NameNotDefined
 
-from edsl.Base import Base
-from edsl.utilities.remove_edsl_version import remove_edsl_version
-from edsl.exceptions.agents import AgentListError
-from edsl.utilities.is_notebook import is_notebook
-from edsl.results.ResultsExportMixin import ResultsExportMixin
-import logging
+from ..base import Base
+from ..utilities import is_notebook, remove_edsl_version, dict_hash
+from ..dataset.dataset_operations_mixin import AgentListOperationsMixin
+
+from .agent import Agent
+
+from .exceptions import AgentListError
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from edsl.scenarios.ScenarioList import ScenarioList
-    from edsl.agents.Agent import Agent
+    from ..scenarios import ScenarioList
+    from ..agents import Agent
     from pandas import DataFrame
+    from ..dataset import Dataset
 
 
 def is_iterable(obj):
@@ -36,30 +43,66 @@ class EmptyAgentList:
 
 
 # ResultsExportMixin,
-class AgentList(UserList, Base):
-    """A list of Agents."""
+class AgentList(UserList, Base, AgentListOperationsMixin):
+    """A list of Agents with additional functionality for manipulation and analysis.
+
+    The AgentList class extends Python's UserList to provide a container for Agent objects
+    with methods for filtering, transforming, and analyzing collections of agents.
+
+
+    >>> AgentList.example().to_scenario_list()
+    ScenarioList([Scenario({'age': 22, 'hair': 'brown', 'height': 5.5}), Scenario({'age': 22, 'hair': 'brown', 'height': 5.5})])
+
+    >>> AgentList.example().to_dataset()
+    Dataset([{'age': [22, 22]}, {'hair': ['brown', 'brown']}, {'height': [5.5, 5.5]}])
+
+    >>> AgentList.example().to_pandas()
+       age   hair  height
+    0   22  brown     5.5
+    1   22  brown     5.5
+
+
+    """
 
     __documentation__ = (
         "https://docs.expectedparrot.com/en/latest/agents.html#agentlist-class"
     )
 
-    def __init__(self, data: Optional[list["Agent"]] = None):
+    def __init__(self, data: Optional[list["Agent"]] = None, codebook: Optional[dict[str, str]] = None):
         """Initialize a new AgentList.
 
-        :param data: A list of Agents.
+        >>> from edsl import Agent
+        >>> al = AgentList([Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5}),
+        ...                Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})])
+        >>> al
+        AgentList([Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5}), Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})])
+        >>> al_with_codebook = AgentList([Agent(traits = {'age': 22})], codebook={'age': 'Age in years'})
+        >>> al_with_codebook[0].codebook
+        {'age': 'Age in years'}
+
+        Args:
+            data: A list of Agent objects. If None, creates an empty AgentList.
+            codebook: Optional dictionary mapping trait names to descriptions. 
+                      If provided, will be applied to all agents in the list.
         """
         if data is not None:
             super().__init__(data)
         else:
             super().__init__()
+            
+        # Apply codebook to all agents if provided
+        if codebook is not None:
+            self.set_codebook(codebook)
 
     def shuffle(self, seed: Optional[str] = None) -> AgentList:
-        """Shuffle the AgentList.
+        """Randomly shuffle the agents in place.
 
-        :param seed: The seed for the random number generator.
-        """
-        import random
+        Args:
+            seed: Optional seed for the random number generator to ensure reproducibility.
 
+        Returns:
+            AgentList: The shuffled AgentList (self).
+        """
         if seed is not None:
             random.seed(seed)
         random.shuffle(self.data)
@@ -68,64 +111,52 @@ class AgentList(UserList, Base):
     def sample(self, n: int, seed: Optional[str] = None) -> AgentList:
         """Return a random sample of agents.
 
-        :param n: The number of agents to sample.
-        :param seed: The seed for the random number generator.
+        Args:
+            n: The number of agents to sample.
+            seed: Optional seed for the random number generator to ensure reproducibility.
+
+        Returns:
+            AgentList: A new AgentList containing the sampled agents.
         """
-        import random
 
         if seed:
             random.seed(seed)
         return AgentList(random.sample(self.data, n))
 
-    def to_pandas(self) -> "DataFrame":
-        """Return a pandas DataFrame.
-
-        >>> from edsl.agents.Agent import Agent
-        >>> al = AgentList([Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5}), Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})])
-        >>> al.to_pandas()
-           age   hair  height
-        0   22  brown     5.5
-        1   22  brown     5.5
-        """
-        return self.to_scenario_list().to_pandas()
-
-    def tally(
-        self, *fields: Optional[str], top_n: Optional[int] = None, output="Dataset"
-    ) -> Union[dict, "Dataset"]:
-        """Tally the values of a field or perform a cross-tab of multiple fields.
+    def duplicate(self) -> AgentList:
+        """Create a deep copy of the AgentList.
 
-        :param fields: The field(s) to tally, multiple fields for cross-tabulation.
+        Returns:
+            AgentList: A new AgentList containing copies of all agents.
 
-        >>> al = AgentList.example()
-        >>> al.tally('age')
-        Dataset([{'age': [22]}, {'count': [2]}])
-        """
-        return self.to_scenario_list().tally(*fields, top_n=top_n, output=output)
-
-    def duplicate(self):
-        """Duplicate the AgentList.
-
-        >>> al = AgentList.example()
-        >>> al2 = al.duplicate()
-        >>> al2 == al
-        True
-        >>> id(al2) == id(al)
-        False
+        Examples:
+            >>> al = AgentList.example()
+            >>> al2 = al.duplicate()
+            >>> al2 == al
+            True
+            >>> id(al2) == id(al)
+            False
         """
         return AgentList([a.duplicate() for a in self.data])
 
-    def rename(self, old_name, new_name) -> AgentList:
-        """Rename a trait in the AgentList.
+    def rename(self, old_name: str, new_name: str) -> AgentList:
+        """Rename a trait across all agents in the list.
 
-        :param old_name: The old name of the trait.
-        :param new_name: The new name of the trait.
-        :param inplace: Whether to rename the trait in place.
+        Args:
+            old_name: The current name of the trait.
+            new_name: The new name to assign to the trait.
 
-        >>> from edsl.agents.Agent import Agent
-        >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}), Agent(traits = {'a': 1, 'b': 2})])
-        >>> al2 = al.rename('a', 'c')
-        >>> assert al2 == AgentList([Agent(traits = {'c': 1, 'b': 1}), Agent(traits = {'c': 1, 'b': 2})])
-        >>> assert al != al2
+        Returns:
+            AgentList: A new AgentList with the renamed trait.
+
+        Examples:
+            >>> from edsl import Agent
+            >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}),
+            ...                Agent(traits = {'a': 1, 'b': 2})])
+            >>> al2 = al.rename('a', 'c')
+            >>> assert al2 == AgentList([Agent(traits = {'c': 1, 'b': 1}),
+            ...                         Agent(traits = {'c': 1, 'b': 2})])
+            >>> assert al != al2
         """
         newagents = []
         for agent in self:
@@ -133,15 +164,21 @@ class AgentList(UserList, Base):
         return AgentList(newagents)
 
     def select(self, *traits) -> AgentList:
-        """Selects agents with only the references traits.
+        """Create a new AgentList with only the specified traits.
 
-        >>> from edsl.agents.Agent import Agent
-        >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}), Agent(traits = {'a': 1, 'b': 2})])
-        >>> al.select('a')
-        AgentList([Agent(traits = {'a': 1}), Agent(traits = {'a': 1})])
+        Args:
+            *traits: Variable number of trait names to keep.
 
-        """
+        Returns:
+            AgentList: A new AgentList containing agents with only the selected traits.
 
+        Examples:
+            >>> from edsl import Agent
+            >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}),
+            ...                Agent(traits = {'a': 1, 'b': 2})])
+            >>> al.select('a')
+            AgentList([Agent(traits = {'a': 1}), Agent(traits = {'a': 1})])
+        """
         if len(traits) == 1:
             traits_to_select = [list(traits)[0]]
         else:
@@ -150,23 +187,27 @@ class AgentList(UserList, Base):
         return AgentList([agent.select(*traits_to_select) for agent in self.data])
 
     def filter(self, expression: str) -> AgentList:
-        """
-        Filter a list of agents based on an expression.
+        """Filter agents based on a boolean expression.
+
+        Args:
+            expression: A string containing a boolean expression to evaluate against
+                each agent's traits.
 
-        >>> from edsl.agents.Agent import Agent
-        >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}), Agent(traits = {'a': 1, 'b': 2})])
-        >>> al.filter("b == 2")
-        AgentList([Agent(traits = {'a': 1, 'b': 2})])
+        Returns:
+            AgentList: A new AgentList containing only agents that satisfy the expression.
+
+        Examples:
+            >>> from edsl import Agent
+            >>> al = AgentList([Agent(traits = {'a': 1, 'b': 1}),
+            ...                Agent(traits = {'a': 1, 'b': 2})])
+            >>> al.filter("b == 2")
+            AgentList([Agent(traits = {'a': 1, 'b': 2})])
         """
 
         def create_evaluator(agent: "Agent"):
-            """Create an evaluator for the given result.
-            The 'combined_dict' is a mapping of all values for that Result object.
-            """
+            """Create an evaluator for the given agent."""
             return EvalWithCompoundTypes(names=agent.traits)
 
-            # iterates through all the results and evaluates the expression
-
         try:
             new_data = [
                 agent for agent in self.data if create_evaluator(agent).eval(expression)
@@ -188,7 +229,7 @@ class AgentList(UserList, Base):
     @property
     def all_traits(self) -> list[str]:
         """Return all traits in the AgentList.
-        >>> from edsl.agents.Agent import Agent
+        >>> from edsl import Agent
         >>> agent_1 = Agent(traits = {'age': 22})
         >>> agent_2 = Agent(traits = {'hair': 'brown'})
         >>> al = AgentList([agent_1, agent_2])
@@ -201,7 +242,7 @@ class AgentList(UserList, Base):
         return list(d.keys())
 
     @classmethod
-    def from_csv(cls, file_path: str, name_field: Optional[str] = None):
+    def from_csv(cls, file_path: str, name_field: Optional[str] = None, codebook: Optional[dict[str, str]] = None):
         """Load AgentList from a CSV file.
 
         >>> import csv
@@ -216,12 +257,16 @@ class AgentList(UserList, Base):
         >>> al = AgentList.from_csv('/tmp/agents.csv', name_field='hair')
         >>> al
         AgentList([Agent(name = \"""brown\""", traits = {'age': '22', 'height': '5.5'})])
+        >>> al = AgentList.from_csv('/tmp/agents.csv', codebook={'age': 'Age in years'})
+        >>> al[0].codebook
+        {'age': 'Age in years'}
         >>> os.remove('/tmp/agents.csv')
 
         :param file_path: The path to the CSV file.
         :param name_field: The name of the field to use as the agent name.
+        :param codebook: Optional dictionary mapping trait names to descriptions.
         """
-        from edsl.agents.Agent import Agent
+        from .agent import Agent
 
         agent_list = []
         with open(file_path, "r") as f:
@@ -234,9 +279,9 @@ class AgentList(UserList, Base):
                     name_field = "name"
                 if name_field is not None:
                     agent_name = row.pop(name_field)
-                    agent_list.append(Agent(traits=row, name=agent_name))
+                    agent_list.append(Agent(traits=row, name=agent_name, codebook=codebook))
                 else:
-                    agent_list.append(Agent(traits=row))
+                    agent_list.append(Agent(traits=row, codebook=codebook))
         return cls(agent_list)
 
     def translate_traits(self, codebook: dict[str, str]):
@@ -258,7 +303,7 @@ class AgentList(UserList, Base):
         """Remove traits from the AgentList.
 
         :param traits: The traits to remove.
-        >>> from edsl.agents.Agent import Agent
+        >>> from edsl import Agent
         >>> al = AgentList([Agent({'age': 22, 'hair': 'brown', 'height': 5.5}), Agent({'age': 22, 'hair': 'brown', 'height': 5.5})])
         >>> al.remove_trait('age')
         AgentList([Agent(traits = {'hair': 'brown', 'height': 5.5}), Agent(traits = {'hair': 'brown', 'height': 5.5})])
@@ -282,7 +327,7 @@ class AgentList(UserList, Base):
         >>> al.add_trait('new_trait', [1, 2, 3])
         Traceback (most recent call last):
         ...
-        edsl.exceptions.agents.AgentListError: The passed values have to be the same length as the agent list.
+        edsl.agents.exceptions.AgentListError: The passed values have to be the same length as the agent list.
         ...
         """
         if not is_iterable(values):
@@ -322,8 +367,6 @@ class AgentList(UserList, Base):
         >>> hash(al)
         1681154913465662422
         """
-        from edsl.utilities.utilities import dict_hash
-
         return dict_hash(self.to_dict(add_edsl_version=False, sorted=True))
 
     def to_dict(self, sorted=False, add_edsl_version=True):
@@ -331,7 +374,13 @@ class AgentList(UserList, Base):
 
         >>> AgentList.example().to_dict(add_edsl_version=False)
         {'agent_list': [{'traits': {'age': 22, 'hair': 'brown', 'height': 5.5}}, {'traits': {'age': 22, 'hair': 'brown', 'height': 5.5}}]}
-
+        >>> example_codebook = {'age': 'Age in years'}
+        >>> al = AgentList.example().set_codebook(example_codebook)
+        >>> result = al.to_dict(add_edsl_version=False)
+        >>> 'codebook' in result
+        True
+        >>> result['codebook'] == example_codebook
+        True
         """
         if sorted:
             data = self.data[:]
@@ -344,6 +393,19 @@ class AgentList(UserList, Base):
                 agent.to_dict(add_edsl_version=add_edsl_version) for agent in data
             ]
         }
+        
+        # Add codebook if all agents have the same codebook
+        if len(self.data) > 0:
+            # Get the first agent's codebook
+            first_codebook = self.data[0].codebook
+            
+            # Check if all agents have the same codebook
+            all_same = all(agent.codebook == first_codebook for agent in self.data)
+            
+            # Only include codebook if it's non-empty and consistent across all agents
+            if all_same and first_codebook:
+                d["codebook"] = first_codebook
+                
         if add_edsl_version:
             from edsl import __version__
 
@@ -360,15 +422,15 @@ class AgentList(UserList, Base):
     def __repr__(self):
         return f"AgentList({self.data})"
 
-    def _summary(self):
+    def _summary(self) -> dict:
         return {
             "agents": len(self),
         }
-
+    
     def set_codebook(self, codebook: dict[str, str]) -> AgentList:
         """Set the codebook for the AgentList.
 
-        >>> from edsl.agents.Agent import Agent
+        >>> from edsl import Agent
         >>> a = Agent(traits = {'hair': 'brown'})
         >>> al = AgentList([a, a])
         >>> _ = al.set_codebook({'hair': "Color of hair on driver's license"})
@@ -383,51 +445,13 @@ class AgentList(UserList, Base):
 
         return self
 
-    def to_csv(self, file_path: str):
-        """Save the AgentList to a CSV file.
-
-        :param file_path: The path to the CSV file.
-        """
-        self.to_scenario_list().to_csv(file_path)
-
-    def to_list(self, include_agent_name=False) -> list[tuple]:
-        """Return a list of tuples."""
-        return self.to_scenario_list(include_agent_name).to_list()
-
-    def to_scenario_list(
-        self, include_agent_name: bool = False, include_instruction: bool = False
-    ) -> ScenarioList:
-        """Converts the agent to a scenario list."""
-        from edsl.scenarios.ScenarioList import ScenarioList
-        from edsl.scenarios.Scenario import Scenario
-
-        # raise NotImplementedError("This method is not implemented yet.")
-
-        scenario_list = ScenarioList()
-        for agent in self.data:
-            d = agent.traits
-            if include_agent_name:
-                d["agent_name"] = agent.name
-            if include_instruction:
-                d["instruction"] = agent.instruction
-            scenario_list.append(Scenario(d))
-        return scenario_list
-
-        # if include_agent_name:
-        #     return ScenarioList(
-        #         [
-        #             Scenario(agent.traits | {"agent_name": agent.name} | })
-        #             for agent in self.data
-        #         ]
-        #     )
-        # return ScenarioList([Scenario(agent.traits) for agent in self.data])
 
     def table(
         self,
         *fields,
         tablefmt: Optional[str] = None,
         pretty_labels: Optional[dict] = None,
-    ) -> Table:
+    ) -> "Table":
         if len(self) == 0:
             e = AgentListError("Cannot create a table from an empty AgentList.")
             if is_notebook():
@@ -442,18 +466,24 @@ class AgentList(UserList, Base):
         )
 
     def to_dataset(self, traits_only: bool = True):
+        """Convert the AgentList to a Dataset.
+
+        Args:
+            traits_only: If True, only include agent traits. If False, also include
+                agent parameters like instructions and names.
+
+        Returns:
+            Dataset: A dataset containing the agents' traits and optionally their parameters.
+
+        Examples:
+            >>> from edsl import AgentList
+            >>> al = AgentList.example()
+            >>> al.to_dataset()
+            Dataset([{'age': [22, 22]}, {'hair': ['brown', 'brown']}, {'height': [5.5, 5.5]}])
+            >>> al.to_dataset(traits_only=False)  # doctest: +NORMALIZE_WHITESPACE
+            Dataset([{'age': [22, 22]}, {'hair': ['brown', 'brown']}, {'height': [5.5, 5.5]}, {'agent_parameters': [{'instruction': 'You are answering questions as if you were a human. Do not break character.', 'agent_name': None}, {'instruction': 'You are answering questions as if you were a human. Do not break character.', 'agent_name': None}]}])
         """
-        Convert the AgentList to a Dataset.
-
-        >>> from edsl.agents.AgentList import AgentList
-        >>> al = AgentList.example()
-        >>> al.to_dataset()
-        Dataset([{'age': [22, 22]}, {'hair': ['brown', 'brown']}, {'height': [5.5, 5.5]}])
-        >>> al.to_dataset(traits_only = False)
-        Dataset([{'age': [22, 22]}, {'hair': ['brown', 'brown']}, {'height': [5.5, 5.5]}, {'agent_parameters': [{'instruction': 'You are answering questions as if you were a human. Do not break character.', 'agent_name': None}, {'instruction': 'You are answering questions as if you were a human. Do not break character.', 'agent_name': None}]}])
-        """
-        from edsl.results.Dataset import Dataset
-        from collections import defaultdict
+        from ..dataset import Dataset
 
         agent_trait_keys = []
         for agent in self:
@@ -472,55 +502,84 @@ class AgentList(UserList, Base):
                 )
         return Dataset([{key: entry} for key, entry in data.items()])
 
-    def tree(self, node_order: Optional[List[str]] = None):
-        return self.to_scenario_list().tree(node_order)
-
     @classmethod
     @remove_edsl_version
     def from_dict(cls, data: dict) -> "AgentList":
         """Deserialize the dictionary back to an AgentList object.
 
         :param: data: A dictionary representing an AgentList.
-        >>> from edsl.agents.Agent import Agent
+
+        >>> from edsl import Agent
         >>> al = AgentList([Agent.example(), Agent.example()])
         >>> al2 = AgentList.from_dict(al.to_dict())
         >>> al2 == al
         True
+        >>> example_codebook = {'age': 'Age in years'}
+        >>> al = AgentList([Agent.example()]).set_codebook(example_codebook)
+        >>> al2 = AgentList.from_dict(al.to_dict())
+        >>> al2[0].codebook == example_codebook
+        True
         """
-        from edsl.agents.Agent import Agent
+        from .agent import Agent
 
         agents = [Agent.from_dict(agent_dict) for agent_dict in data["agent_list"]]
-        return cls(agents)
+        agent_list = cls(agents)
+        
+        # Apply codebook if present in the dictionary
+        if "codebook" in data and data["codebook"]:
+            agent_list.set_codebook(data["codebook"])
+            
+        return agent_list
 
     @classmethod
-    def example(cls, randomize: bool = False) -> AgentList:
+    def example(cls, randomize: bool = False, codebook: Optional[dict[str, str]] = None) -> AgentList:
         """
         Returns an example AgentList instance.
 
         :param randomize: If True, uses Agent's randomize method.
+        :param codebook: Optional dictionary mapping trait names to descriptions.
+        
+        >>> al = AgentList.example()
+        >>> al
+        AgentList([Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5}), Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})])
+        >>> al = AgentList.example(codebook={'age': 'Age in years'})
+        >>> al[0].codebook
+        {'age': 'Age in years'}
         """
-        from edsl.agents.Agent import Agent
+        from .agent import Agent
 
-        return cls([Agent.example(randomize), Agent.example(randomize)])
+        agent_list = cls([Agent.example(randomize), Agent.example(randomize)])
+        
+        if codebook:
+            agent_list.set_codebook(codebook)
+            
+        return agent_list
 
     @classmethod
-    def from_list(self, trait_name: str, values: List[Any]):
+    def from_list(self, trait_name: str, values: List[Any], codebook: Optional[dict[str, str]] = None) -> "AgentList":
         """Create an AgentList from a list of values.
 
         :param trait_name: The name of the trait.
         :param values: A list of values.
+        :param codebook: Optional dictionary mapping trait names to descriptions.
 
         >>> AgentList.from_list('age', [22, 23])
         AgentList([Agent(traits = {'age': 22}), Agent(traits = {'age': 23})])
+        >>> al = AgentList.from_list('age', [22], codebook={'age': 'Age in years'})
+        >>> al[0].codebook
+        {'age': 'Age in years'}
         """
-        from edsl.agents.Agent import Agent
+        from .agent import Agent
 
-        return AgentList([Agent({trait_name: value}) for value in values])
+        agent_list = AgentList([Agent({trait_name: value}) for value in values])
+        
+        if codebook:
+            agent_list.set_codebook(codebook)
+            
+        return agent_list
 
     def __mul__(self, other: AgentList) -> AgentList:
         """Takes the cross product of two AgentLists."""
-        from itertools import product
-
         new_sl = []
         for s1, s2 in list(product(self, other)):
             new_sl.append(s1 + s2)
@@ -531,13 +590,13 @@ class AgentList(UserList, Base):
 
         >>> al = AgentList.example()
         >>> print(al.code())
-        from edsl.agents.Agent import Agent
-        from edsl.agents.AgentList import AgentList
+        from edsl import Agent
+        from edsl import AgentList
         agent_list = AgentList([Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5}), Agent(traits = {'age': 22, 'hair': 'brown', 'height': 5.5})])
         """
         lines = [
-            "from edsl.agents.Agent import Agent",
-            "from edsl.agents.AgentList import AgentList",
+            "from edsl import Agent",
+            "from edsl import AgentList",
         ]
         lines.append(f"agent_list = AgentList({self.data})")
         if string:
@@ -547,5 +606,4 @@ class AgentList(UserList, Base):
 
 if __name__ == "__main__":
     import doctest
-
-    doctest.testmod(optionflags=doctest.ELLIPSIS)
+    doctest.testmod(optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE)