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

edsl/dataset/dataset.py

@@ -1,24 +1,25 @@
-
-
 from __future__ import annotations
 import sys
 import json
 import random
 from collections import UserList
-from typing import Any, Union, Optional, TYPE_CHECKING
+from typing import Any, Union, Optional, TYPE_CHECKING, Callable
 
 from ..base import PersistenceMixin, HashingMixin
 
 from .dataset_tree import Tree
+from .exceptions import DatasetKeyError, DatasetValueError, DatasetTypeError
+
 
 from .display.table_display import TableDisplay
-from .smart_objects import FirstObject
-from .r.ggplot import GGPlotMethod
+#from .smart_objects import FirstObject
 from .dataset_operations_mixin import DatasetOperationsMixin
 
 if TYPE_CHECKING:
     from ..surveys import Survey
-    from ..questions.QuestionBase import QuestionBase
+    from ..questions import QuestionBase
+    from ..jobs import Job  # noqa: F401
+
 
 class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
     """
@@ -76,6 +77,7 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
             Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}])
         """
         super().__init__(data)
+        #self.data = data
         self.print_parameters = print_parameters
 
 
@@ -118,19 +120,9 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
             new_data.append({key: values[:n]})
         return Dataset(new_data)
 
-    def expand(self, field):
-        return self.to_scenario_list().expand(field)
+    # def expand(self, field):
+    #     return self.to_scenario_list().expand(field)
 
-    def view(self):
-        from perspective.widget import PerspectiveWidget
-
-        w = PerspectiveWidget(
-            self.to_pandas(),
-            plugin="Datagrid",
-            aggregates={"datetime": "any"},
-            sort=[["date", "desc"]],
-        )
-        return w
 
     def keys(self) -> list[str]:
         """Return the keys of the dataset.
@@ -212,7 +204,7 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         values = value_dict["value"]
 
         if not (len(rows) == len(keys) == len(values)):
-            raise ValueError("All input arrays must have the same length")
+            raise DatasetValueError("All input arrays must have the same length")
 
         # Get unique keys and row indices
         unique_keys = sorted(set(keys))
@@ -272,12 +264,6 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         >>> d = Dataset([{'a.b':[1,2,3,4]}])
         >>> d._key_to_value('a.b')
         [1, 2, 3, 4]
-
-        >>> d._key_to_value('a')
-        Traceback (most recent call last):
-        ...
-        KeyError: "Key 'a' not found in any of the dictionaries."
-
         """
         potential_matches = []
         for data_dict in self.data:
@@ -290,11 +276,13 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         if len(potential_matches) == 1:
             return potential_matches[0][1]
         elif len(potential_matches) > 1:
-            raise KeyError(
+            from .exceptions import DatasetKeyError
+            raise DatasetKeyError(
                 f"Key '{key}' found in more than one location: {[m[0] for m in potential_matches]}"
             )
 
-        raise KeyError(f"Key '{key}' not found in any of the dictionaries.")
+        from .exceptions import DatasetKeyError
+        raise DatasetKeyError(f"Key '{key}' not found in any of the dictionaries.")
 
     def first(self) -> dict[str, Any]:
         """Get the first value of the first key in the first dictionary.
@@ -308,7 +296,7 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
             """Get the values of the first key in the dictionary."""
             return list(d.values())[0]
 
-        return FirstObject(get_values(self.data[0])[0])
+        return get_values(self.data[0])[0]
 
     def latex(self, **kwargs):
         return self.table().latex()
@@ -338,7 +326,7 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         """
         if "format" in kwargs:
             if kwargs["format"] not in ["html", "markdown", "rich", "latex"]:
-                raise ValueError(f"Format '{kwargs['format']}' not supported.")
+                raise DatasetValueError(f"Format '{kwargs['format']}' not supported.")
             
             # If rich format is requested, set tablefmt accordingly
             if kwargs["format"] == "rich":
@@ -371,10 +359,18 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         merged_df = df1.merge(df2, how="left", left_on=by_x, right_on=by_y)
         return Dataset.from_pandas_dataframe(merged_df)
 
-    def to(self, survey_or_question: Union["Survey", "QuestionBase"]) -> "Jobs":
-        """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
+    def to(self, survey_or_question: Union["Survey", "QuestionBase"]) -> "Job":
+        """Return a new dataset with the observations transformed by the given survey or question.
+        
+        >>> d = Dataset([{'person_name':["John"]}])
+        >>> from edsl import QuestionFreeText 
+        >>> q = QuestionFreeText(question_text = "How are you, {{ person_name ?}}?", question_name = "how_feeling")
+        >>> jobs = d.to(q)
+        >>> isinstance(jobs, object)
+        True
+        """
+        from ..surveys import Survey
+        from ..questions import QuestionBase
 
         if isinstance(survey_or_question, Survey):
             return survey_or_question.by(self.to_scenario_list())
@@ -396,9 +392,10 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         """
         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()}"
-                                 )
+                from .exceptions import DatasetValueError
+                raise DatasetValueError(f"Key '{key}' not found in the dataset. "
+                                        f"Available keys: {self.keys()}"
+                                       )
             
         if isinstance(keys, str):
             keys = [keys]
