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

edsl/{results/Dataset.py → dataset/dataset.py}

@@ -1,46 +1,84 @@
-"""A module to represent a dataset of observations."""
+
 
 from __future__ import annotations
 import sys
 import json
 import random
 from collections import UserList
-from typing import Any, Union, Optional
-
-from edsl.results.ResultsExportMixin import ResultsExportMixin
-from edsl.results.DatasetTree import Tree
-from edsl.results.TableDisplay import TableDisplay
-from edsl.Base import PersistenceMixin, HashingMixin
-
-
-from edsl.results.smart_objects import FirstObject
-
-from edsl.results.ResultsGGMixin import GGPlotMethod
-
-class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
-    """A class to represent a dataset of observations."""
+from typing import Any, Union, Optional, TYPE_CHECKING
+
+from ..base import PersistenceMixin, HashingMixin
+
+from .dataset_tree import Tree
+
+from .display.table_display import TableDisplay
+from .smart_objects import FirstObject
+from .r.ggplot import GGPlotMethod
+from .dataset_operations_mixin import DatasetOperationsMixin
+
+if TYPE_CHECKING:
+    from ..surveys import Survey
+    from ..questions.QuestionBase import QuestionBase
+
+class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
+    """
+    A versatile data container for tabular data with powerful manipulation capabilities.
+    
+    The Dataset class is a fundamental data structure in EDSL that represents tabular data
+    in a column-oriented format. It provides a rich set of methods for data manipulation,
+    transformation, analysis, visualization, and export through the DatasetOperationsMixin.
+    
+    Key features:
+    
+    1. Column-oriented data structure optimized for LLM experiment results
+    2. Rich data manipulation API similar to dplyr/pandas (filter, select, mutate, etc.)
+    3. Visualization capabilities including tables, plots, and reports
+    4. Export to various formats (CSV, Excel, SQLite, pandas, etc.)
+    5. Serialization for storage and transport
+    6. Tree-based data exploration
+    
+    A Dataset typically contains multiple columns, each represented as a dictionary
+    with a single key-value pair. The key is the column name and the value is a list
+    of values for that column. All columns must have the same length.
+    
+    The Dataset class inherits from:
+    - UserList: Provides list-like behavior for storing column data
+    - DatasetOperationsMixin: Provides data manipulation methods
+    - PersistenceMixin: Provides serialization capabilities
+    - HashingMixin: Provides hashing functionality for comparison and storage
+    
+    Datasets are typically created by transforming other EDSL container types like
+    Results, AgentList, or ScenarioList, but can also be created directly from data.
+    """
 
     def __init__(
         self, data: list[dict[str, Any]] = None, print_parameters: Optional[dict] = None
     ):
-        """Initialize the dataset with the given data."""
+        """
+        Initialize a new Dataset instance.
+        
+        Parameters:
+            data: A list of dictionaries, where each dictionary represents a column
+                 in the dataset. Each dictionary should have a single key-value pair,
+                 where the key is the column name and the value is a list of values.
+                 All value lists must have the same length.
+            print_parameters: Optional dictionary of parameters controlling how the
+                             dataset is displayed when printed.
+                
+        Examples:
+            >>> # Create a dataset with two columns
+            >>> d = Dataset([{'a': [1, 2, 3]}, {'b': [4, 5, 6]}])
+            >>> len(d)
+            3
+            
+            >>> # Dataset with a single column
+            >>> Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}])
+            Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}])
+        """
         super().__init__(data)
         self.print_parameters = print_parameters
 
 
-    def ggplot2(
-        self,
-        ggplot_code: str,
-        shape="wide",
-        sql: str = None,
-        remove_prefix: bool = True,
-        debug: bool = False,
-        height=4,
-        width=6,
-        factor_orders: Optional[dict] = None,
-    ):
-        return GGPlotMethod(self).ggplot2(ggplot_code, shape, sql, remove_prefix, debug, height, width, factor_orders)
-
     def __len__(self) -> int:
         """Return the number of observations in the dataset.
 
@@ -95,16 +133,29 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
         return w
 
     def keys(self) -> list[str]:
-        """Return the keys of the first observation in the dataset.
+        """Return the keys of the dataset.
 
         >>> d = Dataset([{'a.b':[1,2,3,4]}])
         >>> d.keys()
+        ['a.b']
+
+        >>> d = Dataset([{'a.b':[1,2,3,4]}, {'c.d':[5,6,7,8]}])
+        >>> d.keys()
+        ['a.b', 'c.d']
+
+
         ['a.b']
         """
         return [list(o.keys())[0] for o in self]
 
     def filter(self, expression):
         return self.to_scenario_list().filter(expression).to_dataset()
