]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Added binary typecasting of int, float, bool.
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 2 Apr 2020 05:55:10 +0000 (18:55 +1300)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 2 Apr 2020 05:55:10 +0000 (18:55 +1300)
psycopg3/types/numeric.py
psycopg3/types/oids.py
tests/types/test_numeric.py

index 5744fdb67387469c9e9ae5e7c43e0739227df57b..efb6b46a1e0ef141e1ea40746fe6ed729f8b616d 100644 (file)
@@ -5,6 +5,7 @@ Adapters of numeric types.
 # Copyright (C) 2020 The Psycopg Team
 
 import codecs
+import struct
 from decimal import Decimal
 from typing import Tuple
 
@@ -15,6 +16,13 @@ from .oids import type_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")
+
 
 @Adapter.text(int)
 def adapt_int(obj: int) -> Tuple[bytes, Oid]:
@@ -53,6 +61,30 @@ def cast_int(data: bytes) -> int:
     return int(_decode(data)[0])
 
 
+@Typecaster.binary(type_oid["int2"])
+def cast_binary_int2(data: bytes) -> int:
+    rv: int = _int2_struct.unpack(data)[0]
+    return rv
+
+
+@Typecaster.binary(type_oid["int4"])
+def cast_binary_int4(data: bytes) -> int:
+    rv: int = _int4_struct.unpack(data)[0]
+    return rv
+
+
+@Typecaster.binary(type_oid["int8"])
+def cast_binary_int8(data: bytes) -> int:
+    rv: int = _int8_struct.unpack(data)[0]
+    return rv
+
+
+@Typecaster.binary(type_oid["oid"])
+def cast_binary_oid(data: bytes) -> int:
+    rv: int = _oid_struct.unpack(data)[0]
+    return rv
+
+
 @Typecaster.text(type_oid["float4"])
 @Typecaster.text(type_oid["float8"])
 def cast_float(data: bytes) -> float:
@@ -60,14 +92,32 @@ def cast_float(data: bytes) -> float:
     return float(data)
 
 
+@Typecaster.binary(type_oid["float4"])
+def cast_binary_float4(data: bytes) -> float:
+    rv: float = _float4_struct.unpack(data)[0]
+    return rv
+
+
+@Typecaster.binary(type_oid["float8"])
+def cast_binary_float8(data: bytes) -> float:
+    rv: float = _float8_struct.unpack(data)[0]
+    return rv
+
+
 @Typecaster.text(type_oid["numeric"])
 def cast_numeric(data: bytes) -> Decimal:
     return Decimal(_decode(data)[0])
 
 
 _bool_casts = {b"t": True, b"f": False}
+_bool_binary_casts = {b"\x01": True, b"\x00": False}
 
 
 @Typecaster.text(type_oid["bool"])
 def cast_bool(data: bytes) -> bool:
     return _bool_casts[data]