@@ -442,7 +439,11 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
 
         return self
 
-    def expand(self, field):
+    def expand_field(self, field):
+        """Expand a field in the dataset.
+        
+        Renamed to avoid conflict with the expand method defined earlier.
+        """
         return self.to_scenario_list().expand(field).to_dataset()
 
     def sample(
@@ -462,21 +463,18 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         >>> d = Dataset([{'a.b':[1,2,3,4]}])
         >>> d.sample(n=2, seed=0, with_replacement=True)
         Dataset([{'a.b': [4, 4]}])
-
-        >>> d.sample(n = 10, seed=0, with_replacement=False)
-        Traceback (most recent call last):
-        ...
-        ValueError: Sample size cannot be greater than the number of available elements when sampling without replacement.
         """
         if seed is not None:
             random.seed(seed)
 
         # Validate the input for sampling parameters
         if n is None and frac is None:
-            raise ValueError("Either 'n' or 'frac' must be provided for sampling.")
+            from .exceptions import DatasetValueError
+            raise DatasetValueError("Either 'n' or 'frac' must be provided for sampling.")
 
         if n is not None and frac is not None:
-            raise ValueError("Only one of 'n' or 'frac' should be specified.")
+            from .exceptions import DatasetValueError
+            raise DatasetValueError("Only one of 'n' or 'frac' should be specified.")
 
         # Get the length of the lists from the first entry
         first_key, first_values = list(self[0].items())[0]
@@ -487,7 +485,8 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
             n = int(total_length * frac)
 
         if not with_replacement and n > total_length:
-            raise ValueError(
+            from .exceptions import DatasetValueError
+            raise DatasetValueError(
                 "Sample size cannot be greater than the number of available elements when sampling without replacement."
             )
 
@@ -504,60 +503,72 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
 
         return self
 
-    def order_by(self, sort_key: str, reverse: bool = False) -> Dataset:
-        """Return a new dataset with the observations sorted by the given key.
-
-        :param sort_key: The key to sort the observations by.
-        :param reverse: Whether to sort in reverse order.
-
-        >>> d = Dataset([{'a':[1,2,3,4]}, {'b':[4,3,2,1]}])
-        >>> d.order_by('a')
-        Dataset([{'a': [1, 2, 3, 4]}, {'b': [4, 3, 2, 1]}])
-
-        >>> d.order_by('a', reverse=True)
-        Dataset([{'a': [4, 3, 2, 1]}, {'b': [1, 2, 3, 4]}])
-
-        >>> d = Dataset([{'X.a':[1,2,3,4]}, {'X.b':[4,3,2,1]}])
-        >>> d.order_by('a')
-        Dataset([{'X.a': [1, 2, 3, 4]}, {'X.b': [4, 3, 2, 1]}])
-
-
+    def get_sort_indices(self, lst: list[Any], reverse: bool = False, use_numpy: bool = True) -> list[int]:
         """
-        import numpy as np
+        Return the indices that would sort the list, using either numpy or pure Python.
+        None values are placed at the end of the sorted list.
 
-        def sort_indices(lst: list[Any]) -> list[int]:
-            """
-            Return the indices that would sort the list.
+        Args:
+            lst: The list to be sorted
+            reverse: Whether to sort in descending order
+            use_numpy: Whether to use numpy implementation (falls back to pure Python if numpy is unavailable)
 
-            :param lst: The list to be sorted.
-            :return: A list of indices that would sort the list.
-            """
-            indices = np.argsort(lst).tolist()
-            if reverse:
-                indices.reverse()
-            return indices
+        Returns:
+            A list of indices that would sort the list
+        """
+        if use_numpy:
+            try:
+                import numpy as np
+                # Convert list to numpy array
+                arr = np.array(lst, dtype=object)
+                # Get mask of non-None values
+                mask = ~(arr is None)
+                # Get indices of non-None and None values
+                non_none_indices = np.where(mask)[0]
+                none_indices = np.where(~mask)[0]
+                # Sort non-None values
+                sorted_indices = non_none_indices[np.argsort(arr[mask])]
+                # Combine sorted non-None indices with None indices
+                indices = np.concatenate([sorted_indices, none_indices]).tolist()
+                if reverse:
+                    # When reversing, keep None values at end
+                    indices = sorted_indices[::-1].tolist() + none_indices.tolist()
+                return indices
+            except ImportError:
+                # Fallback to pure Python if numpy is not available
+                pass
+        
+        # Pure Python implementation
+        enumerated = list(enumerate(lst))
+        # Sort None values to end by using (is_none, value) as sort key
+        sorted_pairs = sorted(enumerated, 
+                            key=lambda x: (x[1] is None, x[1]), 
+                            reverse=reverse)
+        return [index for index, _ in sorted_pairs]
+
+    def order_by(self, sort_key: str, reverse: bool = False, use_numpy: bool = True) -> Dataset:
+        """Return a new dataset with the observations sorted by the given key.
 
+        Args:
+            sort_key: The key to sort the observations by
+            reverse: Whether to sort in reverse order
+            use_numpy: Whether to use numpy for sorting (faster for large lists)
+        """
         number_found = 0
         for obs in self.data:
             key, values = list(obs.items())[0]
-            # an obseration is {'a':[1,2,3,4]}
-            # key = list(obs.keys())[0]
-            if (
-                sort_key == key or sort_key == key.split(".")[-1]
-            ):  # e.g., "age" in "scenario.age"
+            if sort_key == key or sort_key == key.split(".")[-1]:
                 relevant_values = values
                 number_found += 1
 
         if number_found == 0:
-            raise ValueError(f"Key '{sort_key}' not found in any of the dictionaries.")
+            raise DatasetKeyError(f"Key '{sort_key}' not found in any of the dictionaries.")
         elif number_found > 1:
-            raise ValueError(f"Key '{sort_key}' found in more than one dictionary.")
+            raise DatasetKeyError(f"Key '{sort_key}' found in more than one dictionary.")
 
-        # relevant_values = self._key_to_value(sort_key)
-        sort_indices_list = sort_indices(relevant_values)
+        sort_indices_list = self.get_sort_indices(relevant_values, reverse=reverse, use_numpy=use_numpy)
         new_data = []
         for observation in self.data:
-            # print(observation)
             key, values = list(observation.items())[0]
             new_values = [values[i] for i in sort_indices_list]
             new_data.append({key: new_values})
@@ -578,7 +589,7 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
     def table(
         self,
         *fields,
-        tablefmt: Optional[str] = None,
+        tablefmt: Optional[str] = "rich",
         max_rows: Optional[int] = None,
         pretty_labels=None,
         print_parameters: Optional[dict] = None,
@@ -637,7 +648,8 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
 
         if max_rows is not None:
             if max_rows > len(data):
-                raise ValueError(
+                from .exceptions import DatasetValueError
+                raise DatasetValueError(
                     "max_rows cannot be greater than the number of rows in the dataset."
                 )
             last_line = data[-1]
@@ -675,6 +687,19 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
     def from_pandas_dataframe(cls, df):
         result = cls([{col: df[col].tolist()} for col in df.columns])
         return result
+    
+    def to_dict(self) -> dict:
+        """
+        Convert the dataset to a dictionary.
+        """
+        return {'data': self.data}
+
+    @classmethod
+    def from_dict(cls, data: dict) -> 'Dataset':
+        """
+        Convert a dictionary to a dataset.
+        """
+        return cls(data['data'])
 
     def to_docx(self, output_file: str, title: str = None) -> None:
         """
@@ -726,6 +751,72 @@ class Dataset(UserList, DatasetOperationsMixin, PersistenceMixin, HashingMixin):
         # Save the document
         doc.save(output_file)
 
+    def expand(self, field: str, number_field: bool = False) -> "Dataset":
+        """
+        Expand a field containing lists into multiple rows.
+        
+        Args:
+            field: The field containing lists to expand
+            number_field: If True, adds a number field indicating the position in the original list
+            
+        Returns:
+            A new Dataset with the expanded rows
+            
+        Example:
+            >>> from edsl.dataset import Dataset
+            >>> d = Dataset([{'a': [[1, 2, 3], [4, 5, 6]]}, {'b': ['x', 'y']}])
+            >>> d.expand('a')
+            Dataset([{'a': [1, 2, 3, 4, 5, 6]}, {'b': ['x', 'x', 'x', 'y', 'y', 'y']}])
+        """
+        from collections.abc import Iterable
+        
+        # Find the field in the dataset
+        field_data = None
+        for entry in self.data:
+            key = list(entry.keys())[0]
+            if key == field:
+                field_data = entry[key]
+                break
+            
+        if field_data is None:
+            raise DatasetKeyError(f"Field '{field}' not found in dataset. Available fields are: {self.keys()}")
+
+
+        # Validate that the field contains lists
+        if not all(isinstance(v, list) for v in field_data):
+            raise DatasetTypeError(f"Field '{field}' must contain lists in all entries")
+        
+        # Create new expanded data structure
+        new_data = []
+        
+        # Process each field
+        for entry in self.data:
+            key, values = list(entry.items())[0]
+            new_values = []
+            
+            if key == field:
+                # This is the field to expand - flatten all sublists
+                for row_values in values:
+                    if not isinstance(row_values, Iterable) or isinstance(row_values, str):
+                        row_values = [row_values]
+                    new_values.extend(row_values)
+            else:
+                # For other fields, repeat each value the appropriate number of times
+                for i, row_value in enumerate(values):
+                    expand_length = len(field_data[i]) if i < len(field_data) else 0
+                    new_values.extend([row_value] * expand_length)
+            
+            new_data.append({key: new_values})
+        
+        # Add number field if requested
+        if number_field:
+            number_values = []
+            for i, lst in enumerate(field_data):
+                number_values.extend(range(1, len(lst) + 1))
+            new_data.append({f"{field}_number": number_values})
+        
+        return Dataset(new_data)
+
 
 if __name__ == "__main__":
     import doctest