edsl 0.1.30.dev4__py3-none-any.whl → 0.1.31__py3-none-any.whl

edsl/results/DatasetExportMixin.py

@@ -3,18 +3,22 @@
 import base64
 import csv
 import io
+import html
 
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, Union, List
 
 
 class DatasetExportMixin:
-    """Mixin class"""
+    """Mixin class for exporting Dataset objects."""
 
     def relevant_columns(
         self, data_type: Optional[str] = None, remove_prefix=False
     ) -> list:
         """Return the set of keys that are present in the dataset.
 
+        :param data_type: The data type to filter by.
+        :param remove_prefix: Whether to remove the prefix from the column names.
+
         >>> from edsl.results.Dataset import Dataset
         >>> d = Dataset([{'a.b':[1,2,3,4]}])
         >>> d.relevant_columns()
@@ -23,22 +27,70 @@ class DatasetExportMixin:
         >>> d.relevant_columns(remove_prefix=True)
         ['b']
 
+        >>> d = Dataset([{'a':[1,2,3,4]}, {'b':[5,6,7,8]}])
+        >>> d.relevant_columns()
+        ['a', 'b']
+
         >>> from edsl.results import Results; Results.example().select('how_feeling', 'how_feeling_yesterday').relevant_columns()
         ['answer.how_feeling', 'answer.how_feeling_yesterday']
+
+        >>> from edsl.results import Results
+        >>> sorted(Results.example().select().relevant_columns(data_type = "model"))
+        ['model.frequency_penalty', 'model.logprobs', 'model.max_tokens', 'model.model', 'model.presence_penalty', 'model.temperature', 'model.top_logprobs', 'model.top_p']
+
+        >>> Results.example().relevant_columns(data_type = "flimflam")
+        Traceback (most recent call last):
+        ...
+        ValueError: No columns found for data type: flimflam. Available data types are: ['agent', 'answer', 'comment', 'model', 'prompt', 'question_options', 'question_text', 'question_type', 'raw_model_response', 'scenario'].
         """
         columns = [list(x.keys())[0] for x in self]
-        # columns = set([list(result.keys())[0] for result in self.data])
         if remove_prefix:
             columns = [column.split(".")[-1] for column in columns]
 
+        def get_data_type(column):
+            if "." in column:
+                return column.split(".")[0]
+            else:
+                return None
+
         if data_type:
+            all_columns = columns[:]
             columns = [
-                column for column in columns if column.split(".")[0] == data_type
+                column for column in columns if get_data_type(column) == data_type
             ]
+            if len(columns) == 0:
+                all_data_types = sorted(
+                    list(set(get_data_type(column) for column in all_columns))
+                )
+                raise ValueError(
+                    f"No columns found for data type: {data_type}. Available data types are: {all_data_types}."
+                )
 
         return columns
 
-    def _make_tabular(self, remove_prefix: bool, pretty_labels: Optional[dict] = None):
+    def num_observations(self):
+        """Return the number of observations in the dataset.
+
+        >>> from edsl.results import Results
+        >>> Results.example().num_observations()
+        4
+        """
+        _num_observations = None
+        for entry in self:
+            key, values = list(entry.items())[0]
+            if _num_observations is None:
+                _num_observations = len(values)
+            else:
+                if len(values) != _num_observations:
+                    raise ValueError(
+                        "The number of observations is not consistent across columns."
+                    )
+
+        return _num_observations
+
+    def _make_tabular(
+        self, remove_prefix: bool, pretty_labels: Optional[dict] = None
+    ) -> tuple[list, List[list]]:
         """Turn the results into a tabular format.
 
         :param remove_prefix: Whether to remove the prefix from the column names.
@@ -51,27 +103,41 @@ class DatasetExportMixin:
         >>> r.select('how_feeling')._make_tabular(remove_prefix = True, pretty_labels = {'how_feeling': "How are you feeling"})
         (['How are you feeling'], [['OK'], ['Great'], ['Terrible'], ['OK']])
         """
