dycw-utilities 0.129.10__py3-none-any.whl → 0.175.17__py3-none-any.whl

utilities/sqlalchemy_polars.py

@@ -4,7 +4,7 @@ import datetime as dt
 import decimal
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal, assert_never, cast, overload, override
+from typing import TYPE_CHECKING, Any, cast, overload, override
 from uuid import UUID
 
 import polars as pl
@@ -25,9 +25,8 @@ from polars import (
 )
 from sqlalchemy import Column, Select, select
 from sqlalchemy.exc import DuplicateColumnError
-from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
-from utilities.asyncio import timeout_dur
+from utilities.asyncio import timeout_td
 from utilities.functions import identity
 from utilities.iterables import (
     CheckDuplicatesError,
@@ -36,7 +35,7 @@ from utilities.iterables import (
     chunked,
     one,
 )
-from utilities.polars import zoned_datetime
+from utilities.polars import zoned_date_time_dtype
 from utilities.reprlib import get_repr
 from utilities.sqlalchemy import (
     CHUNK_SIZE_FRAC,
@@ -45,7 +44,6 @@ from utilities.sqlalchemy import (
     get_chunk_size,
     get_columns,
     insert_items,
-    upsert_items,
 )
 from utilities.text import snake_case
 from utilities.typing import is_subclass_gen
@@ -61,14 +59,12 @@ if TYPE_CHECKING:
     )
 
     from polars._typing import PolarsDataType, SchemaDict
+    from sqlalchemy.ext.asyncio import AsyncEngine
     from sqlalchemy.sql import ColumnCollection
     from sqlalchemy.sql.base import ReadOnlyColumnCollection
-    from tenacity.retry import RetryBaseT
-    from tenacity.stop import StopBaseT
-    from tenacity.wait import WaitBaseT
+    from whenever import TimeDelta
 
-    import utilities.types
-    from utilities.types import TimeZoneLike
+    from utilities.types import Delta, MaybeType, TimeZoneLike
 
 
 async def insert_dataframe(
@@ -78,55 +74,37 @@ async def insert_dataframe(
     /,
     *,
     snake: bool = False,
+    is_upsert: bool = False,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
     assume_tables_exist: bool = False,
-    upsert: Literal["selected", "all"] | None = None,
-    timeout_create: utilities.types.Duration | None = None,
+    timeout_create: TimeDelta | None = None,
     error_create: type[Exception] = TimeoutError,
-    timeout_insert: utilities.types.Duration | None = None,
+    timeout_insert: TimeDelta | None = None,
     error_insert: type[Exception] = TimeoutError,
 ) -> None:
     """Insert/upsert a DataFrame into a database."""
     mapping = _insert_dataframe_map_df_schema_to_table(
         df.schema, table_or_orm, snake=snake
     )
-    items = df.select(mapping).rename(mapping).to_dicts()
+    items = df.select(*mapping).rename(mapping).rows(named=True)
     if len(items) == 0:
-        if not df.is_empty():
-            raise InsertDataFrameError(df=df)
         if not assume_tables_exist:
             await ensure_tables_created(
                 engine, table_or_orm, timeout=timeout_create, error=error_create
             )
         return
-    match upsert:
-        case None:
-            await insert_items(
-                engine,
-                (items, table_or_orm),
-                snake=snake,
-                chunk_size_frac=chunk_size_frac,
-                assume_tables_exist=assume_tables_exist,
-                timeout_create=timeout_create,
-                error_create=error_create,
-                timeout_insert=timeout_insert,
-                error_insert=error_insert,
-            )
-        case "selected" | "all" as selected_or_all:  # skipif-ci-and-not-linux
-            await upsert_items(
-                engine,
-                (items, table_or_orm),
-                snake=snake,
-                chunk_size_frac=chunk_size_frac,
-                selected_or_all=selected_or_all,
-                assume_tables_exist=assume_tables_exist,
-                timeout_create=timeout_create,
-                error_create=error_create,
-                timeout_insert=timeout_insert,
-                error_insert=error_insert,
-            )
-        case _ as never:
-            assert_never(never)
+    await insert_items(
+        engine,
+        (items, table_or_orm),
+        snake=snake,
+        is_upsert=is_upsert,
+        chunk_size_frac=chunk_size_frac,
+        assume_tables_exist=assume_tables_exist,
+        timeout_create=timeout_create,
+        error_create=error_create,
+        timeout_insert=timeout_insert,
+        error_insert=error_insert,
+    )
 
 
 def _insert_dataframe_map_df_schema_to_table(
@@ -210,15 +188,6 @@ def _insert_dataframe_check_df_and_db_types(
     )
 
 
-@dataclass(kw_only=True, slots=True)
-class InsertDataFrameError(Exception):
-    df: DataFrame
-
-    @override
-    def __str__(self) -> str:
-        return f"Non-empty DataFrame must resolve to at least 1 item\n\n{self.df}"
-
-
 @overload
 async def select_to_dataframe(
     sel: Select[Any],
@@ -231,10 +200,8 @@ async def select_to_dataframe(
     in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    stop: StopBaseT | None = None,
-    wait: WaitBaseT | None = None,
-    retry: RetryBaseT | None = None,
-    timeout: utilities.types.Duration | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame: ...
 @overload
@@ -249,10 +216,8 @@ async def select_to_dataframe(
     in_clauses: None = None,
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    stop: StopBaseT | None = None,
-    wait: WaitBaseT | None = None,
-    retry: RetryBaseT | None = None,
-    timeout: utilities.types.Duration | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> Iterable[DataFrame]: ...
 @overload
@@ -267,10 +232,8 @@ async def select_to_dataframe(
     in_clauses: tuple[Column[Any], Iterable[Any]],
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    stop: StopBaseT | None = None,
-    wait: WaitBaseT | None = None,
-    retry: RetryBaseT | None = None,
-    timeout: utilities.types.Duration | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> AsyncIterable[DataFrame]: ...
 @overload
@@ -285,10 +248,8 @@ async def select_to_dataframe(
     in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    stop: StopBaseT | None = None,
-    wait: WaitBaseT | None = None,
-    retry: RetryBaseT | None = None,
-    timeout: utilities.types.Duration | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]: ...
 async def select_to_dataframe(
@@ -302,32 +263,16 @@ async def select_to_dataframe(
     in_clauses: tuple[Column[Any], Iterable[Any]] | None = None,
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
-    timeout: utilities.types.Duration | None = None,
-    error: type[Exception] = TimeoutError,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]:
     """Read a table from a database into a DataFrame."""
-    if not issubclass(AsyncEngine, type(engine)):
-        # for handling testing
-        engine = create_async_engine(engine.url)
-        return await select_to_dataframe(
-            sel,
-            engine,
-            snake=snake,
-            time_zone=time_zone,
-            batch_size=batch_size,
-            in_clauses=in_clauses,
-            in_clauses_chunk_size=in_clauses_chunk_size,
-            chunk_size_frac=chunk_size_frac,
-            timeout=timeout,
-            error=error,
-            **kwargs,
-        )
     if snake:
         sel = _select_to_dataframe_apply_snake(sel)
     schema = _select_to_dataframe_map_select_to_df_schema(sel, time_zone=time_zone)
     if in_clauses is None:
-        async with timeout_dur(duration=timeout, error=error):
+        async with timeout_td(timeout, error=error):
             return read_database(
                 sel,
                 cast("Any", engine),
@@ -344,7 +289,7 @@ async def select_to_dataframe(
         chunk_size_frac=chunk_size_frac,
     )
     if batch_size is None:
-        async with timeout_dur(duration=timeout, error=error):
+        async with timeout_td(timeout, error=error):
             dfs = [
                 await select_to_dataframe(
                     sel,
@@ -365,7 +310,7 @@ async def select_to_dataframe(
             return DataFrame(schema=schema)
 
     async def yield_dfs() -> AsyncIterator[DataFrame]:
-        async with timeout_dur(duration=timeout, error=error):
+        async with timeout_td(timeout, error=error):
             for sel_i in sels:
                 for df in await select_to_dataframe(
                     sel_i,
@@ -421,7 +366,7 @@ def _select_to_dataframe_map_table_column_type_to_dtype(
         return pl.Date
     if is_subclass_gen(py_type, dt.datetime):
         has_tz: bool = type_use.timezone
-        return zoned_datetime(time_zone=time_zone) if has_tz else Datetime()
+        return zoned_date_time_dtype(time_zone=time_zone) if has_tz else Datetime()
     if issubclass(py_type, dt.time):
         return Time
     if issubclass(py_type, dt.timedelta):
@@ -460,15 +405,14 @@ def _select_to_dataframe_yield_selects_with_in_clauses(
     in_clauses_chunk_size: int | None = None,
     chunk_size_frac: float = CHUNK_SIZE_FRAC,
 ) -> Iterator[Select[Any]]:
-    max_length = len(sel.selected_columns)
     in_col, in_values = in_clauses
     if in_clauses_chunk_size is None:
         chunk_size = get_chunk_size(
-            engine, chunk_size_frac=chunk_size_frac, scaling=max_length
+            engine, sel.selected_columns, chunk_size_frac=chunk_size_frac
         )
     else:
         chunk_size = in_clauses_chunk_size
     return (sel.where(in_col.in_(values)) for values in chunked(in_values, chunk_size))
 
 
-__all__ = ["InsertDataFrameError", "insert_dataframe", "select_to_dataframe"]
+__all__ = ["insert_dataframe", "select_to_dataframe"]

utilities/string.py

@@ -10,11 +10,10 @@ def substitute_environ(path_or_text: Path | str, /, **kwargs: Any) -> str:
     """Substitute the environment variables in a file."""
     match path_or_text:
         case Path() as path:
-            with path.open() as fh:
-                return substitute_environ(fh.read(), **kwargs)
+            return substitute_environ(path.read_text(), **kwargs)
         case str() as text:
             return Template(text).substitute(environ, **kwargs)
-        case _ as never:
+        case never:
             assert_never(never)