edsl 0.1.38.dev4__py3-none-any.whl → 0.1.39__py3-none-any.whl

edsl/scenarios/handlers/sql.py

@@ -0,0 +1,313 @@
+from edsl.scenarios.file_methods import FileMethods
+import tempfile
+import re
+from typing import List, Optional
+import textwrap
+
+
+class SqlMethods(FileMethods):
+    suffix = "sql"
+
+    def view_system(self):
+        import os
+        import subprocess
+
+        if os.path.exists(self.path):
+            try:
+                if (os_name := os.name) == "posix":
+                    subprocess.run(["open", self.path], check=True)  # macOS
+                elif os_name == "nt":
+                    os.startfile(self.path)  # Windows
+                else:
+                    subprocess.run(["xdg-open", self.path], check=True)  # Linux
+            except Exception as e:
+                print(f"Error opening SQL file: {e}")
+        else:
+            print("SQL file was not found.")
+
+    def view_notebook(self):
+        from IPython.display import FileLink, display, HTML
+        import pygments
+        from pygments.lexers import SqlLexer
+        from pygments.formatters import HtmlFormatter
+
+        try:
+            with open(self.path, "r", encoding="utf-8") as f:
+                content = f.read()
+
+            formatter = HtmlFormatter(style="monokai")
+            highlighted_sql = pygments.highlight(content, SqlLexer(), formatter)
+            css = formatter.get_style_defs(".highlight")
+            display(HTML(f"<style>{css}</style>{highlighted_sql}"))
+            display(FileLink(self.path))
+        except Exception as e:
+            print(f"Error displaying SQL: {e}")
+
+    def _format_keywords(self, sql: str) -> str:
+        """Capitalize SQL keywords."""
+        keywords = {
+            "select",
+            "from",
+            "where",
+            "and",
+            "or",
+            "insert",
+            "update",
+            "delete",
+            "create",
+            "drop",
+            "alter",
+            "table",
+            "into",
+            "values",
+            "group",
+            "by",
+            "having",
+            "order",
+            "limit",
+            "join",
+            "left",
+            "right",
+            "inner",
+            "outer",
+            "on",
+            "as",
+            "distinct",
+            "count",
+            "sum",
+            "avg",
+            "max",
+            "min",
+            "between",
+            "like",
+            "in",
+            "is",
+            "null",
+            "not",
+            "case",
+            "when",
+            "then",
+            "else",
+            "end",
+        }
+
+        words = sql.split()
+        formatted_words = []
+        for word in words:
+            lower_word = word.lower()
+            if lower_word in keywords:
+                formatted_words.append(word.upper())
+            else:
+                formatted_words.append(word.lower())
+        return " ".join(formatted_words)
+
+    def _indent_sql(self, sql: str) -> str:
+        """Add basic indentation to SQL statement."""
+        lines = sql.split("\n")
+        indented_lines = []
+        indent_level = 0
+
+        for line in lines:
+            line = line.strip()
+
+            # Decrease indent for closing parentheses
+            if line.startswith(")"):
+                indent_level = max(0, indent_level - 1)
+
+            # Add indentation
+            if line:
+                indented_lines.append("    " * indent_level + line)
+            else:
+                indented_lines.append("")
+
+            # Increase indent after opening parentheses
+            if line.endswith("("):
+                indent_level += 1
+
+            # Special cases for common SQL clauses
+            lower_line = line.lower()
+            if any(
+                clause in lower_line
+                for clause in [
+                    "select",
+                    "from",
+                    "where",
+                    "group by",
+                    "having",
+                    "order by",
+                ]
+            ):
+                indent_level = 1
+
+        return "\n".join(indented_lines)
+
+    def format_sql(self) -> bool:
+        """Format the SQL file with proper indentation and keyword capitalization."""
+        try:
+            with open(self.path, "r", encoding="utf-8") as f:
+                content = f.read()
+
+            # Remove extra whitespace and format
+            content = " ".join(content.split())
+            content = self._format_keywords(content)
+            content = self._indent_sql(content)
+
+            # Wrap long lines
+            wrapped_content = []
+            for line in content.split("\n"):
+                if len(line) > 80:
+                    wrapped_line = textwrap.fill(
+                        line, width=80, subsequent_indent="    "
+                    )
+                    wrapped_content.append(wrapped_line)
+                else:
+                    wrapped_content.append(line)
+
+            formatted_sql = "\n".join(wrapped_content)
+
+            with open(self.path, "w", encoding="utf-8") as f:
+                f.write(formatted_sql)
+
+            return True
+        except Exception as e:
+            print(f"Error formatting SQL: {e}")
+            return False
+
+    def split_statements(self) -> List[str]:
+        """Split the SQL file into individual statements."""
+        try:
+            with open(self.path, "r", encoding="utf-8") as f:
+                content = f.read()
+
+            # Handle both semicolon and GO statement terminators
+            statements = []
+            current_stmt = []
+
+            for line in content.split("\n"):
+                line = line.strip()
+
+                # Skip empty lines and comments
+                if not line or line.startswith("--"):
+                    continue
+
+                if line.endswith(";"):
+                    current_stmt.append(line[:-1])  # Remove semicolon
+                    statements.append(" ".join(current_stmt))
+                    current_stmt = []
+                elif line.upper() == "GO":
+                    if current_stmt:
+                        statements.append(" ".join(current_stmt))
+                        current_stmt = []
+                else:
+                    current_stmt.append(line)
+
+            # Add any remaining statement
+            if current_stmt:
+                statements.append(" ".join(current_stmt))
+
+            return [stmt.strip() for stmt in statements if stmt.strip()]
+        except Exception as e:
+            print(f"Error splitting SQL statements: {e}")
+            return []
+
+    def validate_basic_syntax(self) -> bool:
+        """
+        Perform basic SQL syntax validation.
+        This is a simple check and doesn't replace proper SQL parsing.
+        """
+        try:
+            with open(self.path, "r", encoding="utf-8") as f:
+                content = f.read()
+
+            statements = self.split_statements()
+            for stmt in statements:
+                # Check for basic SQL keywords
+                stmt_upper = stmt.upper()
+                if not any(
+                    keyword in stmt_upper
+                    for keyword in [
+                        "SELECT",
+                        "INSERT",
+                        "UPDATE",
+                        "DELETE",
+                        "CREATE",
+                        "DROP",
+                        "ALTER",
+                    ]
+                ):
+                    print(f"Warning: Statement might be incomplete: {stmt}")
+
+                # Check for basic parentheses matching
+                if stmt.count("(") != stmt.count(")"):
+                    print(f"Error: Unmatched parentheses in statement: {stmt}")
+                    return False
+
+                # Check for basic quote matching
+                if stmt.count("'") % 2 != 0:
+                    print(f"Error: Unmatched quotes in statement: {stmt}")
+                    return False
+
+            return True
+        except Exception as e:
+            print(f"Error validating SQL: {e}")
+            return False
+
+    def extract_table_names(self) -> List[str]:
+        """Extract table names from the SQL file."""
+        tables = set()
+        try:
+            with open(self.path, "r", encoding="utf-8") as f:
+                content = f.read()
+
+            patterns = [
+                r"FROM\s+([a-zA-Z_][a-zA-Z0-9_]*)",
+                r"JOIN\s+([a-zA-Z_][a-zA-Z0-9_]*)",
+                r"UPDATE\s+([a-zA-Z_][a-zA-Z0-9_]*)",
+                r"INSERT\s+INTO\s+([a-zA-Z_][a-zA-Z0-9_]*)",
+                r"CREATE\s+TABLE\s+([a-zA-Z_][a-zA-Z0-9_]*)",
+            ]
+
+            for pattern in patterns:
+                tables.update(re.findall(pattern, content, re.IGNORECASE))
+
+            return sorted(list(tables))
+        except Exception as e:
+            print(f"Error extracting table names: {e}")
+            return []
+
+    def example(self):
+        sample_sql = """-- Sample SQL file with common operations
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name VARCHAR(100) NOT NULL,
+    department VARCHAR(50),
+    salary DECIMAL(10,2),
+    hire_date DATE
+);
+
+INSERT INTO employees (name, department, salary, hire_date)
+VALUES 
+    ('John Doe', 'Engineering', 75000.00, '2023-01-15'),
+    ('Jane Smith', 'Marketing', 65000.00, '2023-02-01');
+
+-- Query to analyze employee data
+SELECT 
+    department,
+    COUNT(*) as employee_count,
+    AVG(salary) as avg_salary
+FROM employees
+GROUP BY department
+HAVING COUNT(*) > 0
+ORDER BY avg_salary DESC;
+
+-- Update salary with conditions
+UPDATE employees
+SET salary = salary * 1.1
+WHERE department = 'Engineering'
+    AND hire_date < '2024-01-01';
+"""
+        with tempfile.NamedTemporaryFile(
+            delete=False, suffix=".sql", mode="w", encoding="utf-8"
+        ) as f:
+            f.write(sample_sql)
+        return f.name