-        d = {}
-        full_header = sorted(list(self.relevant_columns()))
-        for entry in self.data:
-            key, list_of_values = list(entry.items())[0]
-            d[key] = list_of_values
+
+        def create_dict_from_list_of_dicts(list_of_dicts):
+            for entry in list_of_dicts:
+                key, list_of_values = list(entry.items())[0]
+                yield key, list_of_values
+
+        tabular_repr = dict(create_dict_from_list_of_dicts(self.data))
+
+        full_header = [list(x.keys())[0] for x in self]
+
+        rows = []
+        for i in range(self.num_observations()):
+            row = [tabular_repr[h][i] for h in full_header]
+            rows.append(row)
+
         if remove_prefix:
             header = [h.split(".")[-1] for h in full_header]
         else:
             header = full_header
-        num_observations = len(list(self[0].values())[0])
-        rows = []
-        # rows.append(header)
-        for i in range(num_observations):
-            row = [d[h][i] for h in full_header]
-            rows.append(row)
+
         if pretty_labels is not None:
             header = [pretty_labels.get(h, h) for h in header]
+
         return header, rows
 
     def print_long(self):
-        """Print the results in a long format."""
+        """Print the results in a long format.
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r.select('how_feeling').print_long()
+        answer.how_feeling: OK
+        answer.how_feeling: Great
+        answer.how_feeling: Terrible
+        answer.how_feeling: OK
+        """
         for entry in self:
             key, list_of_values = list(entry.items())[0]
             for value in list_of_values:
@@ -81,7 +147,7 @@ class DatasetExportMixin:
         self,
         pretty_labels: Optional[dict] = None,
         filename: Optional[str] = None,
-        format: Literal["rich", "html", "markdown", "latex"] = None,
+        format: Optional[Literal["rich", "html", "markdown", "latex"]] = None,
         interactive: bool = False,
         split_at_dot: bool = True,
         max_rows=None,
@@ -98,6 +164,12 @@ class DatasetExportMixin:
         :param format: The format to print the results in. Options are 'rich', 'html', or 'markdown'.
         :param interactive: Whether to print the results interactively in a Jupyter notebook.
         :param split_at_dot: Whether to split the column names at the last dot w/ a newline.
+        :param max_rows: The maximum number of rows to print.
+        :param tee: Whether to return the dataset.
+        :param iframe: Whether to display the table in an iframe.
+        :param iframe_height: The height of the iframe.
+        :param iframe_width: The width of the iframe.
+        :param web: Whether to display the table in a web browser.
 
         Example: Print in rich format at the terminal
 
@@ -117,6 +189,42 @@ class DatasetExportMixin:
         │ OK           │
         └──────────────┘
 
+        >>> r = Results.example()
+        >>> r2 = r.select("how_feeling").print(format = "rich", tee = True, max_rows = 2)
+        ┏━━━━━━━━━━━━━━┓
+        ┃ answer       ┃
+        ┃ .how_feeling ┃
+        ┡━━━━━━━━━━━━━━┩
+        │ OK           │
+        ├──────────────┤
+        │ Great        │
+        └──────────────┘
+        >>> r2
+        Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}])
+
+        >>> r.select('how_feeling').print(format = "rich", max_rows = 2)
+        ┏━━━━━━━━━━━━━━┓
+        ┃ answer       ┃
+        ┃ .how_feeling ┃
+        ┡━━━━━━━━━━━━━━┩
+        │ OK           │
+        ├──────────────┤
+        │ Great        │
+        └──────────────┘
+
+        >>> r.select('how_feeling').print(format = "rich", split_at_dot = False)
+        ┏━━━━━━━━━━━━━━━━━━━━┓
+        ┃ answer.how_feeling ┃
+        ┡━━━━━━━━━━━━━━━━━━━━┩
+        │ OK                 │
+        ├────────────────────┤
+        │ Great              │
+        ├────────────────────┤
+        │ Terrible           │
+        ├────────────────────┤
+        │ OK                 │
+        └────────────────────┘
+
         Example: using the pretty_labels parameter
 
         >>> r.select('how_feeling').print(format="rich", pretty_labels = {'answer.how_feeling': "How are you feeling"})
