dycw-utilities 0.149.11__py3-none-any.whl → 0.150.1__py3-none-any.whl

{dycw_utilities-0.149.11.dist-info → dycw_utilities-0.150.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.149.11
+Version: 0.150.1
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.149.11.dist-info → dycw_utilities-0.150.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=ODa44fHOmk12CX4kZdMwXIvOfgua6psr4jhzHKVTZf0,61
+utilities/__init__.py,sha256=-ff8DCsQu3pztocOWVj0HdCbLUhyIgoNX_Q2gilrFmI,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=2m2a2C-Qgc6OHTTHL332-t66A7xDITt_SORT7a1DJWo,16792
 utilities/atomicwrites.py,sha256=xcOWenTBRS0oat3kg7Sqe51AohNThMQ2ixPL7QCG8hw,5795
@@ -66,7 +66,7 @@ utilities/sentinel.py,sha256=3jIwgpMekWgDAxPDA_hXMP2St43cPhciKN3LWiZ7kv0,1248
 utilities/shelve.py,sha256=4OzjQI6kGuUbJciqf535rwnao-_IBv66gsT6tRGiUt0,759
 utilities/slack_sdk.py,sha256=ppFBvKgfg5IRWiIoKPtpTyzBtBF4XmwEvU3I5wLJikM,2140
 utilities/socket.py,sha256=K77vfREvzoVTrpYKo6MZakol0EYu2q1sWJnnZqL0So0,118
-utilities/sqlalchemy.py,sha256=q2aYUDAC3SE88Lt6XaKa3CLzT_ePaWvQu_OuRk19x9g,35520
+utilities/sqlalchemy.py,sha256=5TPqPvJUCuxIeiMD7IKTptYg11UXNy7ThGw81reCIRU,37290
 utilities/sqlalchemy_polars.py,sha256=18AoEbeNJUKF3-5hroNy9J5LQwS_QJAXbMfKc9sChtk,14250
 utilities/statsmodels.py,sha256=koyiBHvpMcSiBfh99wFUfSggLNx7cuAw3rwyfAhoKpQ,3410
 utilities/string.py,sha256=MB0X6UPTUc06JdAdj-PctZ238IXeCjE5dAJibNw6ZrU,587
@@ -89,8 +89,8 @@ utilities/zoneinfo.py,sha256=oEH-nL3t4h9uawyZqWDtNtDAl6M-CLpLYGI_nI6DulM,1971
 utilities/pytest_plugins/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 utilities/pytest_plugins/pytest_randomly.py,sha256=NXzCcGKbpgYouz5yehKb4jmxmi2SexKKpgF4M65bi10,414
 utilities/pytest_plugins/pytest_regressions.py,sha256=Iwhfv_OJH7UCPZCfoh7ugZ2Xjqjil-BBBsOb8sDwiGI,1471
-dycw_utilities-0.149.11.dist-info/METADATA,sha256=wXNAZioNfRMvkYH6pANvwG5RbGy1D-9z7hSwmANMABA,1698
-dycw_utilities-0.149.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.149.11.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.149.11.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.149.11.dist-info/RECORD,,
+dycw_utilities-0.150.1.dist-info/METADATA,sha256=VjJ5-xcBd6dgd7vO_L77EpWzPfx021Ic9JkmEV8KoVI,1697
+dycw_utilities-0.150.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.150.1.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.150.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.150.1.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.149.11"
+__version__ = "0.150.1"

utilities/sqlalchemy.py

@@ -18,8 +18,19 @@ from itertools import chain
 from math import floor
 from operator import ge, le
 from re import search
-from typing import TYPE_CHECKING, Any, Literal, TypeGuard, assert_never, cast, override
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    TypeGuard,
+    assert_never,
+    cast,
+    overload,
+    override,
+)
 
+import sqlalchemy
+import sqlalchemy.ext.asyncio
 from sqlalchemy import (
     URL,
     Column,
@@ -42,12 +53,17 @@ from sqlalchemy.dialects.postgresql import Insert as postgresql_Insert
 from sqlalchemy.dialects.postgresql import dialect as postgresql_dialect
 from sqlalchemy.dialects.postgresql import insert as postgresql_insert
 from sqlalchemy.dialects.postgresql.asyncpg import PGDialect_asyncpg
+from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.dialects.sqlite import Insert as sqlite_Insert
 from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-from sqlalchemy.exc import ArgumentError, DatabaseError
+from sqlalchemy.exc import (
+    ArgumentError,
+    DatabaseError,
+    OperationalError,
+    ProgrammingError,
+)
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
-from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine
 from sqlalchemy.orm import (
     DeclarativeBase,
     InstrumentedAttribute,
@@ -99,6 +115,18 @@ CHUNK_SIZE_FRAC = 0.8
 ##
 
 
+def check_connect(engine: Engine, /) -> bool:
+    """Check if an engine can connect."""
+    try:
+        with engine.connect() as conn:
+            return bool(conn.execute(text("SELECT 1")).scalar_one())
+    except (OperationalError, ProgrammingError):
+        return False
+
+
+##
+
+
 async def check_engine(
     engine: AsyncEngine,
     /,
@@ -187,7 +215,8 @@ def _columnwise_minmax(*columns: Any, op: Callable[[Any, Any], Any]) -> Any:
 ##
 
 
-def create_async_engine(
+@overload
+def create_engine(
     drivername: str,
     /,
     *,
@@ -198,7 +227,49 @@ def create_async_engine(
     database: str | None = None,
     query: StrMapping | None = None,
     poolclass: type[Pool] | None = NullPool,
-) -> AsyncEngine:
+    async_: Literal[True],
+) -> AsyncEngine: ...
+@overload
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: Literal[False] = False,
+) -> Engine: ...
+@overload
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: bool = False,
+) -> Engine | AsyncEngine: ...
+def create_engine(
+    drivername: str,
+    /,
+    *,
+    username: str | None = None,
+    password: str | None = None,
+    host: str | None = None,
+    port: int | None = None,
+    database: str | None = None,
+    query: StrMapping | None = None,
+    poolclass: type[Pool] | None = NullPool,
+    async_: bool = False,
+) -> Engine | AsyncEngine:
     """Create a SQLAlchemy engine."""
     if query is None:
         kwargs = {}
@@ -217,7 +288,13 @@ def create_async_engine(
         database=database,
         **kwargs,
     )
-    return _create_async_engine(url, poolclass=poolclass)
+    match async_:
+        case False:
+            return sqlalchemy.create_engine(url, poolclass=poolclass)
+        case True:
+            return sqlalchemy.ext.asyncio.create_async_engine(url, poolclass=poolclass)
+        case _ as never:
+            assert_never(never)
 
 
 ##
@@ -825,7 +902,7 @@ def _get_dialect(engine_or_conn: EngineOrConnectionOrAsync, /) -> Dialect:
     if isinstance(dialect, oracle_dialect):  # pragma: no cover
         return "oracle"
     if isinstance(  # skipif-ci-and-not-linux
-        dialect, postgresql_dialect | PGDialect_asyncpg
+        dialect, (postgresql_dialect, PGDialect_asyncpg, PGDialect_psycopg)
     ):
         return "postgresql"
     if isinstance(dialect, sqlite_dialect):
@@ -1107,10 +1184,11 @@ __all__ = [
     "InsertItemsError",
     "TablenameMixin",
     "UpsertItemsError",
+    "check_connect",
     "check_engine",
     "columnwise_max",
     "columnwise_min",
-    "create_async_engine",
+    "create_engine",
     "ensure_tables_created",
     "ensure_tables_dropped",
     "enum_name",