edsl/scenarios/handlers/sqlite.py

@@ -0,0 +1,149 @@
+from edsl.scenarios.file_methods import FileMethods
+import os
+import tempfile
+import sqlite3
+
+
+class SQLiteMethods(FileMethods):
+    suffix = "db"  # or "sqlite", depending on your preference
+
+    def extract_text(self):
+        """
+        Extracts a text representation of the database schema and table contents.
+        """
+        with sqlite3.connect(self.path) as conn:
+            cursor = conn.cursor()
+
+            # Get all table names
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+            tables = cursor.fetchall()
+
+            full_text = []
+
+            # For each table, get schema and contents
+            for (table_name,) in tables:
+                # Get table schema
+                cursor.execute(
+                    f"SELECT sql FROM sqlite_master WHERE type='table' AND name='{table_name}';"
+                )
+                schema = cursor.fetchone()[0]
+                full_text.append(f"Table: {table_name}")
+                full_text.append(f"Schema: {schema}")
+
+                # Get table contents
+                cursor.execute(f"SELECT * FROM {table_name};")
+                rows = cursor.fetchall()
+
+                # Get column names
+                column_names = [description[0] for description in cursor.description]
+                full_text.append(f"Columns: {', '.join(column_names)}")
+
+                # Add row data
+                for row in rows:
+                    full_text.append(str(row))
+                full_text.append("\n")
+
+        return "\n".join(full_text)
+
+    def view_system(self):
+        """
+        Opens the database with the system's default SQLite viewer if available.
+        """
+        import os
+        import subprocess
+
+        if os.path.exists(self.path):
+            try:
+                if (os_name := os.name) == "posix":
+                    # Try DB Browser for SQLite on macOS
+                    subprocess.run(
+                        ["open", "-a", "DB Browser for SQLite", self.path], check=True
+                    )
+                elif os_name == "nt":
+                    # Try DB Browser for SQLite on Windows
+                    subprocess.run(["DB Browser for SQLite.exe", self.path], check=True)
+                else:
+                    # Try sqlitebrowser on Linux
+                    subprocess.run(["sqlitebrowser", self.path], check=True)
+            except Exception as e:
+                print(f"Error opening SQLite database: {e}")
+        else:
+            print("SQLite database file was not found.")
+
+    def view_notebook(self):
+        """
+        Displays database contents in a Jupyter notebook.
+        """
+        import pandas as pd
+        from IPython.display import HTML, display
+
+        with sqlite3.connect(self.path) as conn:
+            # Get all table names
+            cursor = conn.cursor()
+            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+            tables = cursor.fetchall()
+
+            html_parts = []
+            for (table_name,) in tables:
+                # Read table into pandas DataFrame
+                df = pd.read_sql_query(f"SELECT * FROM {table_name}", conn)
+
+                # Convert to HTML with styling
+                table_html = f"""
+                <div style="margin-bottom: 20px;">
+                    <h3>{table_name}</h3>
+                    {df.to_html(index=False)}
+                </div>
+                """
+                html_parts.append(table_html)
+
+            # Combine all tables into one scrollable div
+            html = f"""
+            <div style="width: 800px; height: 800px; padding: 20px; 
+                       border: 1px solid #ccc; overflow-y: auto;">
+                {''.join(html_parts)}
+            </div>
+            """
+            display(HTML(html))
+
+    def example(self):
+        """
+        Creates an example SQLite database for testing.
+        """
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as tmp:
+            conn = sqlite3.connect(tmp.name)
+            cursor = conn.cursor()
+
+            # Create a sample table
+            cursor.execute(
+                """
+                CREATE TABLE survey_responses (
+                    id INTEGER PRIMARY KEY,
+                    question TEXT,
+                    response TEXT
+                )
+            """
+            )
+
+            # Insert some sample data
+            sample_data = [
+                (1, "First Survey Question", "Response 1"),
+                (2, "Second Survey Question", "Response 2"),
+            ]
+            cursor.executemany(
+                "INSERT INTO survey_responses (id, question, response) VALUES (?, ?, ?)",
+                sample_data,
+            )
+
+            conn.commit()
+            conn.close()
+            tmp.close()
+
+            return tmp.name
+
+
+if __name__ == "__main__":
+    sqlite_temp = SQLiteMethods.example()
+    from edsl.scenarios.FileStore import FileStore
+
+    fs = FileStore(sqlite_temp)