@@ -142,33 +250,40 @@ class DatasetExportMixin:
         | Terrible |
         | OK |
         ...
+
+        >>> r.select('how_feeling').print(format='latex')
+        \\begin{tabular}{l}
+        \\toprule
+        ...
         """
         from IPython.display import HTML, display
         from edsl.utilities.utilities import is_notebook
 
-        if format is None:
-            if is_notebook():
-                format = "html"
-            else:
-                format = "rich"
+        def _determine_format(format):
+            if format is None:
+                if is_notebook():
+                    format = "html"
+                else:
+                    format = "rich"
+            if format not in ["rich", "html", "markdown", "latex"]:
+                raise ValueError("format must be one of 'rich', 'html', or 'markdown'.")
+
+            return format
+
+        format = _determine_format(format)
 
         if pretty_labels is None:
             pretty_labels = {}
 
-        if format not in ["rich", "html", "markdown", "latex"]:
-            raise ValueError("format must be one of 'rich', 'html', or 'markdown'.")
+        if pretty_labels != {}:  # only split at dot if there are no pretty labels
+            split_at_dot = False
 
-        new_data = []
-        for index, entry in enumerate(self):
-            key, list_of_values = list(entry.items())[0]
-            new_data.append({pretty_labels.get(key, key): list_of_values})
+        def _create_data():
+            for index, entry in enumerate(self):
+                key, list_of_values = list(entry.items())[0]
+                yield {pretty_labels.get(key, key): list_of_values[:max_rows]}
 
-        if max_rows is not None:
-            for entry in new_data:
-                for key in entry:
-                    actual_rows = len(entry[key])
-                    entry[key] = entry[key][:max_rows]
-            # print(f"Showing only the first {max_rows} rows of {actual_rows} rows.")
+        new_data = list(_create_data())
 
         if format == "rich":
             from edsl.utilities.interface import print_dataset_with_rich
@@ -176,55 +291,54 @@ class DatasetExportMixin:
             print_dataset_with_rich(
                 new_data, filename=filename, split_at_dot=split_at_dot
             )
-        elif format == "html":
-            notebook = is_notebook()
+            return self if tee else None
+
+        if format == "markdown":
+            from edsl.utilities.interface import print_list_of_dicts_as_markdown_table
+
+            print_list_of_dicts_as_markdown_table(new_data, filename=filename)
+            return self if tee else None
+
+        if format == "latex":
+            df = self.to_pandas()
+            df.columns = [col.replace("_", " ") for col in df.columns]
+            latex_string = df.to_latex(index=False)
+
+            if filename is not None:
+                with open(filename, "w") as f:
+                    f.write(latex_string)
+            else:
+                print(latex_string)
+
+            return self if tee else None
+
+        if format == "html":
             from edsl.utilities.interface import print_list_of_dicts_as_html_table
 
             html_source = print_list_of_dicts_as_html_table(
                 new_data, interactive=interactive
             )
-            if iframe:
-                import html
 
-                height = iframe_height
-                width = iframe_width
-                escaped_output = html.escape(html_source)
-                # escaped_output = html_source
+            # if download_link:
+            #     from IPython.display import HTML, display
+            #     csv_file = output.getvalue()
+            #     b64 = base64.b64encode(csv_file.encode()).decode()
+            #     download_link = f'<a href="data:file/csv;base64,{b64}" download="my_data.csv">Download CSV file</a>'
+            #     #display(HTML(download_link))
+
+            if iframe:
                 iframe = f""""
