From: Daniele Varrazzo Date: Thu, 2 Apr 2020 05:55:10 +0000 (+1300) Subject: Added binary typecasting of int, float, bool. X-Git-Tag: 3.0.dev0~628 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=2153c5b27bd27927f3197234e4d4eecc0232134c;p=thirdparty%2Fpsycopg.git Added binary typecasting of int, float, bool. --- diff --git a/psycopg3/types/numeric.py b/psycopg3/types/numeric.py index 5744fdb67..efb6b46a1 100644 --- a/psycopg3/types/numeric.py +++ b/psycopg3/types/numeric.py @@ -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] diff --git a/psycopg3/types/oids.py b/psycopg3/types/oids.py index 6a4651bec..32b71b452 100644 --- a/psycopg3/types/oids.py +++ b/psycopg3/types/oids.py @@ -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 diff --git a/tests/types/test_numeric.py b/tests/types/test_numeric.py index 7fed199f4..7fd6648d3 100644 --- a/tests/types/test_numeric.py +++ b/tests/types/test_numeric.py @@ -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"])