+    
+    def mutate(self, new_var_string: str, functions_dict: Optional[dict[str, Callable]] = None) -> "Dataset":
+        return self.to_scenario_list().mutate(new_var_string, functions_dict).to_dataset()
+    
+    def collapse(self, field:str, separator: Optional[str] = None) -> "Dataset":
+        return self.to_scenario_list().collapse(field, separator).to_dataset()
 
     def long(self, exclude_fields: list[str] = None) -> Dataset:
         headers, data = self._tabular()
@@ -274,14 +325,33 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
         return Dataset(new_data)
 
     def print(self, pretty_labels=None, **kwargs):
+        """
+        Print the dataset in a formatted way.
+        
+        Args:
+            pretty_labels: A dictionary mapping column names to their display names
+            **kwargs: Additional arguments
+                format: The output format ("html", "markdown", "rich", "latex")
+                
+        Returns:
+            TableDisplay object
+        """
         if "format" in kwargs:
             if kwargs["format"] not in ["html", "markdown", "rich", "latex"]:
                 raise ValueError(f"Format '{kwargs['format']}' not supported.")
+            
+            # If rich format is requested, set tablefmt accordingly
+            if kwargs["format"] == "rich":
+                kwargs["tablefmt"] = "rich"
+                
         if pretty_labels is None:
             pretty_labels = {}
         else:
             return self.rename(pretty_labels).print(**kwargs)
-        return self.table()
+            
+        # Pass through any tablefmt parameter
+        tablefmt = kwargs.get("tablefmt", None)
+        return self.table(tablefmt=tablefmt)
 
     def rename(self, rename_dic) -> Dataset:
         new_data = []
@@ -302,7 +372,8 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
         return Dataset.from_pandas_dataframe(merged_df)
 
     def to(self, survey_or_question: Union["Survey", "QuestionBase"]) -> "Jobs":
-        from edsl.surveys.Survey import Survey
+        """Return a new dataset with the observations transformed by the given survey or question."""
+        from edsl.surveys import Survey
         from edsl.questions.QuestionBase import QuestionBase
 
         if isinstance(survey_or_question, Survey):
@@ -321,7 +392,14 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
 
         >>> d.select('a.b', 'c.d')
         Dataset([{'a.b': [1, 2, 3, 4]}, {'c.d': [5, 6, 7, 8]}])
+
         """
+        for key in keys:
+            if key not in self.keys():
+                raise ValueError(f"Key '{key}' not found in the dataset."
+                                 f"Available keys: {self.keys()}"
+                                 )
+            
         if isinstance(keys, str):
             keys = [keys]
 
@@ -491,8 +569,10 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
 
         >>> d = Dataset([{'a':[1,2,3,4]}, {'b':[4,3,2,1]}])
         >>> d.tree()
-        Tree(Dataset({'a': [1, 2, 3, 4], 'b': [4, 3, 2, 1]}))
+        Tree(Dataset({'a': [1, 2, 3, 4], 'b': [4, 3, 2, 1]}), node_order=['a', 'b'])
         """
+        if node_order is None:
+            node_order = self.keys()
         return Tree(self, node_order=node_order)
 
     def table(
@@ -515,7 +595,8 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
 
         headers, data = self._tabular()
 
-        if tablefmt is not None:
+        if tablefmt is not None and tablefmt != "rich":
+            # Rich format is handled separately, so we don't validate it against tabulate_formats
             from tabulate import tabulate_formats
 
             if tablefmt not in tabulate_formats:
@@ -523,7 +604,7 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
                     f"Error: The following table format is not supported: {tablefmt}",
                     file=sys.stderr,
                 )
-                print(f"\nAvailable formats are: {tabulate_formats}", file=sys.stderr)
+                print(f"\nAvailable formats are: {tabulate_formats} and 'rich'", file=sys.stderr)
                 return None
 
         if max_rows:
@@ -648,5 +729,4 @@ class Dataset(UserList, ResultsExportMixin, PersistenceMixin, HashingMixin):
 
 if __name__ == "__main__":
     import doctest
-
     doctest.testmod(optionflags=doctest.ELLIPSIS)