-                <iframe srcdoc="{ escaped_output }" style="width: {width}px; height: {height}px;"></iframe>
+                <iframe srcdoc="{ html.escape(html_source) }" style="width: {iframe_width}px; height: {iframe_height}px;"></iframe>
                 """
                 display(HTML(iframe))
-            elif notebook:
+            elif is_notebook():
                 display(HTML(html_source))
             else:
                 from edsl.utilities.interface import view_html
 
                 view_html(html_source)
 
-        elif format == "markdown":
-            from edsl.utilities.interface import print_list_of_dicts_as_markdown_table
-
-            print_list_of_dicts_as_markdown_table(new_data, filename=filename)
-        elif format == "latex":
-            df = self.to_pandas()
-            df.columns = [col.replace("_", " ") for col in df.columns]
-            latex_string = df.to_latex()
-            if filename is not None:
-                with open(filename, "w") as f:
-                    f.write(latex_string)
-            else:
-                return latex_string
-            # raise NotImplementedError("Latex format not yet implemented.")
-            # latex_string = create_latex_table_from_data(new_data, filename=filename)
-            # if filename is None:
-            #     return latex_string
-            # Not working quite
-
-        else:
-            raise ValueError("format not recognized.")
-
-        if tee:
-            return self
+            return self if tee else None
 
     def to_csv(
         self,
@@ -245,6 +359,25 @@ class DatasetExportMixin:
         >>> r = Results.example()
         >>> r.select('how_feeling').to_csv()
         'answer.how_feeling\\r\\nOK\\r\\nGreat\\r\\nTerrible\\r\\nOK\\r\\n'
+
+        >>> r.select('how_feeling').to_csv(pretty_labels = {'answer.how_feeling': "How are you feeling"})
+        'How are you feeling\\r\\nOK\\r\\nGreat\\r\\nTerrible\\r\\nOK\\r\\n'
+
+        >>> import tempfile
+        >>> filename = tempfile.NamedTemporaryFile(delete=False).name
+        >>> r.select('how_feeling').to_csv(filename = filename)
+        >>> import os
+        >>> import csv
+        >>> with open(filename, newline='') as f:
+        ...     reader = csv.reader(f)
+        ...     for row in reader:
+        ...         print(row)
+        ['answer.how_feeling']
+        ['OK']
+        ['Great']
+        ['Terrible']
+        ['OK']
+
         """
         if pretty_labels is None:
             pretty_labels = {}
@@ -264,6 +397,8 @@ class DatasetExportMixin:
             writer.writerows(rows)
 
             if download_link:
+                from IPython.display import HTML, display
+
                 csv_file = output.getvalue()
                 b64 = base64.b64encode(csv_file.encode()).decode()
                 download_link = f'<a href="data:file/csv;base64,{b64}" download="my_data.csv">Download CSV file</a>'
@@ -271,6 +406,22 @@ class DatasetExportMixin:
             else:
                 return output.getvalue()
 
+    def download_link(self, pretty_labels: Optional[dict] = None) -> str:
+        """Return a download link for the results.
+
+        :param pretty_labels: A dictionary of pretty labels for the columns.
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r.select('how_feeling').download_link()
+        '<a href="data:file/csv;base64,YW5zd2VyLmhvd19mZWVsaW5nDQpPSw0KR3JlYXQNClRlcnJpYmxlDQpPSw0K" download="my_data.csv">Download CSV file</a>'
+        """
+        import base64
+
+        csv_string = self.to_csv(pretty_labels=pretty_labels)
+        b64 = base64.b64encode(csv_string.encode()).decode()
+        return f'<a href="data:file/csv;base64,{b64}" download="my_data.csv">Download CSV file</a>'
+
     def to_pandas(self, remove_prefix: bool = False) -> "pd.DataFrame":
         """Convert the results to a pandas DataFrame.
 
@@ -290,8 +441,8 @@ class DatasetExportMixin:
         csv_string = self.to_csv(remove_prefix=remove_prefix)
         csv_buffer = io.StringIO(csv_string)
         df = pd.read_csv(csv_buffer)
-        df_sorted = df.sort_index(axis=1)  # Sort columns alphabetically
-        return df_sorted
+        # df_sorted = df.sort_index(axis=1)  # Sort columns alphabetically
+        return df
 
     def to_scenario_list(self, remove_prefix: bool = True) -> list[dict]:
         """Convert the results to a list of dictionaries, one per scenario.
