From e334244ca7fddc8a07dafce3d301b5994c94d9ef Mon Sep 17 00:00:00 2001 From: Thijs Damsma Date: Mon, 9 Nov 2020 21:25:36 +0100 Subject: [PATCH] Support UUID=True --- sqlalchemy-stubs/dialects/postgresql/base.pyi | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/sqlalchemy-stubs/dialects/postgresql/base.pyi b/sqlalchemy-stubs/dialects/postgresql/base.pyi index d910e5c..cb31ae0 100644 --- a/sqlalchemy-stubs/dialects/postgresql/base.pyi +++ b/sqlalchemy-stubs/dialects/postgresql/base.pyi @@ -1,8 +1,9 @@ from ... import schema from ...engine import default, reflection from ...sql import compiler, expression, sqltypes, type_api -from typing import Any, Optional, Set, Type, Text, Pattern, Dict +from typing import Any, Optional, Set, Type, Text, Pattern, Dict, TypeVar, overload from datetime import timedelta +import uuid from sqlalchemy.types import INTEGER as INTEGER, BIGINT as BIGINT, SMALLINT as SMALLINT, VARCHAR as VARCHAR, \ CHAR as CHAR, TEXT as TEXT, FLOAT as FLOAT, NUMERIC as NUMERIC, \ @@ -66,10 +67,16 @@ class BIT(sqltypes.TypeEngine[str]): def __init__(self, length: Optional[int] = ..., varying: bool = ...) -> None: ... PGBit = BIT -class UUID(sqltypes.TypeEngine[str]): +_T = TypeVar("_T") + +class UUID(sqltypes.TypeEngine[_T]): __visit_name__: str = ... as_uuid: bool = ... - def __init__(self, as_uuid: bool = ...) -> None: ... + @overload + def __new__(self, as_uuid: Literal[True]) -> UUID[sqltypes.TypeEngine[uuid.UUID]]:... + @overload + def __new__(self, as_uuid: Literal[False, None]) -> UUID[sqltypes.TypeEngine[str]]:... + def __new__(self) -> UUID[sqltypes.TypeEngine[str]]:... def bind_processor(self, dialect: Any): ... def result_processor(self, dialect: Any, coltype: Any): ... PGUuid = UUID