sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/core/result.py

@@ -10,17 +10,20 @@ Classes:
 """
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, cast, overload
 
 from mypy_extensions import mypyc_attr
 from typing_extensions import TypeVar
 
 from sqlspec.core.compiler import OperationType
+from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow
+from sqlspec.utils.schema import to_schema
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     from sqlspec.core.statement import SQL
+    from sqlspec.typing import ArrowTable, PandasDataFrame, PolarsDataFrame, SchemaT
 
 
 __all__ = ("ArrowResult", "SQLResult", "StatementResult")
@@ -52,8 +55,8 @@ class StatementResult(ABC):
         statement: "SQL",
         data: Any = None,
         rows_affected: int = 0,
-        last_inserted_id: Optional[Union[int, str]] = None,
-        execution_time: Optional[float] = None,
+        last_inserted_id: int | str | None = None,
+        execution_time: float | None = None,
         metadata: Optional["dict[str, Any]"] = None,
     ) -> None:
         """Initialize statement result.
@@ -129,19 +132,6 @@ class SQLResult(StatementResult):
     The operation_type attribute indicates the nature of the operation.
 
     For script execution, tracks multiple statement results and errors.
-
-    Attributes:
-        column_names: Names of columns in the result set.
-        error: Exception that occurred during execution.
-        errors: List of error messages for script execution.
-        has_more: Whether there are additional result pages available.
-        inserted_ids: List of IDs from INSERT operations.
-        operation_index: Index of operation in a script.
-        parameters: Parameters used for the query.
-        statement_results: Results from individual statements in a script.
-        successful_statements: Count of successful statements in a script.
-        total_count: Total number of rows in the complete result set.
-        total_statements: Total number of statements in a script.
     """
 
     __slots__ = (
@@ -164,19 +154,19 @@ class SQLResult(StatementResult):
     def __init__(
         self,
         statement: "SQL",
-        data: Optional[list[dict[str, Any]]] = None,
+        data: list[dict[str, Any]] | None = None,
         rows_affected: int = 0,
-        last_inserted_id: Optional[Union[int, str]] = None,
-        execution_time: Optional[float] = None,
+        last_inserted_id: int | str | None = None,
+        execution_time: float | None = None,
         metadata: Optional["dict[str, Any]"] = None,
-        error: Optional[Exception] = None,
+        error: Exception | None = None,
         operation_type: OperationType = "SELECT",
-        operation_index: Optional[int] = None,
-        parameters: Optional[Any] = None,
+        operation_index: int | None = None,
+        parameters: Any | None = None,
         column_names: Optional["list[str]"] = None,
-        total_count: Optional[int] = None,
+        total_count: int | None = None,
         has_more: bool = False,
-        inserted_ids: Optional["list[Union[int, str]]"] = None,
+        inserted_ids: Optional["list[int | str]"] = None,
         statement_results: Optional["list[SQLResult]"] = None,
         errors: Optional["list[str]"] = None,
         total_statements: int = 0,
@@ -280,15 +270,25 @@ class SQLResult(StatementResult):
 
         return False
 
-    def get_data(self) -> "list[dict[str,Any]]":
+    @overload
+    def get_data(self, *, schema_type: "type[SchemaT]") -> "list[SchemaT]": ...
+
+    @overload
+    def get_data(self, *, schema_type: None = None) -> "list[dict[str, Any]]": ...
+
+    def get_data(self, *, schema_type: "type[SchemaT] | None" = None) -> "list[SchemaT] | list[dict[str, Any]]":
         """Get the data from the result.
 
         For regular operations, returns the list of rows.
         For script operations, returns a summary dictionary containing
         execution statistics and results.
 
+        Args:
+            schema_type: Optional schema type to transform the data into.
+                Supports Pydantic models, dataclasses, msgspec structs, attrs classes, and TypedDict.
+
         Returns:
-            List of result rows or script summary.
+            List of result rows (optionally transformed to schema_type) or script summary.
         """
         op_type_upper = self.operation_type.upper()
         if op_type_upper == "SCRIPT":
@@ -303,7 +303,10 @@ class SQLResult(StatementResult):
                     "total_rows_affected": self.get_total_rows_affected(),
                 }
             ]
-        return self.data or []
+        data = cast("list[dict[str, Any]]", self.data or [])
+        if schema_type:
+            return cast("list[SchemaT]", to_schema(data, schema_type=schema_type))
+        return data
 
     def add_statement_result(self, result: "SQLResult") -> None:
         """Add a statement result to the script execution results.