+
+
+@Typecaster.binary(type_oid["bool"])
+def cast_binary_bool(data: bytes) -> bool:
+    return _bool_binary_casts[data]
index 6a4651becad28027c7ee2bf6fb0d6976a4117de9..32b71b452bad8ed6626c4ce86848cf620860bdf2 100644 (file)
@@ -93,6 +93,11 @@ _oids_table = [
 
 type_oid: Dict[str, Oid] = {name: Oid(oid) for name, oid, _, _ in _oids_table}
 
+# add aliases too
+for r in _oids_table:
+    if r[3] not in type_oid:
+        type_oid[r[3]] = Oid(r[1])
+
 
 def self_update() -> None:
     import subprocess as sp
index 7fed199f48b6ee85f6b299b5705532bf1cf17cfb..7fd6648d3789e6a016c97e529ab68f46b3b01097 100644 (file)
@@ -3,6 +3,10 @@ from math import isnan, isinf, exp
 
 import pytest
 
+from psycopg3.adapt import Typecaster, Format
+from psycopg3.types import type_oid
+from psycopg3.types.numeric import cast_float
+
 
 #
 # Tests with int
@@ -31,9 +35,9 @@ def test_adapt_int(conn, val, expr):
 @pytest.mark.parametrize(
     "val, pgtype, want",
     [
-        ("0", "int", 0),
-        ("1", "int", 1),
-        ("-1", "int", -1),
+        ("0", "integer", 0),
+        ("1", "integer", 1),
+        ("-1", "integer", -1),
         ("0", "int2", 0),
         ("0", "int4", 0),
         ("0", "int8", 0),
@@ -49,18 +53,12 @@ def test_adapt_int(conn, val, expr):
         ("4294967295", "oid", 4294967295),
     ],
 )
-def test_cast_int(conn, val, pgtype, want):
-    cur = conn.cursor()
-    cur.execute("select %%s::%s" % pgtype, (val,))
-    assert cur.pgresult.fformat(0) == 0
-    result = cur.fetchone()[0]
-    assert result == want
-    assert type(result) is type(want)
-
-    # test binary
-    cur.binary = True
+@pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
+def test_cast_int(conn, val, pgtype, want, format):
+    cur = conn.cursor(binary=format == Format.BINARY)
     cur.execute("select %%s::%s" % pgtype, (val,))
-    assert cur.pgresult.fformat(0) == 1
+    assert cur.pgresult.fformat(0) == format
+    assert cur.pgresult.ftype(0) == type_oid[pgtype]
     result = cur.fetchone()[0]
     assert result == want
     assert type(result) is type(want)
@@ -134,9 +132,11 @@ def test_adapt_float_approx(conn, val, expr):
         ("-inf", "float8", -float("inf")),
     ],
 )
-def test_cast_float(conn, val, pgtype, want):
-    cur = conn.cursor()
+@pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
+def test_cast_float(conn, val, pgtype, want, format):
+    cur = conn.cursor(binary=format == Format.BINARY)
     cur.execute("select %%s::%s" % pgtype, (val,))
+    assert cur.pgresult.fformat(0) == format
     result = cur.fetchone()[0]
     assert type(result) is type(want)
     if isnan(want):
@@ -161,9 +161,11 @@ def test_cast_float(conn, val, pgtype, want):
         ("-1.42e40", "float8", -1.42e40),
     ],
 )
-def test_cast_float_approx(conn, expr, pgtype, want):
-    cur = conn.cursor()
+@pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
+def test_cast_float_approx(conn, expr, pgtype, want, format):
+    cur = conn.cursor(binary=format == Format.BINARY)
     cur.execute("select %s::%s" % (expr, pgtype))
+    assert cur.pgresult.fformat(0) == format
     result = cur.fetchone()[0]
     assert result == pytest.approx(want)
 
@@ -206,10 +208,6 @@ def test_roundtrip_numeric(conn, val):
     ],
 )
 def test_numeric_as_float(conn, val):
-    from psycopg3.adapt import Typecaster
-    from psycopg3.types import type_oid
-    from psycopg3.types.numeric import cast_float
-
     cur = conn.cursor()
     Typecaster.register(type_oid["numeric"], cast_float, cur)
 
@@ -228,10 +226,13 @@ def test_numeric_as_float(conn, val):
 #
 
 
+@pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY])
 @pytest.mark.parametrize("b", [True, False, None])
-def test_roundtrip_bool(conn, b):
-    cur = conn.cursor()
-    cur.execute("select %s", (b,)).fetchone()[0] is b
+def test_roundtrip_bool(conn, b, format):
+    cur = conn.cursor(binary=format == Format.BINARY)
+    result = cur.execute("select %s", (b,)).fetchone()[0]
+    assert cur.pgresult.fformat(0) == format
+    assert result is b
 
 
 @pytest.mark.parametrize("pgtype", [None, "float8", "int8", "numeric"])