edsl/scenarios/handlers/txt.py

@@ -0,0 +1,33 @@
+from edsl.scenarios.file_methods import FileMethods
+import tempfile
+
+
+class TxtMethods(FileMethods):
+    suffix = "txt"
+
+    def view_system(self):
+        import os
+        import subprocess
+
+        if os.path.exists(self.path):
+            try:
+                if (os_name := os.name) == "posix":
+                    subprocess.run(["open", self.path], check=True)  # macOS
+                elif os_name == "nt":
+                    os.startfile(self.path)  # Windows
+                else:
+                    subprocess.run(["xdg-open", self.path], check=True)  # Linux
+            except Exception as e:
+                print(f"Error opening TXT: {e}")
+        else:
+            print("TXT file was not found.")
+
+    def view_notebook(self):
+        from IPython.display import FileLink, display
+
+        display(FileLink(self.path))
+
+    def example(self):
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as f:
+            f.write(b"Hello, World!")
+        return f.name

edsl/scenarios/{ScenarioJoin.py → scenario_join.py}

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import Union, TYPE_CHECKING
 
-# if TYPE_CHECKING:
-from edsl.scenarios.ScenarioList import ScenarioList
-from edsl.scenarios.Scenario import Scenario
+if TYPE_CHECKING:
+    from edsl.scenarios.ScenarioList import ScenarioList
+    from edsl.scenarios.Scenario import Scenario
 
 
 class ScenarioJoin:
