edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/results/DatasetExportMixin.py

@@ -1,11 +1,12 @@
 """Mixin class for exporting results."""
 
+import base64
+import csv
 import io
-import warnings
-import textwrap
-from typing import Optional, Tuple, Union, List
+import html
+from typing import Optional
 
-from edsl.results.file_exports import CSVExport, ExcelExport, JSONLExport, SQLiteExport
+from typing import Literal, Optional, Union, List
 
 
 class DatasetExportMixin:
@@ -36,7 +37,7 @@ class DatasetExportMixin:
 
         >>> from edsl.results import Results
         >>> sorted(Results.example().select().relevant_columns(data_type = "model"))
-        ['model.frequency_penalty', ...]
+        ['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):
@@ -71,7 +72,7 @@ class DatasetExportMixin:
     def num_observations(self):
         """Return the number of observations in the dataset.
 
-        >>> from edsl.results.Results import Results
+        >>> from edsl.results import Results
         >>> Results.example().num_observations()
         4
         """
@@ -143,148 +144,299 @@ class DatasetExportMixin:
             for value in list_of_values:
                 print(f"{key}: {value}")
 
-    def _get_tabular_data(
-        self,
-        remove_prefix: bool = False,
-        pretty_labels: Optional[dict] = None,
-    ) -> Tuple[List[str], List[List]]:
-        """Internal method to get tabular data in a standard format.
-
-        Args:
-            remove_prefix: Whether to remove the prefix from column names
-            pretty_labels: Dictionary mapping original column names to pretty labels
-
-        Returns:
-            Tuple containing (header_row, data_rows)
-        """
-        if pretty_labels is None:
-            pretty_labels = {}
-
-        return self._make_tabular(
-            remove_prefix=remove_prefix, pretty_labels=pretty_labels
-        )
-
-    def to_jsonl(self, filename: Optional[str] = None) -> Optional["FileStore"]:
-        """Export the results to a FileStore instance containing JSONL data."""
-        exporter = JSONLExport(data=self, filename=filename)
-        return exporter.export()
-
-    def to_sqlite(
-        self,
-        filename: Optional[str] = None,
-        remove_prefix: bool = False,
-        pretty_labels: Optional[dict] = None,
-        table_name: str = "results",
-        if_exists: str = "replace",
-    ) -> Optional["FileStore"]:
-        """Export the results to a SQLite database file."""
-        exporter = SQLiteExport(
-            data=self,
-            filename=filename,
-            remove_prefix=remove_prefix,
-            pretty_labels=pretty_labels,
-            table_name=table_name,
-            if_exists=if_exists,
-        )
-        return exporter.export()
+    # 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,
         filename: Optional[str] = None,
         remove_prefix: bool = False,
+        download_link: bool = False,
         pretty_labels: Optional[dict] = None,
-    ) -> Optional["FileStore"]:
-        """Export the results to a FileStore instance containing CSV data."""
-        exporter = CSVExport(
-            data=self,
-            filename=filename,
-            remove_prefix=remove_prefix,
-            pretty_labels=pretty_labels,
-        )
-        return exporter.export()
+    ):
+        """Export the results to a CSV file.
 
-    def to_excel(
-        self,
-        filename: Optional[str] = None,
-        remove_prefix: bool = False,
-        pretty_labels: Optional[dict] = None,
-        sheet_name: Optional[str] = None,
-    ) -> Optional["FileStore"]:
-        """Export the results to a FileStore instance containing Excel data."""
-        exporter = ExcelExport(
-            data=self,
-            filename=filename,
-            remove_prefix=remove_prefix,
-            pretty_labels=pretty_labels,
-            sheet_name=sheet_name,
-        )
-        return exporter.export()
+        :param filename: The filename to save the CSV file to.
+        :param remove_prefix: Whether to remove the prefix from the column names.
+        :param download_link: Whether to display a download link in a Jupyter notebook.
 
-    def _db(self, remove_prefix: bool = True):
-        """Create a SQLite database in memory and return the connection.
+        Example:
 
-        Args:
-            shape: The shape of the data in the database (wide or long)
-            remove_prefix: Whether to remove the prefix from the column names
+        >>> from edsl.results import Results
+        >>> 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']
 
