dycw-utilities 0.135.0__py3-none-any.whl → 0.178.1__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
@@ -35,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,
@@ -44,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
@@ -65,7 +64,7 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.base import ReadOnlyColumnCollection
     from whenever import TimeDelta
 
-    from utilities.types import MaybeType, TimeZoneLike
+    from utilities.types import Delta, MaybeType, TimeZoneLike
 
 
 async def insert_dataframe(
@@ -75,9 +74,9 @@ 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: TimeDelta | None = None,
     error_create: type[Exception] = TimeoutError,
     timeout_insert: TimeDelta | None = None,
@@ -87,43 +86,25 @@ async def insert_dataframe(
     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(
@@ -207,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],
@@ -228,7 +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,
-    timeout: TimeDelta | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame: ...
 @overload
@@ -243,7 +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,
-    timeout: TimeDelta | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> Iterable[DataFrame]: ...
 @overload
@@ -258,7 +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,
-    timeout: TimeDelta | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> AsyncIterable[DataFrame]: ...
 @overload
@@ -273,7 +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,
-    timeout: TimeDelta | None = None,
+    timeout: Delta | None = None,
+    error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]: ...
 async def select_to_dataframe(
@@ -287,7 +263,7 @@ 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: TimeDelta | None = None,
+    timeout: Delta | None = None,
     error: MaybeType[BaseException] = TimeoutError,
     **kwargs: Any,
 ) -> DataFrame | Iterable[DataFrame] | AsyncIterable[DataFrame]:
@@ -390,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):
@@ -429,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, max_length=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)