@@ -348,13 +351,28 @@ class SQLResult(StatementResult):
         """
         return len(self.column_names) if self.column_names else 0
 
-    def get_first(self) -> "Optional[dict[str, Any]]":
+    @overload
+    def get_first(self, *, schema_type: "type[SchemaT]") -> "SchemaT | None": ...
+
+    @overload
+    def get_first(self, *, schema_type: None = None) -> "dict[str, Any] | None": ...
+
+    def get_first(self, *, schema_type: "type[SchemaT] | None" = None) -> "SchemaT | dict[str, Any] | None":
         """Get the first row from the result, if any.
 
+        Args:
+            schema_type: Optional schema type to transform the data into.
+                Supports Pydantic models, dataclasses, msgspec structs, attrs classes, and TypedDict.
+
         Returns:
-            First row or None if no data.
+            First row (optionally transformed to schema_type) or None if no data.
         """
-        return self.data[0] if self.data else None
+        if not self.data:
+            return None
+        row = cast("dict[str, Any]", self.data[0])
+        if schema_type:
+            return to_schema(row, schema_type=schema_type)
+        return row
 
     def get_count(self) -> int:
         """Get the number of rows in the current result set (e.g., a page of data).
@@ -434,19 +452,42 @@ class SQLResult(StatementResult):
         """
         return iter(self.data or [])
 
-    def all(self) -> list[dict[str, Any]]:
+    @overload
+    def all(self, *, schema_type: "type[SchemaT]") -> "list[SchemaT]": ...
+
+    @overload
+    def all(self, *, schema_type: None = None) -> list[dict[str, Any]]: ...
+
+    def all(self, *, schema_type: "type[SchemaT] | None" = None) -> "list[SchemaT] | list[dict[str, Any]]":
         """Return all rows as a list.
 
+        Args:
+            schema_type: Optional schema type to transform the data into.
+                Supports Pydantic models, dataclasses, msgspec structs, attrs classes, and TypedDict.
+
         Returns:
-            List of all rows in the result
+            List of all rows (optionally transformed to schema_type)
         """
-        return self.data or []
+        data = cast("list[dict[str, Any]]", self.data or [])
+        if schema_type:
+            return cast("list[SchemaT]", to_schema(data, schema_type=schema_type))
+        return data
+
+    @overload
+    def one(self, *, schema_type: "type[SchemaT]") -> "SchemaT": ...
+
+    @overload
+    def one(self, *, schema_type: None = None) -> "dict[str, Any]": ...
 
-    def one(self) -> "dict[str, Any]":
+    def one(self, *, schema_type: "type[SchemaT] | None" = None) -> "SchemaT | dict[str, Any]":
         """Return exactly one row.
 
+        Args:
+            schema_type: Optional schema type to transform the data into.
+                Supports Pydantic models, dataclasses, msgspec structs, attrs classes, and TypedDict.
+
         Returns:
-            The single row
+            The single row (optionally transformed to schema_type)
 
         Raises:
             ValueError: If no results or more than one result
@@ -463,13 +504,26 @@ class SQLResult(StatementResult):
             msg = f"Multiple results found ({data_len}), exactly one row expected"
             raise ValueError(msg)
 
-        return cast("dict[str, Any]", self.data[0])
+        row = cast("dict[str, Any]", self.data[0])
+        if schema_type:
+            return to_schema(row, schema_type=schema_type)
+        return row
+
+    @overload
+    def one_or_none(self, *, schema_type: "type[SchemaT]") -> "SchemaT | None": ...
 
-    def one_or_none(self) -> "Optional[dict[str, Any]]":
+    @overload
+    def one_or_none(self, *, schema_type: None = None) -> "dict[str, Any] | None": ...
+
+    def one_or_none(self, *, schema_type: "type[SchemaT] | None" = None) -> "SchemaT | dict[str, Any] | None":
         """Return at most one row.
 
+        Args:
+            schema_type: Optional schema type to transform the data into.
+                Supports Pydantic models, dataclasses, msgspec structs, attrs classes, and TypedDict.
+
         Returns:
-            The single row or None if no results
+            The single row (optionally transformed to schema_type) or None if no results
 
         Raises:
             ValueError: If more than one result
@@ -484,7 +538,10 @@ class SQLResult(StatementResult):
             msg = f"Multiple results found ({data_len}), at most one row expected"
             raise ValueError(msg)
 
-        return cast("dict[str, Any]", self.data[0])
+        row = cast("dict[str, Any]", self.data[0])
+        if schema_type:
+            return to_schema(row, schema_type=schema_type)
+        return row
 
     def scalar(self) -> Any:
         """Return the first column of the first row.