@@ -23,7 +23,7 @@ class ScenarioJoin:
         self.left = left
         self.right = right
 
-    def left_join(self, by: Union[str, list[str]]) -> ScenarioList:
+    def left_join(self, by: Union[str, list[str]]) -> "ScenarioList":
         """Perform a left join between the two ScenarioLists.
 
         Args:
@@ -35,6 +35,8 @@ class ScenarioJoin:
         Raises:
             ValueError: If by is empty or if any join keys don't exist in both ScenarioLists
         """
+        from edsl.scenarios.ScenarioList import ScenarioList
+
         self._validate_join_keys(by)
         by_keys = [by] if isinstance(by, str) else by
 
@@ -86,6 +88,8 @@ class ScenarioJoin:
         self, by_keys: list[str], other_dict: dict, all_keys: set
     ) -> list[Scenario]:
         """Create the joined scenarios."""
+        from edsl.scenarios.Scenario import Scenario
+
         new_scenarios = []
 
         for scenario in self.left:
@@ -105,8 +109,8 @@ class ScenarioJoin:
     def _handle_matching_scenario(
         self,
         new_scenario: dict,
-        left_scenario: Scenario,
-        right_scenario: Scenario,
+        left_scenario: "Scenario",
+        right_scenario: "Scenario",
         by_keys: list[str],
     ) -> None:
         """Handle merging of matching scenarios and conflict warnings."""

edsl/scenarios/scenario_selector.py

