codekingpro/portable-devtools

1"""2Compatibility objects with DBAPI 2.03"""4 5# Copyright (C) 2020 The Psycopg Team6 7import time8import datetime as dt9from math import floor10from typing import Any, Sequence, Union11 12from . import postgres13from .abc import AdaptContext, Buffer14from .types.string import BytesDumper, BytesBinaryDumper15 16 17class DBAPITypeObject:18    def __init__(self, name: str, type_names: Sequence[str]):19        self.name = name20        self.values = tuple(postgres.types[n].oid for n in type_names)21 22    def __repr__(self) -> str:23        return f"psycopg.{self.name}"24 25    def __eq__(self, other: Any) -> bool:26        if isinstance(other, int):27            return other in self.values28        else:29            return NotImplemented30 31    def __ne__(self, other: Any) -> bool:32        if isinstance(other, int):33            return other not in self.values34        else:35            return NotImplemented36 37 38BINARY = DBAPITypeObject("BINARY", ("bytea",))39DATETIME = DBAPITypeObject(40    "DATETIME", "timestamp timestamptz date time timetz interval".split()41)42NUMBER = DBAPITypeObject("NUMBER", "int2 int4 int8 float4 float8 numeric".split())43ROWID = DBAPITypeObject("ROWID", ("oid",))44STRING = DBAPITypeObject("STRING", "text varchar bpchar".split())45 46 47class Binary:48    def __init__(self, obj: Any):49        self.obj = obj50 51    def __repr__(self) -> str:52        sobj = repr(self.obj)53        if len(sobj) > 40:54            sobj = f"{sobj[:35]} ... ({len(sobj)} byteschars)"55        return f"{self.__class__.__name__}({sobj})"56 57 58class BinaryBinaryDumper(BytesBinaryDumper):59    def dump(self, obj: Union[Buffer, Binary]) -> Buffer:60        if isinstance(obj, Binary):61            return super().dump(obj.obj)62        else:63            return super().dump(obj)64 65 66class BinaryTextDumper(BytesDumper):67    def dump(self, obj: Union[Buffer, Binary]) -> Buffer:68        if isinstance(obj, Binary):69            return super().dump(obj.obj)70        else:71            return super().dump(obj)72 73 74def Date(year: int, month: int, day: int) -> dt.date:75    return dt.date(year, month, day)76 77 78def DateFromTicks(ticks: float) -> dt.date:79    return TimestampFromTicks(ticks).date()80 81 82def Time(hour: int, minute: int, second: int) -> dt.time:83    return dt.time(hour, minute, second)84 85 86def TimeFromTicks(ticks: float) -> dt.time:87    return TimestampFromTicks(ticks).time()88 89 90def Timestamp(91    year: int, month: int, day: int, hour: int, minute: int, second: int92) -> dt.datetime:93    return dt.datetime(year, month, day, hour, minute, second)94 95 96def TimestampFromTicks(ticks: float) -> dt.datetime:97    secs = floor(ticks)98    frac = ticks - secs99    t = time.localtime(ticks)100    tzinfo = dt.timezone(dt.timedelta(seconds=t.tm_gmtoff))101    rv = dt.datetime(*t[:6], round(frac * 1_000_000), tzinfo=tzinfo)102    return rv103 104 105def register_dbapi20_adapters(context: AdaptContext) -> None:106    adapters = context.adapters107    adapters.register_dumper(Binary, BinaryTextDumper)108    adapters.register_dumper(Binary, BinaryBinaryDumper)109 110    # Make them also the default dumpers when dumping by bytea oid111    adapters.register_dumper(None, BinaryTextDumper)112    adapters.register_dumper(None, BinaryBinaryDumper)113