@@ -527,8 +584,8 @@ class ArrowResult(StatementResult):
         statement: "SQL",
         data: Any,
         rows_affected: int = 0,
-        last_inserted_id: Optional[Union[int, str]] = None,
-        execution_time: Optional[float] = None,
+        last_inserted_id: int | str | None = None,
+        execution_time: float | None = None,
         metadata: Optional["dict[str, Any]"] = None,
         schema: Optional["dict[str, Any]"] = None,
     ) -> None:
@@ -562,7 +619,7 @@ class ArrowResult(StatementResult):
         """
         return self.data is not None
 
-    def get_data(self) -> Any:
+    def get_data(self) -> "ArrowTable":
         """Get the Apache Arrow Table from the result.
 
         Returns:
@@ -570,10 +627,19 @@ class ArrowResult(StatementResult):
 
         Raises:
             ValueError: If no Arrow table is available.
+            TypeError: If data is not an Arrow Table.
         """
         if self.data is None:
             msg = "No Arrow table available for this result"
             raise ValueError(msg)
+
+        ensure_pyarrow()
+
+        import pyarrow as pa
+
+        if not isinstance(self.data, pa.Table):
+            msg = f"Expected an Arrow Table, but got {type(self.data).__name__}"
+            raise TypeError(msg)
         return self.data
 
     @property
@@ -624,13 +690,134 @@ class ArrowResult(StatementResult):
 
         return cast("int", self.data.num_columns)
 
+    def to_pandas(self) -> "PandasDataFrame":
+        """Convert Arrow data to pandas DataFrame.
+
+        Returns:
+            pandas DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_pandas()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_pandas()
+
+        import pandas as pd
+
+        result = self.data.to_pandas()
+        if not isinstance(result, pd.DataFrame):
+            msg = f"Expected a pandas DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_polars(self) -> "PolarsDataFrame":
+        """Convert Arrow data to Polars DataFrame.
+
+        Returns:
+            Polars DataFrame containing the result data.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> df = result.to_polars()
+            >>> print(df.head())
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        ensure_polars()
+
+        import polars as pl
+
+        result = pl.from_arrow(self.data)
+        if not isinstance(result, pl.DataFrame):
+            msg = f"Expected a Polars DataFrame, but got {type(result).__name__}"
+            raise TypeError(msg)
+        return result
+
+    def to_dict(self) -> "list[dict[str, Any]]":
+        """Convert Arrow data to list of dictionaries.
+
+        Returns:
+            List of dictionaries, one per row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> rows = result.to_dict()
+            >>> print(rows[0])
+            {'id': 1, 'name': 'Alice'}
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("list[dict[str, Any]]", self.data.to_pylist())
+
+    def __len__(self) -> int:
+        """Return number of rows in the Arrow table.
+
+        Returns:
+            Number of rows.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow("SELECT * FROM users")
+            >>> print(len(result))
+            100
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        return cast("int", self.data.num_rows)
+
+    def __iter__(self) -> "Iterator[dict[str, Any]]":
+        """Iterate over rows as dictionaries.
+
+        Yields:
+            Dictionary for each row.
+
+        Raises:
+            ValueError: If no Arrow table is available.
+
+        Examples:
+            >>> result = session.select_to_arrow(
+            ...     "SELECT id, name FROM users"
+            ... )
+            >>> for row in result:
+            ...     print(row["name"])
+        """
+        if self.data is None:
+            msg = "No Arrow table available"
+            raise ValueError(msg)
+
+        yield from self.data.to_pylist()
+
 
 def create_sql_result(
     statement: "SQL",
-    data: Optional[list[dict[str, Any]]] = None,
+    data: list[dict[str, Any]] | None = None,
     rows_affected: int = 0,
-    last_inserted_id: Optional[Union[int, str]] = None,
-    execution_time: Optional[float] = None,
+    last_inserted_id: int | str | None = None,
+    execution_time: float | None = None,
     metadata: Optional["dict[str, Any]"] = None,
     **kwargs: Any,
 ) -> SQLResult:
@@ -663,8 +850,8 @@ def create_arrow_result(
     statement: "SQL",
     data: Any,
     rows_affected: int = 0,
-    last_inserted_id: Optional[Union[int, str]] = None,
-    execution_time: Optional[float] = None,
+    last_inserted_id: int | str | None = None,
+    execution_time: float | None = None,
     metadata: Optional["dict[str, Any]"] = None,
     schema: Optional["dict[str, Any]"] = None,
 ) -> ArrowResult:

sqlspec/core/splitter.py

@@ -16,13 +16,12 @@ MySQL, SQLite, DuckDB, and BigQuery.
 import re
 import threading
 from abc import ABC, abstractmethod
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from enum import Enum
 from re import Pattern
-from typing import Any, Callable, Final, Optional, Union, cast
+from typing import Any, Final, TypeAlias, cast
 
 from mypy_extensions import mypyc_attr
-from typing_extensions import TypeAlias
 
 from sqlspec.core.cache import CacheKey, UnifiedCache
 from sqlspec.utils.logging import get_logger
@@ -98,9 +97,9 @@ class Token:
         return f"Token({self.type.value}, {self.value!r}, {self.line}:{self.column})"
 
 
-TokenHandler: TypeAlias = Callable[[str, int, int, int], Optional[Token]]
-TokenPattern: TypeAlias = Union[str, TokenHandler]
-CompiledTokenPattern: TypeAlias = Union[Pattern[str], TokenHandler]
+TokenHandler: TypeAlias = Callable[[str, int, int, int], Token | None]
+TokenPattern: TypeAlias = str | TokenHandler
+CompiledTokenPattern: TypeAlias = Pattern[str] | TokenHandler
 
 
 @mypyc_attr(allow_interpreted_subclasses=False)
@@ -111,13 +110,13 @@ class DialectConfig(ABC):
 
     def __init__(self) -> None:
         """Initialize dialect configuration."""
-        self._name: Optional[str] = None
-        self._block_starters: Optional[set[str]] = None
-        self._block_enders: Optional[set[str]] = None
-        self._statement_terminators: Optional[set[str]] = None
-        self._batch_separators: Optional[set[str]] = None
-        self._special_terminators: Optional[dict[str, Callable[[list[Token], int], bool]]] = None
-        self._max_nesting_depth: Optional[int] = None
+        self._name: str | None = None
+        self._block_starters: set[str] | None = None
+        self._block_enders: set[str] | None = None
+        self._statement_terminators: set[str] | None = None
+        self._batch_separators: set[str] | None = None
+        self._special_terminators: dict[str, Callable[[list[Token], int], bool]] | None = None
+        self._max_nesting_depth: int | None = None
 
     @property
     @abstractmethod
@@ -436,7 +435,7 @@ class PostgreSQLDialectConfig(DialectConfig):
         return [(TokenType.STRING_LITERAL, self._handle_dollar_quoted_string)]
 
     @staticmethod
-    def _handle_dollar_quoted_string(text: str, position: int, line: int, column: int) -> Optional[Token]:
+    def _handle_dollar_quoted_string(text: str, position: int, line: int, column: int) -> Token | None:
         """Handle PostgreSQL dollar-quoted string literals.
 
         Parses dollar-quoted strings in the format $tag$content$tag$
@@ -613,8 +612,8 @@ class BigQueryDialectConfig(DialectConfig):
         return self._statement_terminators
 
 
-_pattern_cache: Optional[UnifiedCache] = None
-_result_cache: Optional[UnifiedCache] = None
+_pattern_cache: UnifiedCache | None = None
+_result_cache: UnifiedCache | None = None
 _cache_lock = threading.Lock()
 
 
@@ -878,7 +877,7 @@ class StatementSplitter:
         return False
 
 
-def split_sql_script(script: str, dialect: Optional[str] = None, strip_trailing_terminator: bool = False) -> list[str]:
+def split_sql_script(script: str, dialect: str | None = None, strip_trailing_terminator: bool = False) -> list[str]:
     """Split SQL script into individual statements.
 
     Args: