]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
Added UUID type load/dump
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 29 Oct 2020 22:02:45 +0000 (23:02 +0100)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Thu, 29 Oct 2020 22:02:45 +0000 (23:02 +0100)
psycopg3/psycopg3/types/__init__.py
psycopg3/psycopg3/types/uuid.py [new file with mode: 0644]
tests/types/test_uuid.py [new file with mode: 0644]

index 0ee0378c17ee998565a5730d79d430a7d697a57c..7ccda4c7ad453c0afe0fd17d3f0dd58729e845ca 100644 (file)
@@ -9,6 +9,7 @@ from .oids import builtins
 
 # Register default adapters
 from . import array, composite, date, json, numeric, singletons, text  # noqa
+from . import uuid  # noqa
 
 # Register associations with array oids
 array.register_all_arrays()
diff --git a/psycopg3/psycopg3/types/uuid.py b/psycopg3/psycopg3/types/uuid.py
new file mode 100644 (file)
index 0000000..fc3b957
--- /dev/null
@@ -0,0 +1,48 @@
+"""
+Adapters for the UUID type.
+"""
+
+# Copyright (C) 2020 The Psycopg Team
+
+# TODO: importing uuid is slow. Don't import it at module level.
+# Should implement lazy dumper registration.
+from uuid import UUID
+
+import codecs
+
+from ..adapt import Dumper, Loader
+from ..proto import DecodeFunc, EncodeFunc
+from .oids import builtins
+
+_encode_ascii = codecs.lookup("ascii").encode
+_decode_ascii = codecs.lookup("ascii").decode
+
+
+@Dumper.text(UUID)
+class UUIDDumper(Dumper):
+
+    oid = builtins["uuid"].oid
+
+    def dump(self, obj: UUID, __encode: EncodeFunc = _encode_ascii) -> bytes:
+        return __encode(obj.hex)[0]
+
+
+@Dumper.binary(UUID)
+class UUIDBinaryDumper(Dumper):
+
+    oid = builtins["uuid"].oid
+
+    def dump(self, obj: UUID) -> bytes:
+        return obj.bytes
+
+
+@Loader.text(builtins["uuid"].oid)
+class UUIDLoader(Loader):
+    def load(self, data: bytes, __decode: DecodeFunc = _decode_ascii) -> UUID:
+        return UUID(__decode(data)[0])
+
+
+@Loader.binary(builtins["uuid"].oid)
+class UUIDBinaryLoader(Loader):
+    def load(self, data: bytes) -> UUID:
+        return UUID(bytes=data)
diff --git a/tests/types/test_uuid.py b/tests/types/test_uuid.py
new file mode 100644 (file)
index 0000000..7352f3c
--- /dev/null
@@ -0,0 +1,22 @@
+from uuid import UUID
+
+import pytest
+
+from psycopg3.adapt import Format
+
+
+@pytest.mark.parametrize("fmt_in", [Format.TEXT, Format.BINARY])
+def test_uuid_dump(conn, fmt_in):
+    ph = "%s" if fmt_in == Format.TEXT else "%b"
+    val = "12345678123456781234567812345679"
+    cur = conn.cursor()
+    cur.execute(f"select {ph} = %s::uuid", (UUID(val), val))
+    assert cur.fetchone()[0] is True
+
+
+@pytest.mark.parametrize("fmt_out", [Format.TEXT, Format.BINARY])
+def test_uuid_load(conn, fmt_out):
+    cur = conn.cursor(format=fmt_out)
+    val = "12345678123456781234567812345679"
+    cur.execute("select %s::uuid", (val,))
+    assert cur.fetchone()[0] == UUID(val)