edsl 0.1.39.dev2__py3-none-any.whl → 0.1.39.dev3__py3-none-any.whl

edsl/results/ResultsDBMixin.py

@@ -0,0 +1,238 @@
+"""Mixin for working with SQLite respresentation of a 'Results' object."""
+
+import sqlite3
+from enum import Enum
+from typing import Literal, Union, Optional
+
+
+class SQLDataShape(Enum):
+    """Enum for the shape of the data in the SQL database."""
+
+    WIDE = "wide"
+    LONG = "long"
+
+
+class ResultsDBMixin:
+    """Mixin for interacting with a Results object as if it were a SQL database."""
+
+    def _rows(self):
+        """Return the rows of the `Results` object as a list of tuples."""
+        for index, result in enumerate(self):
+            yield from result.rows(index)
+
+    def export_sql_dump(self, shape: Literal["wide", "long"], filename: str):
+        """Export the SQL database to a file.
+
+        :param shape: The shape of the data in the database (wide or long)
+        :param filename: The filename to save the database to
+        """
+        shape_enum = self._get_shape_enum(shape)
+        conn = self._db(shape=shape_enum)
+
+        with open(filename, "w") as f:
+            for line in conn.iterdump():
+                f.write(f"{line}\n")
+
+        conn.close()
+
+    def backup_db_to_file(self, shape: Literal["wide", "long"], filename: str):
+        """Backup the in-memory database to a file.
+
+
+        :param shape: The shape of the data in the database (wide or long)
+        :param filename: The filename to save the database to
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r.backup_db_to_file(filename="backup.db", shape="long")
+
+        """
+        shape_enum = self._get_shape_enum(shape)
+        # Source database connection (in-memory)
+        source_conn = self._db(shape=shape_enum)
+
+        # Destination database connection (file)
+        dest_conn = sqlite3.connect(filename)
+
+        # Backup in-memory database to file
+        with source_conn:
+            source_conn.backup(dest_conn)
+
+        # Close both connections
+        source_conn.close()
+        dest_conn.close()
+
+    def _db(self, shape: SQLDataShape, remove_prefix=False):
+        """Create a SQLite database in memory and return the connection.
+
+        :param shape: The shape of the data in the database (wide or long)
+        :param remove_prefix: Whether to remove the prefix from the column names
+
+        """
+        if shape == SQLDataShape.LONG:
+            conn = sqlite3.connect(":memory:")
+
+            create_table_query = """
+            CREATE TABLE self (
+                id INTEGER,
+                data_type TEXT,
+                key TEXT, 
+                value TEXT
+            )
+            """
+            conn.execute(create_table_query)
+
+            list_of_tuples = list(self._rows())
+            insert_query = (
+                "INSERT INTO self (id, data_type, key, value) VALUES (?, ?, ?, ?)"
+            )
+            conn.executemany(insert_query, list_of_tuples)
+            conn.commit()
+            return conn
+        elif shape == SQLDataShape.WIDE:
+            from sqlalchemy import create_engine
+
+            engine = create_engine("sqlite:///:memory:")
+            df = self.to_pandas(remove_prefix=remove_prefix, lists_as_strings=True)
+            df.to_sql("self", engine, index=False, if_exists="replace")
+            return engine.connect()
+        else:
+            raise Exception("Invalid SQLDataShape")
+
+    def _get_shape_enum(self, shape: Literal["wide", "long"]):
+        """Convert the shape string to a SQLDataShape enum."""
+        if shape is None:
+            raise Exception("Must select either 'wide' or 'long' format")
+        elif shape == "wide":
+            return SQLDataShape.WIDE
+        elif shape == "long":
+            return SQLDataShape.LONG
+        else:
+            raise Exception("Invalid shape: must be either 'long' or 'wide'")
+
+    def sql(
+        self,
+        query: str,
+        shape: Literal["wide", "long"] = "wide",
+        remove_prefix: bool = True,
+        transpose: bool = None,
+        transpose_by: str = None,
+        csv: bool = False,
+        to_list=False,
+        to_latex=False,
+        filename: Optional[str] = None,
+    ) -> Union["pd.DataFrame", str]:
+        """Execute a SQL query and return the results as a DataFrame.
+
+        :param query: The SQL query to execute
+        :param shape: The shape of the data in the database (wide or long)
+        :param remove_prefix: Whether to remove the prefix from the column names
+        :param transpose: Whether to transpose the DataFrame
+        :param transpose_by: The column to use as the index when transposing
+        :param csv: Whether to return the DataFrame as a CSV string
+
+
+        Example usage:
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> d = r.sql("select data_type, key, value from self where data_type = 'answer' order by value limit 3", shape="long")
+        >>> sorted(list(d['value']))
+        ['Good', 'Great', 'Great']
+
+        We can also return the data in wide format.
+        Note the use of single quotes to escape the column names, as required by sql.
+
+        >>> from edsl.results import Results
+        >>> Results.example().sql("select how_feeling from self", shape = 'wide', remove_prefix=True)
+          how_feeling
+        0          OK
+        1       Great
+        2    Terrible
+        3          OK
+        """
+        import pandas as pd
+
+        shape_enum = self._get_shape_enum(shape)
+
+        conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
+        df = pd.read_sql_query(query, conn)
+
+        # 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()
+
+        if csv and to_list:
+            raise Exception("Cannot return both CSV and list")
+
+        if to_list:
+            return df.values.tolist()
+
+        if to_latex:
+            df.columns = [col.replace("_", " ") for col in df.columns]
+
+            latex_output = df.to_latex(index=False)
+            if filename:
+                with open(filename, "w") as f:
+                    f.write(latex_output)
+                return None
+            return latex_output
+
+        if csv:
+            if filename:
+                df.to_csv(filename, index=False)
+                return None
+
+            return df.to_csv(index=False)
+
+        return df
+
+    def show_schema(
+        self, shape: Literal["wide", "long"], remove_prefix: bool = False
+    ) -> None:
+        """Show the schema of the Results database.
+
+        :param shape: The shape of the data in the database (wide or long)
+        :param remove_prefix: Whether to remove the prefix from the column names
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r.show_schema(shape="long")
+        Type: table, Name: self, SQL: CREATE TABLE self (
+        ...
+        <BLANKLINE>
+        """
+        import pandas as pd
+
+        shape_enum = self._get_shape_enum(shape)
+        conn = self._db(shape=shape_enum, remove_prefix=remove_prefix)
+
+        if shape_enum == SQLDataShape.LONG:
+            # Query to get the schema of all tables
+            query = "SELECT type, name, sql FROM sqlite_master WHERE type='table'"
+            cursor = conn.execute(query)
+            schema = cursor.fetchall()
+            conn.close()
+
+            # Format and return the schema information
+            schema_info = ""
+            for row in schema:
+                schema_info += f"Type: {row[0]}, Name: {row[1]}, SQL: {row[2]}\n"
+
+            print(schema_info)
+        elif shape_enum == SQLDataShape.WIDE:
+            query = f"PRAGMA table_info(self)"
+            schema = pd.read_sql(query, conn)
+            # print(schema)
+            return schema
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/results/ResultsExportMixin.py