-        Returns:
-            A database connection
         """
-        from sqlalchemy import create_engine
-
-        engine = create_engine("sqlite:///:memory:")
-        if remove_prefix:
-            df = self.remove_prefix().to_pandas(lists_as_strings=True)
-        else:
-            df = self.to_pandas(lists_as_strings=True)
-        df.to_sql(
-            "self",
-            engine,
-            index=False,
-            if_exists="replace",
+        if pretty_labels is None:
+            pretty_labels = {}
+        header, rows = self._make_tabular(
+            remove_prefix=remove_prefix, pretty_labels=pretty_labels
         )
-        return engine.connect()
 
-    def sql(
-        self,
-        query: str,
-        transpose: bool = None,
-        transpose_by: str = None,
-        remove_prefix: bool = True,
-    ) -> Union["pd.DataFrame", str]:
-        """Execute a SQL query and return the results as a DataFrame.
-
-        Args:
-            query: The SQL query to execute
-            shape: The shape of the data in the database (wide or long)
-            remove_prefix: Whether to remove the prefix from the column names
-            transpose: Whether to transpose the DataFrame
-            transpose_by: The column to use as the index when transposing
-            csv: Whether to return the DataFrame as a CSV string
-            to_list: Whether to return the results as a list
-            to_latex: Whether to return the results as LaTeX
-            filename: Optional filename to save the results to
-
-        Returns:
-            DataFrame, CSV string, list, or LaTeX string depending on parameters
+        if filename is not None:
+            with open(filename, "w") as f:
+                writer = csv.writer(f)
+                writer.writerow(header)
+                writer.writerows(rows)
+            # print(f"Saved to {filename}")
+        else:
+            output = io.StringIO()
+            writer = csv.writer(output)
+            writer.writerow(header)
+            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>'
+                display(HTML(download_link))
+            else:
+                return output.getvalue()
 
-        """
-        import pandas as pd
+    def download_link(self, pretty_labels: Optional[dict] = None) -> str:
+        """Return a download link for the results.
 
-        conn = self._db(remove_prefix=remove_prefix)
-        df = pd.read_sql_query(query, conn)
+        :param pretty_labels: A dictionary of pretty labels for the columns.
 
-        # Transpose the DataFrame if transpose is True
-        if transpose or transpose_by:
-            df = pd.DataFrame(df)
-            if transpose_by:
-                df = df.set_index(transpose_by)
-            else:
-                df = df.set_index(df.columns[0])
-            df = df.transpose()
-        from edsl.results.Dataset import Dataset
+        >>> 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
 
-        return Dataset.from_pandas_dataframe(df)
+        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, lists_as_strings=False
@@ -295,6 +447,19 @@ class DatasetExportMixin:
 
         """
         return self._to_pandas_strings(remove_prefix)
+        # if lists_as_strings:
+        #     return self._to_pandas_strings(remove_prefix=remove_prefix)
+
+        # import pandas as pd
+
+        # df = pd.DataFrame(self.data)
+
+        # if remove_prefix:
+        #     # Optionally remove prefixes from column names
+        #     df.columns = [col.split(".")[-1] for col in df.columns]
+
+        # df_sorted = df.sort_index(axis=1)  # Sort columns alphabetically
+        # return df_sorted
 
     def _to_pandas_strings(self, remove_prefix: bool = False) -> "pd.DataFrame":
         """Convert the results to a pandas DataFrame.
@@ -313,32 +478,12 @@ class DatasetExportMixin:
 
         import pandas as pd
 
-        csv_string = self.to_csv(remove_prefix=remove_prefix).text
+        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
 
-    def to_polars(
-        self, remove_prefix: bool = False, lists_as_strings=False
-    ) -> "pl.DataFrame":
-        """Convert the results to a Polars DataFrame.
-
-        :param remove_prefix: Whether to remove the prefix from the column names.
-        """
-        return self._to_polars_strings(remove_prefix)
-
-    def _to_polars_strings(self, remove_prefix: bool = False) -> "pl.DataFrame":
-        """Convert the results to a Polars DataFrame.
-
-        :param remove_prefix: Whether to remove the prefix from the column names.
-        """
-        import polars as pl
-
-        csv_string = self.to_csv(remove_prefix=remove_prefix).text
-        df = pl.read_csv(io.StringIO(csv_string))
-        return df
-
     def to_scenario_list(self, remove_prefix: bool = True) -> list[dict]:
         """Convert the results to a list of dictionaries, one per scenario.
 
@@ -349,14 +494,14 @@ class DatasetExportMixin:
         >>> r.select('how_feeling').to_scenario_list()
         ScenarioList([Scenario({'how_feeling': 'OK'}), Scenario({'how_feeling': 'Great'}), Scenario({'how_feeling': 'Terrible'}), Scenario({'how_feeling': 'OK'})])
         """
-        from edsl.scenarios.ScenarioList import ScenarioList
-        from edsl.scenarios.Scenario import Scenario
+        from edsl import ScenarioList, Scenario
 
         list_of_dicts = self.to_dicts(remove_prefix=remove_prefix)
         scenarios = []
         for d in list_of_dicts:
             scenarios.append(Scenario(d))
         return ScenarioList(scenarios)
+        # 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.
@@ -368,8 +513,7 @@ class DatasetExportMixin:
         >>> 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.agents import Agent
-        from edsl.agents.AgentList import AgentList
+        from edsl import AgentList, Agent
 
         list_of_dicts = self.to_dicts(remove_prefix=remove_prefix)
         agents = []
@@ -377,11 +521,6 @@ class DatasetExportMixin:
             if "name" in d:
                 d["agent_name"] = d.pop("name")
                 agents.append(Agent(d, name=d["agent_name"]))
-            if "agent_parameters" in d:
-                agent_parameters = d.pop("agent_parameters")
-                agent_name = agent_parameters.get("name", None)
-                instruction = agent_parameters.get("instruction", None)
-                agents.append(Agent(d, name=agent_name, instruction=instruction))
             else:
                 agents.append(Agent(d))
         return AgentList(agents)
@@ -469,9 +608,7 @@ class DatasetExportMixin:
                     new_list.append(item)
             list_to_return = new_list
 
-        from edsl.utilities.PrettyList import PrettyList
-
-        return PrettyList(list_to_return)
+        return list_to_return
 
     def html(
         self,
@@ -521,10 +658,8 @@ class DatasetExportMixin:
         >>> r = Results.example()
         >>> r.select('how_feeling').tally('answer.how_feeling', output = "dict")
         {'OK': 2, 'Great': 1, 'Terrible': 1}
-        >>> from edsl.results.Dataset import Dataset
-        >>> expected = Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}, {'count': [2, 1, 1]}])
-        >>> r.select('how_feeling').tally('answer.how_feeling', output = "Dataset") == expected
-        True
+        >>> r.select('how_feeling').tally('answer.how_feeling', output = "Dataset")
+        Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible']}, {'count': [2, 1, 1]}])
         """
         from collections import Counter
 
@@ -556,6 +691,8 @@ class DatasetExportMixin:
         if top_n is not None:
             sorted_tally = dict(list(sorted_tally.items())[:top_n])
 
+        import warnings
+        import textwrap
         from edsl.results.Dataset import Dataset
 
         if output == "dict":

edsl/results/DatasetTree.py

@@ -1,4 +1,8 @@
 from typing import Dict, List, Any, Optional, List
+from docx import Document
+from docx.shared import Inches, Pt
+from docx.enum.text import WD_ALIGN_PARAGRAPH
+from docx.enum.style import WD_STYLE_TYPE
 
 
 def is_hashable(v):
@@ -94,11 +98,6 @@ class Tree:
         if filename is None:
             filename = "tree_structure.docx"
 
-        from docx import Document
-        from docx.shared import Inches, Pt
-        from docx.enum.text import WD_ALIGN_PARAGRAPH
-        from docx.enum.style import WD_STYLE_TYPE
-
         doc = Document()
 
         # Create styles for headings
@@ -116,29 +115,10 @@ class Tree:
         body_style.font.size = Pt(11)
 
         self._add_to_docx(doc, self.root, 0)
-        import base64
-        from io import BytesIO
-        import base64
-
-        # Save document to bytes buffer
-        doc_buffer = BytesIO()
-        doc.save(doc_buffer)
-        doc_buffer.seek(0)
-
-        base64_string = base64.b64encode(doc_buffer.getvalue()).decode("utf-8")
-        from edsl.scenarios.FileStore import FileStore
-
-        # Create and return FileStore instance
-        return FileStore(
-            path="tree_structure.docx",  # Default name
-            mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
-            binary=True,
-            suffix="docx",
-            base64_string=base64_string,
-        )
-        # doc.save(filename)
-        # from edsl.utilities.utilities import file_notice
-        # file_notice(filename)
+        doc.save(filename)
+        from edsl.utilities.utilities import file_notice
+
+        file_notice(filename)
 
     def _repr_html_(self):
         """Returns an interactive HTML representation of the tree with collapsible sections."""