@@ -0,0 +1,156 @@
+from typing import TYPE_CHECKING
+
+
+class ScenarioSelector:
+    """
+    A class for performing advanced field selection on ScenarioList objects,
+    including support for wildcard patterns.
+
+    Args:
+        scenario_list: The ScenarioList object to perform selections on
+
+    Examples:
+        >>> from edsl import Scenario, ScenarioList
+        >>> scenarios = ScenarioList([Scenario({'test_1': 1, 'test_2': 2, 'other': 3}), Scenario({'test_1': 4, 'test_2': 5, 'other': 6})])
+        >>> selector = ScenarioSelector(scenarios)
+        >>> selector.select('test*')
+        ScenarioList([Scenario({'test_1': 1, 'test_2': 2}), Scenario({'test_1': 4, 'test_2': 5})])
+    """
+
+    def __init__(self, scenario_list: "ScenarioList"):
+        """Initialize with a ScenarioList object."""
+        self.scenario_list = scenario_list
+        self.available_fields = (
+            list(scenario_list.data[0].keys()) if scenario_list.data else []
+        )
+
+    def _match_field_pattern(self, pattern: str, field: str) -> bool:
+        """
+        Checks if a field name matches a pattern with wildcards.
+        Supports '*' as wildcard at start or end of pattern.
+
+        Args:
+            pattern: The pattern to match against, may contain '*' at start or end
+            field: The field name to check
+
+        Examples:
+            >>> from edsl.scenarios import ScenarioList, Scenario
+            >>> selector = ScenarioSelector(ScenarioList([]))
+            >>> selector._match_field_pattern('test*', 'test_field')
+            True
+            >>> selector._match_field_pattern('*field', 'test_field')
+            True
+            >>> selector._match_field_pattern('test', 'test')
+            True
+            >>> selector._match_field_pattern('*test*', 'my_test_field')
+            True
+        """
+        if "*" not in pattern:
+            return pattern == field
+
+        if pattern.startswith("*") and pattern.endswith("*"):
+            return pattern[1:-1] in field
+        elif pattern.startswith("*"):
+            return field.endswith(pattern[1:])
+        elif pattern.endswith("*"):
+            return field.startswith(pattern[:-1])
+        return pattern == field
+
+    def _get_matching_fields(self, patterns: list[str]) -> list[str]:
+        """
+        Gets all fields that match any of the given patterns.
+
+        Args:
+            patterns: List of field patterns, may contain wildcards
+
+        Returns:
+            List of field names that match at least one pattern
+
+        Examples:
+            >>> from edsl import Scenario, ScenarioList
+            >>> scenarios = ScenarioList([
+            ...     Scenario({'test_1': 1, 'test_2': 2, 'other': 3})
+            ... ])
+            >>> selector = ScenarioSelector(scenarios)
+            >>> selector._get_matching_fields(['test*'])
+            ['test_1', 'test_2']
+        """
+        matching_fields = set()
+        for pattern in patterns:
+            matches = [
+                field
+                for field in self.available_fields
+                if self._match_field_pattern(pattern, field)
+            ]
+            matching_fields.update(matches)
+        return sorted(list(matching_fields))
+
+    def select(self, *fields) -> "ScenarioList":
+        """
+        Selects scenarios with only the referenced fields.
+        Supports wildcard patterns using '*' at the start or end of field names.
+
+        Args:
+            *fields: Field names or patterns to select. Patterns may include '*' for wildcards.
+
+        Returns:
+            A new ScenarioList containing only the matched fields.
+
+        Raises:
+            ValueError: If no fields match the given patterns.
+
+        Examples:
+            >>> from edsl import Scenario, ScenarioList
+            >>> scenarios = ScenarioList([
+            ...     Scenario({'test_1': 1, 'test_2': 2, 'other': 3}),
+            ...     Scenario({'test_1': 4, 'test_2': 5, 'other': 6})
+            ... ])
+            >>> selector = ScenarioSelector(scenarios)
+            >>> selector.select('test*')  # Selects all fields starting with 'test'
+            ScenarioList([Scenario({'test_1': 1, 'test_2': 2}), Scenario({'test_1': 4, 'test_2': 5})])
+            >>> selector.select('*_1')  # Selects all fields ending with '_1'
+            ScenarioList([Scenario({'test_1': 1}), Scenario({'test_1': 4})])
+            >>> selector.select('test_1', '*_2')  # Multiple patterns
+            ScenarioList([Scenario({'test_1': 1, 'test_2': 2}), Scenario({'test_1': 4, 'test_2': 5})])
+        """
+        if not self.scenario_list.data:
+            return self.scenario_list.__class__([])
+
+        # Convert single string to list for consistent processing
+        patterns = list(fields)
+
+        # Get all fields that match the patterns
+        fields_to_select = self._get_matching_fields(patterns)
+
+        # If no fields match, raise an informative error
+        if not fields_to_select:
+            raise ValueError(
+                f"No fields matched the given patterns: {patterns}. "
+                f"Available fields are: {self.available_fields}"
+            )
+
+        return self.scenario_list.__class__(
+            [scenario.select(fields_to_select) for scenario in self.scenario_list.data]
+        )
+
+    def get_available_fields(self) -> list[str]:
+        """
+        Returns a list of all available fields in the ScenarioList.
+
+        Returns:
+            List of field names available for selection.
+
+        Examples:
+            >>> from edsl import Scenario, ScenarioList
+            >>> scenarios = ScenarioList([Scenario({'test_1': 1, 'test_2': 2, 'other': 3})])
+            >>> selector = ScenarioSelector(scenarios)
+            >>> selector.get_available_fields()
+            ['other', 'test_1', 'test_2']
+        """
+        return sorted(self.available_fields)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)