@@ -1,45 +1,43 @@
-"""Mixin class for exporting results."""
-
-from functools import wraps
-from typing import Literal, Optional, Union
-
-from edsl.results.DatasetExportMixin import DatasetExportMixin
-
-
-def to_dataset(func):
-    """Convert the Results object to a Dataset object before calling the function."""
-
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        """Return the function with the Results object converted to a Dataset object."""
-        if self.__class__.__name__ == "Results":
-            return func(self.select(), *args, **kwargs)
-        elif self.__class__.__name__ == "AgentList":
-            return func(self.to_dataset(), *args, **kwargs)
-        else:
-            return func(self, *args, **kwargs)
-
-    wrapper._is_wrapped = True
-    return wrapper
-
-
-def decorate_methods_from_mixin(cls, mixin_cls):
-    for attr_name, attr_value in mixin_cls.__dict__.items():
-        if callable(attr_value) and not attr_name.startswith("__"):
-            setattr(cls, attr_name, to_dataset(attr_value))
-    return cls
-
-
-class ResultsExportMixin(DatasetExportMixin):
-    """Mixin class for exporting Results objects."""
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        decorate_methods_from_mixin(cls, DatasetExportMixin)
-
-
-if __name__ == "__main__":
-    # pass
-    import doctest
-
-    doctest.testmod(optionflags=doctest.ELLIPSIS)
+"""Mixin class for exporting results."""
+
+from functools import wraps
+from typing import Literal, Optional, Union
+
+from edsl.results.DatasetExportMixin import DatasetExportMixin
+
+
+def to_dataset(func):
+    """Convert the Results object to a Dataset object before calling the function."""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        """Return the function with the Results object converted to a Dataset object."""
+        if self.__class__.__name__ == "Results":
+            return func(self.select(), *args, **kwargs)
+        else:
+            return func(self, *args, **kwargs)
+
+    wrapper._is_wrapped = True
+    return wrapper
+
+
+def decorate_methods_from_mixin(cls, mixin_cls):
+    for attr_name, attr_value in mixin_cls.__dict__.items():
+        if callable(attr_value) and not attr_name.startswith("__"):
+            setattr(cls, attr_name, to_dataset(attr_value))
+    return cls
+
+
+class ResultsExportMixin(DatasetExportMixin):
+    """Mixin class for exporting Results objects."""
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        decorate_methods_from_mixin(cls, DatasetExportMixin)
+
+
+if __name__ == "__main__":
+    # pass
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/results/ResultsFetchMixin.py

@@ -1,33 +1,33 @@
-"""Mixin for fetching data from results."""
-
-from functools import partial
-from itertools import chain
-
-
-class ResultsFetchMixin:
-    """Mixin for fetching data from results."""
-
-    def _fetch_list(self, data_type: str, key: str) -> list:
-        """
-        Return a list of values from the data for a given data type and key.
-
-        Uses the filtered data, not the original data.
-
-        Example:
-
-        >>> from edsl.results import Results
-        >>> r = Results.example()
-        >>> r._fetch_list('answer', 'how_feeling')
-        ['OK', 'Great', 'Terrible', 'OK']
-        """
-        returned_list = []
-        for row in self.data:
-            returned_list.append(row.sub_dicts[data_type].get(key, None))
-
-        return returned_list
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod(optionflags=doctest.ELLIPSIS)
+"""Mixin for fetching data from results."""
+
+from functools import partial
+from itertools import chain
+
+
+class ResultsFetchMixin:
+    """Mixin for fetching data from results."""
+
+    def _fetch_list(self, data_type: str, key: str) -> list:
+        """
+        Return a list of values from the data for a given data type and key.
+
+        Uses the filtered data, not the original data.
+
+        Example:
+
+        >>> from edsl.results import Results
+        >>> r = Results.example()
+        >>> r._fetch_list('answer', 'how_feeling')
+        ['OK', 'Great', 'Terrible', 'OK']
+        """
+        returned_list = []
+        for row in self.data:
+            returned_list.append(row.sub_dicts[data_type].get(key, None))
+
+        return returned_list
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/results/ResultsGGMixin.py

