Added support for any custom python class in QueryResult

Author: chandr-andr

python/psqlpy/_internal/__init__.pyi

@@ -1,20 +1,101 @@
 import types
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import Self
 
+_CustomClass = TypeVar(
+    "_CustomClass",
+)
+
 class QueryResult:
     """Result."""
 
     def result(self: Self) -> list[dict[Any, Any]]:
         """Return result from database as a list of dicts."""
+    def as_class(
+        self: Self,
+        as_class: Callable[..., _CustomClass],
+    ) -> list[_CustomClass]:
+        """Convert results to passed class.
+
+        The main goal of this method is pydantic,
+        msgspec and dataclasses support.
+
+        ### Parameters:
+        - `as_class`: Any callable python class for the results.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        class ExampleOfAsClass:
+            def __init__(self, username: str) -> None:
+                self.username = username
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            query_result: QueryResult = await db_pool.execute(
+                "SELECT username FROM users WHERE id = $1",
+                [100],
+            )
+            class_results: List[ExampleOfAsClass] = query_result.as_class(
+                as_class=ExampleOfAsClass,
+            )
+        ```
+        """
 
 class SingleQueryResult:
     """Single result."""
 
     def result(self: Self) -> dict[Any, Any]:
         """Return result from database as a dict."""
+    def as_class(
+        self: Self,
+        as_class: Callable[..., _CustomClass],
+    ) -> list[_CustomClass]:
+        """Convert results to passed class.
+
+        The main goal of this method is pydantic,
+        msgspec and dataclasses support.
+
+        ### Parameters:
+        - `as_class`: Any callable python class for the results.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        class ExampleOfAsClass:
+            def __init__(self, username: str) -> None:
+                self.username = username
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            connection = await db_pool.connection()
+            async with connection.transaction() as trans:
+                query_result: SingleQueryResult = await trans.fetch_row(
+                    "SELECT username FROM users WHERE id = $1",
+                    [100],
+                )
+
+            class_result: ExampleOfAsClass = query_result.as_class(
+                as_class=ExampleOfAsClass,
+            )
+        ```
+        """
 
 class IsolationLevel(Enum):
     """Class for Isolation Level for transactions."""

src/query_result.rs

@@ -1,7 +1,26 @@
 use pyo3::{pyclass, pymethods, types::PyDict, Py, PyAny, Python, ToPyObject};
 use tokio_postgres::Row;
 
-use crate::{exceptions::rust_errors::RustPSQLDriverPyResult, value_converter::postgres_to_py};
+use crate::{
+    exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
+    value_converter::postgres_to_py,
+};
+
+/// Convert postgres `Row` into Python Dict.
+///
+/// # Errors
+///
+/// May return Err Result if can not convert
+/// postgres type to python or set new key-value pair
+/// in python dict.
+fn row_to_dict<'a>(py: Python<'a>, postgres_row: &'a Row) -> RustPSQLDriverPyResult<&'a PyDict> {
+    let python_dict = PyDict::new(py);
+    for (column_idx, column) in postgres_row.columns().iter().enumerate() {
+        let python_type = postgres_to_py(py, postgres_row, column, column_idx)?;
+        python_dict.set_item(column.name().to_object(py), python_type)?;
+    }
+    Ok(python_dict)
+}
 
 #[pyclass(name = "QueryResult")]
 #[allow(clippy::module_name_repetitions)]
@@ -33,15 +52,31 @@ impl PSQLDriverPyQueryResult {
     pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
         let mut result: Vec<&PyDict> = vec![];
         for row in &self.inner {
-            let python_dict = PyDict::new(py);
-            for (column_idx, column) in row.columns().iter().enumerate() {
-                let python_type = postgres_to_py(py, row, column, column_idx)?;
-                python_dict.set_item(column.name().to_object(py), python_type)?;
-            }
-            result.push(python_dict);
+            result.push(row_to_dict(py, row)?);
         }
         Ok(result.to_object(py))
     }
+
+    /// Convert result from database to any class passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type to python or create new Python class.
+    pub fn as_class<'a>(
+        &'a self,
+        py: Python<'a>,
+        as_class: &'a PyAny,
+    ) -> RustPSQLDriverPyResult<Py<PyAny>> {
+        let mut res: Vec<&PyAny> = vec![];
+        for row in &self.inner {
+            let pydict: &PyDict = row_to_dict(py, row)?;
+            let convert_class_inst = as_class.call((), Some(pydict))?;
+            res.push(convert_class_inst);
+        }
+
+        Ok(res.to_object(py))
+    }
 }
 
 #[pyclass(name = "SingleQueryResult")]
@@ -68,16 +103,35 @@ impl PSQLDriverSinglePyQueryResult {
     /// # Errors
     ///
     /// May return Err Result if can not convert
-    /// postgres type to python or set new key-value pair
-    /// in python dict.
+    /// postgres type to python, can not set new key-value pair
+    /// in python dict or there are no result.
     pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
-        let python_dict = PyDict::new(py);
         if let Some(row) = self.inner.first() {
-            for (column_idx, column) in row.columns().iter().enumerate() {
-                let python_type = postgres_to_py(py, row, column, column_idx)?;
-                python_dict.set_item(column.name().to_object(py), python_type)?;
-            }
+            return Ok(row_to_dict(py, row)?.to_object(py));
+        }
+        Err(RustPSQLDriverError::RustToPyValueConversionError(
+            "There are not results from the query, can't return first row.".into(),
+        ))
+    }
+
+    /// Convert result from database to any class passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type to python, can not create new Python class
+    /// or there are no results.
+    pub fn as_class<'a>(
+        &'a self,
+        py: Python<'a>,
+        as_class: &'a PyAny,
+    ) -> RustPSQLDriverPyResult<&'a PyAny> {
+        if let Some(row) = self.inner.first() {
+            let pydict: &PyDict = row_to_dict(py, row)?;
+            return Ok(as_class.call((), Some(pydict))?);
         }
-        Ok(python_dict.to_object(py))
+        Err(RustPSQLDriverError::RustToPyValueConversionError(
+            "There are not results from the query, can't convert first row.".into(),
+        ))
     }
 }

tox.ini

@@ -21,6 +21,7 @@ deps =
     pytest>=7,<8
     anyio>=3,<4
     maturin>=1,<2
+    pydantic>=2
 allowlist_externals = maturin
 commands_pre =
     maturin develop