edsl 0.1.38.dev2__py3-none-any.whl → 0.1.38.dev4__py3-none-any.whl

edsl/results/DatasetExportMixin.py

@@ -144,216 +144,216 @@ class DatasetExportMixin:
             for value in list_of_values:
                 print(f"{key}: {value}")
 
-    def print(
-        self,
-        pretty_labels: Optional[dict] = None,
-        filename: Optional[str] = None,
-        format: Optional[Literal["rich", "html", "markdown", "latex"]] = None,
-        interactive: bool = False,
-        split_at_dot: bool = True,
-        max_rows=None,
-        tee=False,
-        iframe=False,
-        iframe_height: int = 200,
-        iframe_width: int = 600,
-        web=False,
-        return_string: bool = False,
-    ) -> Union[None, str, "Results"]:
-        """Print the results in a pretty format.
-
-        :param pretty_labels: A dictionary of pretty labels for the columns.
-        :param filename: The filename to save the results to.
-        :param format: The format to print the results in. Options are 'rich', 'html', 'markdown', or 'latex'.
-        :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.
-        :param return_string: Whether to return the output as a string instead of printing.
-
-        :return: None if tee is False and return_string is False, the dataset if tee is True, or a string if return_string is True.
-
-        Example: Print in rich format at the terminal
-
-        >>> from edsl.results import Results
-        >>> r = Results.example()
-        >>> r.select('how_feeling').print(format = "rich")
-        ┏━━━━━━━━━━━━━━┓
-        ┃ answer       ┃
-        ┃ .how_feeling ┃
-        ┡━━━━━━━━━━━━━━┩
-        │ OK           │
-        ├──────────────┤
-        │ Great        │
-        ├──────────────┤
-        │ Terrible     │
-        ├──────────────┤
-        │ 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"})
-        ┏━━━━━━━━━━━━━━━━━━━━━┓
-        ┃ How are you feeling ┃
-        ┡━━━━━━━━━━━━━━━━━━━━━┩
-        │ OK                  │
-        ├─────────────────────┤
-        │ Great               │
-        ├─────────────────────┤
-        │ Terrible            │
-        ├─────────────────────┤
-        │ OK                  │
-        └─────────────────────┘
-
-        Example: printing in markdown format
-
-        >>> r.select('how_feeling').print(format='markdown')
-        | answer.how_feeling |
-        |--|
-        | OK |
-        | Great |
-        | Terrible |
-        | OK |
-        ...
-
-        >>> r.select('how_feeling').print(format='latex')
-        \\begin{tabular}{l}
-        ...
-        \\end{tabular}
-        <BLANKLINE>
-        """
-        from IPython.display import HTML, display
-        from edsl.utilities.utilities import is_notebook
-        import io
-        import sys
-
-        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', 'markdown', or 'latex'."
-                )
-
-            return format
-
-        format = _determine_format(format)
-
-        if pretty_labels is None:
-            pretty_labels = {}
-
-        if pretty_labels != {}:  # only split at dot if there are no pretty labels
-            split_at_dot = False
-
-        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]}
-
-        new_data = list(_create_data())
-
-        # Capture output if return_string is True
-        if return_string:
-            old_stdout = sys.stdout
-            sys.stdout = io.StringIO()
-
-        output = None
-
-        if format == "rich":
-            from edsl.utilities.interface import print_dataset_with_rich
-
-            output = print_dataset_with_rich(
-                new_data, filename=filename, split_at_dot=split_at_dot
-            )
-        elif format == "markdown":
-            from edsl.utilities.interface import print_list_of_dicts_as_markdown_table
-
-            output = 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(index=False)
-
-            if filename is not None:
-                with open(filename, "w") as f:
-                    f.write(latex_string)
-            else:
-                print(latex_string)
-                output = latex_string
-        elif 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:
-                iframe = f""""
-                <iframe srcdoc="{ html.escape(html_source) }" style="width: {iframe_width}px; height: {iframe_height}px;"></iframe>
-                """
-                display(HTML(iframe))
-            elif is_notebook():
-                display(HTML(html_source))
-            else:
-                from edsl.utilities.interface import view_html
-
-                view_html(html_source)
-
-            output = html_source
-
-        # Restore stdout and get captured output if return_string is True
-        if return_string:
-            captured_output = sys.stdout.getvalue()
-            sys.stdout = old_stdout
-            return captured_output or output
-
-        if tee:
-            return self
-
-        return None
+    # def print(
+    #     self,
+    #     pretty_labels: Optional[dict] = None,
+    #     filename: Optional[str] = None,
+    #     format: Optional[Literal["rich", "html", "markdown", "latex"]] = None,
+    #     interactive: bool = False,
+    #     split_at_dot: bool = True,
+    #     max_rows=None,
+    #     tee=False,
+    #     iframe=False,
+    #     iframe_height: int = 200,
+    #     iframe_width: int = 600,
+    #     web=False,
+    #     return_string: bool = False,
+    # ) -> Union[None, str, "Results"]:
+    #     """Print the results in a pretty format.
+
+    #     :param pretty_labels: A dictionary of pretty labels for the columns.
+    #     :param filename: The filename to save the results to.
+    #     :param format: The format to print the results in. Options are 'rich', 'html', 'markdown', or 'latex'.
+    #     :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.
+    #     :param return_string: Whether to return the output as a string instead of printing.
+
+    #     :return: None if tee is False and return_string is False, the dataset if tee is True, or a string if return_string is True.
+
+    #     Example: Print in rich format at the terminal
+
+    #     >>> from edsl.results import Results
+    #     >>> r = Results.example()
+    #     >>> r.select('how_feeling').print(format = "rich")
+    #     ┏━━━━━━━━━━━━━━┓
+    #     ┃ answer       ┃
+    #     ┃ .how_feeling ┃
+    #     ┡━━━━━━━━━━━━━━┩
+    #     │ OK           │
+    #     ├──────────────┤
+    #     │ Great        │
+    #     ├──────────────┤
+    #     │ Terrible     │
+    #     ├──────────────┤
+    #     │ 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"})
+    #     ┏━━━━━━━━━━━━━━━━━━━━━┓
+    #     ┃ How are you feeling ┃
+    #     ┡━━━━━━━━━━━━━━━━━━━━━┩
+    #     │ OK                  │
+    #     ├─────────────────────┤
+    #     │ Great               │
+    #     ├─────────────────────┤
+    #     │ Terrible            │
+    #     ├─────────────────────┤
+    #     │ OK                  │
+    #     └─────────────────────┘
+
+    #     Example: printing in markdown format
+
+    #     >>> r.select('how_feeling').print(format='markdown')
+    #     | answer.how_feeling |
+    #     |--|
+    #     | OK |
+    #     | Great |
+    #     | Terrible |
+    #     | OK |
+    #     ...
+
+    #     >>> r.select('how_feeling').print(format='latex')
+    #     \\begin{tabular}{l}
+    #     ...
+    #     \\end{tabular}
+    #     <BLANKLINE>
+    #     """
+    #     from IPython.display import HTML, display
+    #     from edsl.utilities.utilities import is_notebook
+    #     import io
+    #     import sys
+
+    #     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', 'markdown', or 'latex'."
+    #             )
+
+    #         return format
+
+    #     format = _determine_format(format)
+
+    #     if pretty_labels is None:
+    #         pretty_labels = {}
+
+    #     if pretty_labels != {}:  # only split at dot if there are no pretty labels
+    #         split_at_dot = False
+
+    #     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]}
+
+    #     new_data = list(_create_data())
+
+    #     # Capture output if return_string is True
+    #     if return_string:
+    #         old_stdout = sys.stdout
+    #         sys.stdout = io.StringIO()
+
+    #     output = None
+
+    #     if format == "rich":
+    #         from edsl.utilities.interface import print_dataset_with_rich
+
+    #         output = print_dataset_with_rich(
+    #             new_data, filename=filename, split_at_dot=split_at_dot
+    #         )
+    #     elif format == "markdown":
+    #         from edsl.utilities.interface import print_list_of_dicts_as_markdown_table
+
+    #         output = 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(index=False)
+
+    #         if filename is not None:
+    #             with open(filename, "w") as f:
+    #                 f.write(latex_string)
+    #         else:
+    #             print(latex_string)
+    #             output = latex_string
+    #     elif 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:
+    #             iframe = f""""
+    #             <iframe srcdoc="{ html.escape(html_source) }" style="width: {iframe_width}px; height: {iframe_height}px;"></iframe>
+    #             """
+    #             display(HTML(iframe))
+    #         elif is_notebook():
+    #             display(HTML(html_source))
+    #         else:
+    #             from edsl.utilities.interface import view_html
+
+    #             view_html(html_source)
+
+    #         output = html_source
+
+    #     # Restore stdout and get captured output if return_string is True
+    #     if return_string:
+    #         captured_output = sys.stdout.getvalue()
+    #         sys.stdout = old_stdout
+    #         return captured_output or output
+
+    #     if tee:
+    #         return self
+
+    #     return None
 
     def to_csv(
         self,
@@ -405,6 +405,7 @@ class DatasetExportMixin:
                 writer = csv.writer(f)
                 writer.writerow(header)
                 writer.writerows(rows)
+            # print(f"Saved to {filename}")
         else:
             output = io.StringIO()
             writer = csv.writer(output)
@@ -515,7 +516,14 @@ class DatasetExportMixin:
         from edsl import AgentList, Agent
 
         list_of_dicts = self.to_dicts(remove_prefix=remove_prefix)
-        return AgentList([Agent(d) for d in list_of_dicts])
+        agents = []
+        for d in list_of_dicts:
+            if "name" in d:
+                d["agent_name"] = d.pop("name")
+                agents.append(Agent(d, name=d["agent_name"]))
+            else:
+                agents.append(Agent(d))
+        return AgentList(agents)
 
     def to_dicts(self, remove_prefix: bool = True) -> list[dict]:
         """Convert the results to a list of dictionaries.
@@ -651,9 +659,7 @@ class DatasetExportMixin:
         >>> r.select('how_feeling').tally('answer.how_feeling', output = "dict")
         {'OK': 2, 'Great': 1, 'Terrible': 1}
         >>> 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}
+        Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}, {'count': [2, 1, 1]}])
         """
         from collections import Counter
 
@@ -664,8 +670,6 @@ 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
@@ -703,12 +707,22 @@ class DatasetExportMixin:
             )
             return sorted_tally
         elif output == "Dataset":
-            return Dataset(
+            dataset = Dataset(
                 [
                     {"value": list(sorted_tally.keys())},
                     {"count": list(sorted_tally.values())},
                 ]
             )
+            # return dataset
+            sl = dataset.to_scenario_list().unpack(
+                "value",
+                new_names=[fields] if isinstance(fields, str) else fields,
+                keep_original=False,
+            )
+            keys = list(sl[0].keys())
+            keys.remove("count")
+            keys.append("count")
+            return sl.reorder_keys(keys).to_dataset()
 
 
 if __name__ == "__main__":