import codecs
import struct
+from typing import Callable, Dict, Tuple, cast
from decimal import Decimal
from ..adapt import Dumper, Loader
+from ..proto import EncodeFunc, DecodeFunc
from .oids import builtins
+UnpackInt = Callable[[bytes], Tuple[int]]
+UnpackFloat = Callable[[bytes], Tuple[float]]
+
FLOAT8_OID = builtins["float8"].oid
NUMERIC_OID = builtins["numeric"].oid
BOOL_OID = builtins["bool"].oid
-_encode = codecs.lookup("ascii").encode
-_decode = codecs.lookup("ascii").decode
-
-_int2_struct = struct.Struct("!h")
-_int4_struct = struct.Struct("!i")
-_int8_struct = struct.Struct("!q")
-_oid_struct = struct.Struct("!I")
-_float4_struct = struct.Struct("!f")
-_float8_struct = struct.Struct("!d")
-
@Dumper.text(int)
class TextIntDumper(Dumper):
- def dump(self, obj: int) -> bytes:
- # We don't know the size of it, so we have to return a type big enough
- return _encode(str(obj))[0]
+ def dump(
+ self, obj: int, __encode: EncodeFunc = codecs.lookup("ascii").encode
+ ) -> bytes:
+ return __encode(str(obj))[0]
@property
def oid(self) -> int:
+ # We don't know the size of it, so we have to return a type big enough
return NUMERIC_OID
@Dumper.text(float)
class TextFloatDumper(Dumper):
- def dump(self, obj: float) -> bytes:
- # Float can't be bigger than this instead
- return _encode(str(obj))[0]
+ def dump(
+ self, obj: float, __encode: EncodeFunc = codecs.lookup("ascii").encode
+ ) -> bytes:
+ return __encode(str(obj))[0]
@property
def oid(self) -> int:
+ # Float can't be bigger than this instead
return FLOAT8_OID
@Dumper.text(Decimal)
class TextDecimalDumper(Dumper):
- def dump(self, obj: Decimal) -> bytes:
- return _encode(str(obj))[0]
+ def dump(
+ self,
+ obj: Decimal,
+ __encode: EncodeFunc = codecs.lookup("ascii").encode,
+ ) -> bytes:
+ return __encode(str(obj))[0]
@property
def oid(self) -> int:
@Loader.text(builtins["int4"].oid)
@Loader.text(builtins["int8"].oid)
@Loader.text(builtins["oid"].oid)
-def load_int(data: bytes) -> int:
- return int(_decode(data)[0])
+class TextIntLoader(Loader):
+ def load(
+ self, data: bytes, __decode: DecodeFunc = codecs.lookup("ascii").decode
+ ) -> int:
+ return int(__decode(data)[0])
@Loader.binary(builtins["int2"].oid)
-def load_int2_binary(data: bytes) -> int:
- rv: int = _int2_struct.unpack(data)[0]
- return rv
+class BinaryInt2Loader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!h").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.binary(builtins["int4"].oid)
-def load_int4_binary(data: bytes) -> int:
- rv: int = _int4_struct.unpack(data)[0]
- return rv
+class BinaryInt4Loader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!i").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.binary(builtins["int8"].oid)
-def load_int8_binary(data: bytes) -> int:
- rv: int = _int8_struct.unpack(data)[0]
- return rv
+class BinaryInt8Loader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!q").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.binary(builtins["oid"].oid)
-def load_oid_binary(data: bytes) -> int:
- rv: int = _oid_struct.unpack(data)[0]
- return rv
+class BinaryOidLoader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!I").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.text(builtins["float4"].oid)
@Loader.text(builtins["float8"].oid)
-def load_float(data: bytes) -> float:
- # it supports bytes directly
- return float(data)
+class TextFloatLoader(Loader):
+ def load(self, data: bytes) -> float:
+ # it supports bytes directly
+ return float(data)
@Loader.binary(builtins["float4"].oid)
-def load_float4_binary(data: bytes) -> float:
- rv: float = _float4_struct.unpack(data)[0]
- return rv
+class BinaryFloat4Loader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!f").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.binary(builtins["float8"].oid)
-def load_float8_binary(data: bytes) -> float:
- rv: float = _float8_struct.unpack(data)[0]
- return rv
+class BinaryFloat8Loader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __unpack: UnpackInt = cast(UnpackInt, struct.Struct("!d").unpack),
+ ) -> int:
+ return __unpack(data)[0]
@Loader.text(builtins["numeric"].oid)
-def load_numeric(data: bytes) -> Decimal:
- return Decimal(_decode(data)[0])
-
-
-_bool_loads = {b"t": True, b"f": False}
-_bool_binary_loads = {b"\x01": True, b"\x00": False}
+class TextNumericLoader(Loader):
+ def load(
+ self, data: bytes, __decode: DecodeFunc = codecs.lookup("ascii").decode
+ ) -> Decimal:
+ return Decimal(__decode(data)[0])
@Loader.text(builtins["bool"].oid)
-def load_bool(data: bytes) -> bool:
- return _bool_loads[data]
+class TextBoolLoader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __values: Dict[bytes, bool] = {b"t": True, b"f": False},
+ ) -> bool:
+ return __values[data]
@Loader.binary(builtins["bool"].oid)
-def load_bool_binary(data: bytes) -> bool:
- return _bool_binary_loads[data]
+class BinaryBoolLoader(Loader):
+ def load(
+ self,
+ data: bytes,
+ __values: Dict[bytes, bool] = {b"\x01": True, b"\x00": False},
+ ) -> bool:
+ return __values[data]
# Copyright (C) 2020 The Psycopg Team
import codecs
-from typing import Optional, Union
+from typing import Optional, Union, TYPE_CHECKING
from ..adapt import Dumper, Loader
from ..proto import AdaptContext, EncodeFunc, DecodeFunc
from ..pq import Escaping
from .oids import builtins, INVALID_OID
+if TYPE_CHECKING:
+ from ..pq.proto import Escaping as EscapingProto
+
TEXT_OID = builtins["text"].oid
BYTEA_OID = builtins["bytea"].oid
@Loader.text(builtins["bytea"].oid)
-def load_bytea_text(data: bytes) -> bytes:
- return Escaping().unescape_bytea(data)
+class TextByteaLoader(Loader):
+ _escaping: "EscapingProto"
+
+ def __init__(self, oid: int, context: AdaptContext = None):
+ super().__init__(oid, context)
+ if not hasattr(self.__class__, "_escaping"):
+ self.__class__._escaping = Escaping()
+
+ def load(self, data: bytes) -> bytes:
+ return self._escaping.unescape_bytea(data)
@Loader.binary(builtins["bytea"].oid)
@Loader.binary(INVALID_OID)
-def load_bytea_binary(data: bytes) -> bytes:
- return data
+class BinaryByteaLoader(Loader):
+ def load(self, data: bytes) -> bytes:
+ return data