@@ -309,6 +460,15 @@ class DatasetExportMixin:
         return ScenarioList([Scenario(d) for d in list_of_dicts])
 
     def to_agent_list(self, remove_prefix: bool = True):
+        """Convert the results to a list of dictionaries, one per agent.
+
+        :param remove_prefix: Whether to remove the prefix from the column names.
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r.select('how_feeling').to_agent_list()
+        AgentList([Agent(traits = {'how_feeling': 'OK'}), Agent(traits = {'how_feeling': 'Great'}), Agent(traits = {'how_feeling': 'Terrible'}), Agent(traits = {'how_feeling': 'OK'})])
+        """
         from edsl import AgentList, Agent
 
         list_of_dicts = self.to_dicts(remove_prefix=remove_prefix)
@@ -344,6 +504,9 @@ class DatasetExportMixin:
     def to_list(self, flatten=False, remove_none=False) -> list[list]:
         """Convert the results to a list of lists.
 
+        :param flatten: Whether to flatten the list of lists.
+        :param remove_none: Whether to remove None values from the list.
+
         >>> from edsl.results import Results
         >>> Results.example().select('how_feeling', 'how_feeling_yesterday')
         Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}, {'answer.how_feeling_yesterday': ['Great', 'Good', 'OK', 'Terrible']}])
@@ -354,6 +517,18 @@ class DatasetExportMixin:
         >>> r = Results.example()
         >>> r.select('how_feeling').to_list()
         ['OK', 'Great', 'Terrible', 'OK']
+
+        >>> from edsl.results.Dataset import Dataset
+        >>> Dataset([{'a.b': [[1, 9], 2, 3, 4]}]).select('a.b').to_list(flatten = True)
+        [1, 9, 2, 3, 4]
+
+        >>> from edsl.results.Dataset import Dataset
+        >>> Dataset([{'a.b': [[1, 9], 2, 3, 4]}, {'c': [6, 2, 3, 4]}]).select('a.b', 'c').to_list(flatten = True)
+        Traceback (most recent call last):
+        ...
+        ValueError: Cannot flatten a list of lists when there are multiple columns selected.
+
+
         """
         if len(self.relevant_columns()) > 1 and flatten:
             raise ValueError(
@@ -385,7 +560,10 @@ class DatasetExportMixin:
         return list_to_return
 
     def html(
-        self, filename: str = None, cta: str = "Open in browser", return_link=False
+        self,
+        filename: Optional[str] = None,
+        cta: str = "Open in browser",
+        return_link: bool = False,
     ):
         import os
         import tempfile
@@ -419,7 +597,7 @@ class DatasetExportMixin:
             return filename
 
     def tally(
-        self, *fields: Optional[str], top_n=None, output="dict"
+        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.
 
@@ -427,9 +605,11 @@ class DatasetExportMixin:
 
         >>> from edsl.results import Results
         >>> r = Results.example()
-        >>> r.select('how_feeling').tally('answer.how_feeling')
+        >>> r.select('how_feeling').tally('answer.how_feeling', output = "dict")
         {'OK': 2, 'Great': 1, 'Terrible': 1}
-        >>> r.select('how_feeling', 'period').tally('how_feeling', 'period')
+        >>> r.select('how_feeling').tally('answer.how_feeling', output = "Dataset")
+        Dataset([{'value': ['OK', 'Great', 'Terrible']}, {'count': [2, 1, 1]}])
+        >>> r.select('how_feeling', 'period').tally('how_feeling', 'period', output = "dict")
         {('OK', 'morning'): 1, ('Great', 'afternoon'): 1, ('Terrible', 'morning'): 1, ('OK', 'afternoon'): 1}
         """
         from collections import Counter
@@ -441,6 +621,8 @@ class DatasetExportMixin:
             column.split(".")[-1] for column in self.relevant_columns()
         ]
 
+        # breakpoint()
+
         if not all(
             f in self.relevant_columns() or f in relevant_columns_without_prefix
             for f in fields
@@ -467,6 +649,7 @@ class DatasetExportMixin:
         from edsl.results.Dataset import Dataset
 
         if output == "dict":
+            # why did I do this?
             warnings.warn(
                 textwrap.dedent(
                     """\

edsl/results/Result.py

@@ -126,6 +126,9 @@ class Result(Base, UserDict):
         self.survey = survey
         self.question_to_attributes = question_to_attributes
 
+        self._combined_dict = None
+        self._problem_keys = None
+
     ###############
     # Used in Results
     ###############
@@ -164,25 +167,64 @@ class Result(Base, UserDict):
             "answer": self.answer,
             "prompt": self.prompt,
             "raw_model_response": self.raw_model_response,
-            "iteration": {"iteration": self.iteration},
+            #            "iteration": {"iteration": self.iteration},
             "question_text": question_text_dict,
             "question_options": question_options_dict,
             "question_type": question_type_dict,
             "comment": comments_dict,
         }
 
+    def check_expression(self, expression) -> None:
+        for key in self.problem_keys:
+            if key in expression and not key + "." in expression:
+                raise ValueError(
+                    f"Key by iself {key} is problematic. Use the full key {key + '.' + key} name instead."
+                )
+        return None
+
     def code(self):
         """Return a string of code that can be used to recreate the Result object."""
         raise NotImplementedError
 
     @property
-    def combined_dict(self) -> dict[str, Any]:
-        """Return a dictionary that includes all sub_dicts, but also puts the key-value pairs in each sub_dict as a key_value pair in the combined dictionary."""
+    def problem_keys(self):
+        """Return a list of keys that are problematic."""
+        return self._problem_keys
+
+    def _compute_combined_dict_and_problem_keys(self) -> None:
         combined = {}
+        problem_keys = []
         for key, sub_dict in self.sub_dicts.items():
             combined.update(sub_dict)
+            # in some cases, the sub_dict might have keys that conflict with the main dict
+            if key in combined:
+                # The key is already in the combined dict
+                problem_keys = problem_keys + [key]
+
             combined.update({key: sub_dict})
-        return combined
+            # I *think* this allows us to do do things like "answer.how_feelling" i.e., that the evaluator can use
+            # dot notation to access the subdicts.
+        self._combined_dict = combined
+        self._problem_keys = problem_keys
+
+    @property
+    def combined_dict(self) -> dict[str, Any]:
+        """Return a dictionary that includes all sub_dicts, but also puts the key-value pairs in each sub_dict as a key_value pair in the combined dictionary.
+
+        >>> r = Result.example()
+        >>> r.combined_dict['how_feeling']
+        'OK'
+        """
+        if self._combined_dict is None or self._problem_keys is None:
+            self._compute_combined_dict_and_problem_keys()
+        return self._combined_dict
+
+    @property
+    def problem_keys(self):
+        """Return a list of keys that are problematic."""
+        if self._combined_dict is None or self._problem_keys is None:
+            self._compute_combined_dict_and_problem_keys()
+        return self._problem_keys
 
     def get_value(self, data_type: str, key: str) -> Any:
         """Return the value for a given data type and key.
@@ -226,7 +268,13 @@ class Result(Base, UserDict):
         return Result.from_dict(self.to_dict())
 
     def __eq__(self, other) -> bool:
-        """Return True if the Result object is equal to another Result object."""
+        """Return True if the Result object is equal to another Result object.
+
+        >>> r = Result.example()
+        >>> r == r
+        True
+
+        """
         return self.to_dict() == other.to_dict()
 
     ###############