@@ -1,121 +1,121 @@
-"""Mixin class for ggplot2 plotting."""
-
-import subprocess
-import tempfile
-from typing import Optional
-
-
-class ResultsGGMixin:
-    """Mixin class for ggplot2 plotting."""
-
-    def ggplot2(
-        self,
-        ggplot_code: str,
-        filename: str = None,
-        shape="wide",
-        sql: str = None,
-        remove_prefix: bool = True,
-        debug: bool = False,
-        height=4,
-        width=6,
-        format="svg",
-        factor_orders: Optional[dict] = None,
-    ):
-        """Create a ggplot2 plot from a DataFrame.
-
-        :param ggplot_code: The ggplot2 code to execute.
-        :param filename: The filename to save the plot to.
-        :param shape: The shape of the data in the DataFrame (wide or long).
-        :param sql: The SQL query to execute beforehand to manipulate the data.
-        :param remove_prefix: Whether to remove the prefix from the column names.
-        :param debug: Whether to print the R code instead of executing it.
-        :param height: The height of the plot in inches.
-        :param width: The width of the plot in inches.
-        :param format: The format to save the plot in (png or svg).
-        :param factor_orders: A dictionary of factor columns and their order.
-        """
-
-        if sql == None:
-            sql = "select * from self"
-
-        if shape == "long":
-            df = self.sql(sql, shape="long")
-        elif shape == "wide":
-            df = self.sql(sql, shape="wide", remove_prefix=remove_prefix)
-
-        # Convert DataFrame to CSV format
-        csv_data = df.to_csv(index=False)
-
-        # Embed the CSV data within the R script
-        csv_data_escaped = csv_data.replace("\n", "\\n").replace("'", "\\'")
-        read_csv_code = f"self <- read.csv(text = '{csv_data_escaped}', sep = ',')\n"
-
-        if factor_orders is not None:
-            for factor, order in factor_orders.items():
-                # read_csv_code += f"""self${{{factor}}} <- factor(self${{{factor}}}, levels=c({','.join(['"{}"'.format(x) for x in order])}))"""
-
-                level_string = ", ".join([f'"{x}"' for x in order])
-                read_csv_code += (
-                    f"self${factor} <- factor(self${factor}, levels=c({level_string}))"
-                )
-                read_csv_code += "\n"
-
-        # Load ggplot2 library
-        load_ggplot2 = "library(ggplot2)\n"
-
-        # Check if a filename is provided for the plot, if not create a temporary one
-        if not filename:
-            filename = tempfile.mktemp(suffix=f".{format}")
-
-        # Combine all R script parts
-        full_r_code = load_ggplot2 + read_csv_code + ggplot_code
-
-        # Add command to save the plot to a file
-        full_r_code += f'\nggsave("{filename}", plot = last_plot(), width = {width}, height = {height}, device = "{format}")'
-
-        if debug:
-            print(full_r_code)
-            return
-
-        result = subprocess.run(
-            ["Rscript", "-"],
-            input=full_r_code,
-            text=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-
-        if result.returncode != 0:
-            if result.returncode == 127:  # 'command not found'
-                raise RuntimeError(
-                    "Rscript is probably not installed. Please install R from https://cran.r-project.org/"
-                )
-            else:
-                raise RuntimeError(
-                    f"An error occurred while running Rscript: {result.stderr}"
-                )
-
-        if result.stderr:
-            print("Error in R script:", result.stderr)
-        else:
-            self._display_plot(filename, width, height)
-
-    def _display_plot(self, filename: str, width: float, height: float):
-        """Display the plot in the notebook."""
-        import matplotlib.pyplot as plt
-        import matplotlib.image as mpimg
-
-        if filename.endswith(".png"):
-            img = mpimg.imread(filename)
-            plt.figure(
-                figsize=(width, height)
-            )  # Set the figure size (width, height) in inches
-            plt.imshow(img)
-            plt.axis("off")
-            plt.show()
-        elif filename.endswith(".svg"):
-            from IPython.display import SVG, display
-
-            display(SVG(filename=filename))
-        else:
-            print("Unsupported file format. Please provide a PNG or SVG file.")
+"""Mixin class for ggplot2 plotting."""
+
+import subprocess
+import tempfile
+from typing import Optional
+
+
+class ResultsGGMixin:
+    """Mixin class for ggplot2 plotting."""
+
+    def ggplot2(
+        self,
+        ggplot_code: str,
+        filename: str = None,
+        shape="wide",
+        sql: str = None,
+        remove_prefix: bool = True,
+        debug: bool = False,
+        height=4,
+        width=6,
+        format="svg",
+        factor_orders: Optional[dict] = None,
+    ):
+        """Create a ggplot2 plot from a DataFrame.
+
+        :param ggplot_code: The ggplot2 code to execute.
+        :param filename: The filename to save the plot to.
+        :param shape: The shape of the data in the DataFrame (wide or long).
+        :param sql: The SQL query to execute beforehand to manipulate the data.
+        :param remove_prefix: Whether to remove the prefix from the column names.
+        :param debug: Whether to print the R code instead of executing it.
+        :param height: The height of the plot in inches.
+        :param width: The width of the plot in inches.
+        :param format: The format to save the plot in (png or svg).
+        :param factor_orders: A dictionary of factor columns and their order.
+        """
+
+        if sql == None:
+            sql = "select * from self"
+
+        if shape == "long":
+            df = self.sql(sql, shape="long")
+        elif shape == "wide":
+            df = self.sql(sql, shape="wide", remove_prefix=remove_prefix)
+
+        # Convert DataFrame to CSV format
+        csv_data = df.to_csv(index=False)
+
+        # Embed the CSV data within the R script
+        csv_data_escaped = csv_data.replace("\n", "\\n").replace("'", "\\'")
+        read_csv_code = f"self <- read.csv(text = '{csv_data_escaped}', sep = ',')\n"
+
+        if factor_orders is not None:
+            for factor, order in factor_orders.items():
+                # read_csv_code += f"""self${{{factor}}} <- factor(self${{{factor}}}, levels=c({','.join(['"{}"'.format(x) for x in order])}))"""
+
+                level_string = ", ".join([f'"{x}"' for x in order])
+                read_csv_code += (
+                    f"self${factor} <- factor(self${factor}, levels=c({level_string}))"
+                )
+                read_csv_code += "\n"
+
+        # Load ggplot2 library
+        load_ggplot2 = "library(ggplot2)\n"
+
+        # Check if a filename is provided for the plot, if not create a temporary one
+        if not filename:
+            filename = tempfile.mktemp(suffix=f".{format}")
+
+        # Combine all R script parts
+        full_r_code = load_ggplot2 + read_csv_code + ggplot_code
+
+        # Add command to save the plot to a file
+        full_r_code += f'\nggsave("{filename}", plot = last_plot(), width = {width}, height = {height}, device = "{format}")'
+
+        if debug:
+            print(full_r_code)
+            return
+
+        result = subprocess.run(
+            ["Rscript", "-"],
+            input=full_r_code,
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        if result.returncode != 0:
+            if result.returncode == 127:  # 'command not found'
+                raise RuntimeError(
+                    "Rscript is probably not installed. Please install R from https://cran.r-project.org/"
+                )
+            else:
+                raise RuntimeError(
+                    f"An error occurred while running Rscript: {result.stderr}"
+                )
+
+        if result.stderr:
+            print("Error in R script:", result.stderr)
+        else:
+            self._display_plot(filename, width, height)
+
+    def _display_plot(self, filename: str, width: float, height: float):
+        """Display the plot in the notebook."""
+        import matplotlib.pyplot as plt
+        import matplotlib.image as mpimg
+
+        if filename.endswith(".png"):
+            img = mpimg.imread(filename)
+            plt.figure(
+                figsize=(width, height)
+            )  # Set the figure size (width, height) in inches
+            plt.imshow(img)
+            plt.axis("off")
+            plt.show()
+        elif filename.endswith(".svg"):
+            from IPython.display import SVG, display
+
+            display(SVG(filename=filename))
+        else:
+            print("Unsupported file format. Please provide a PNG or SVG file.")