]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
All loader functions converted to classes
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 6 Aug 2020 18:41:34 +0000 (19:41 +0100)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Sun, 23 Aug 2020 18:24:02 +0000 (19:24 +0100)
psycopg3/psycopg3/types/numeric.py
psycopg3/psycopg3/types/text.py
tests/types/test_numeric.py

index de021d86c82537d00786241620212bb3c356ccb0..a276fbaf17c7e2f59e30b7db08b773b60e26bfa2 100644 (file)
@@ -6,52 +6,55 @@ Adapers for numeric types.
 
 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:
@@ -82,67 +85,104 @@ class BinaryBoolDumper(Dumper):
 @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]
index 0353ce1080e27fda01eb74e9823b97a146bc808a..e620b4b6814db6a8cf3037b4e6f5acd23392ee73 100644 (file)
@@ -5,13 +5,16 @@ Adapters for textual types.
 # 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
@@ -109,11 +112,20 @@ class BinaryBytesDumper(Dumper):
 
 
 @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
index 39acc883d50d905ba6214b222b5cc2389eb25a7b..02c36558d1e1d8e348e341c28aac720af1176396 100644 (file)
@@ -5,7 +5,7 @@ import pytest
 
 from psycopg3.adapt import Loader, Transformer, Format
 from psycopg3.types import builtins
-from psycopg3.types.numeric import load_float
+from psycopg3.types.numeric import TextFloatLoader
 
 
 #
@@ -261,7 +261,7 @@ def test_load_numeric_binary(conn):
 )
 def test_numeric_as_float(conn, val):
     cur = conn.cursor()
-    Loader.register(builtins["numeric"].oid, load_float, cur)
+    Loader.register(builtins["numeric"].oid, TextFloatLoader, cur)
 
     val = Decimal(val)
     cur.execute("select %s", (val,))