1"""
2Compatibility objects with DBAPI 2.0
3"""
4
5# Copyright (C) 2020-2021 The Psycopg Team
6
7import time
8import datetime as dt
9from math import floor
10from typing import Any, Sequence, Union
11
12from . import postgres
13from .abc import AdaptContext, Buffer
14from .types.string import BytesDumper, BytesBinaryDumper
15
16
17class DBAPITypeObject:
18    def __init__(self, name: str, type_names: Sequence[str]):
19        self.name = name
20        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.values
28        else:
29            return NotImplemented
30
31    def __ne__(self, other: Any) -> bool:
32        if isinstance(other, int):
33            return other not in self.values
34        else:
35            return NotImplemented
36
37
38BINARY = DBAPITypeObject("BINARY", ("bytea",))
39DATETIME = DBAPITypeObject(
40    "DATETIME", "timestamp timestamptz date time timetz interval".split()
41)
42NUMBER = DBAPITypeObject(
43    "NUMBER", "int2 int4 int8 float4 float8 numeric".split()
44)
45ROWID = DBAPITypeObject("ROWID", ("oid",))
46STRING = DBAPITypeObject("STRING", "text varchar bpchar".split())
47
48
49class Binary:
50    def __init__(self, obj: Any):
51        self.obj = obj
52
53    def __repr__(self) -> str:
54        sobj = repr(self.obj)
55        if len(sobj) > 40:
56            sobj = f"{sobj[:35]} ... ({len(sobj)} byteschars)"
57        return f"{self.__class__.__name__}({sobj})"
58
59
60class BinaryBinaryDumper(BytesBinaryDumper):
61    def dump(self, obj: Union[Buffer, Binary]) -> Buffer:
62        if isinstance(obj, Binary):
63            return super().dump(obj.obj)
64        else:
65            return super().dump(obj)
66
67
68class BinaryTextDumper(BytesDumper):
69    def dump(self, obj: Union[Buffer, Binary]) -> Buffer:
70        if isinstance(obj, Binary):
71            return super().dump(obj.obj)
72        else:
73            return super().dump(obj)
74
75
76def Date(year: int, month: int, day: int) -> dt.date:
77    return dt.date(year, month, day)
78
79
80def DateFromTicks(ticks: float) -> dt.date:
81    return TimestampFromTicks(ticks).date()
82
83
84def Time(hour: int, minute: int, second: int) -> dt.time:
85    return dt.time(hour, minute, second)
86
87
88def TimeFromTicks(ticks: float) -> dt.time:
89    return TimestampFromTicks(ticks).time()
90
91
92def Timestamp(
93    year: int, month: int, day: int, hour: int, minute: int, second: int
94) -> dt.datetime:
95    return dt.datetime(year, month, day, hour, minute, second)
96
97
98def TimestampFromTicks(ticks: float) -> dt.datetime:
99    secs = floor(ticks)
100    frac = ticks - secs
101    t = time.localtime(ticks)
102    tzinfo = dt.timezone(dt.timedelta(seconds=t.tm_gmtoff))
103    rv = dt.datetime(*t[:6], round(frac * 1_000_000), tzinfo=tzinfo)
104    return rv
105
106
107def register_dbapi20_adapters(context: AdaptContext) -> None:
108    adapters = context.adapters
109    adapters.register_dumper(Binary, BinaryTextDumper)
110    adapters.register_dumper(Binary, BinaryBinaryDumper)
111
112    # Make them also the default dumpers when dumping by bytea oid
113    adapters.register_dumper(None, BinaryTextDumper)
114    adapters.register_dumper(None, BinaryBinaryDumper)
115