Added codec for built-in type tid

Author: fantix

asyncpg/protocol/codecs/misc.pyx

@@ -31,7 +31,7 @@ cdef init_pseudo_codecs():
 
     # OID and friends
     oid_types = [
-        OIDOID, TIDOID, XIDOID, CIDOID
+        OIDOID, XIDOID, CIDOID
     ]
 
     for oid_type in oid_types:

asyncpg/protocol/codecs/tid.pyx

@@ -0,0 +1,64 @@
+# Copyright (C) 2016-present the asyncpg authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of asyncpg and is released under
+# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
+
+
+cdef tid_encode(ConnectionSettings settings, WriteBuffer buf, obj):
+    cdef int overflow = 0
+    cdef long block, offset
+
+    if not (cpython.PyTuple_Check(obj) or cpython.PyList_Check(obj)):
+        raise TypeError(
+            'list or tuple expected (got type {})'.format(type(obj)))
+
+    if len(obj) != 2:
+        raise ValueError(
+            'invalid number of elements in tid tuple, expecting 2')
+
+    try:
+        block = cpython.PyLong_AsLong(obj[0])
+    except OverflowError:
+        overflow = 1
+
+    # "long" and "long long" have the same size for x86_64, need an extra check
+    if overflow or (sizeof(block) > 4 and (block < -2147483648 or
+                                                 block > 2147483647)):
+        raise OverflowError(
+            'block too big to be encoded as INT4: {!r}'.format(obj[0]))
+
+    try:
+        offset = cpython.PyLong_AsLong(obj[1])
+        overflow = 0
+    except OverflowError:
+        overflow = 1
+
+    if overflow or offset < -32768 or offset > 32767:
+        raise OverflowError(
+            'offset too big to be encoded as INT2: {!r}'.format(obj[1]))
+
+    buf.write_int32(6)
+    buf.write_int32(<int32_t>block)
+    buf.write_int16(<int16_t>offset)
+
+
+cdef tid_decode(ConnectionSettings settings, FastReadBuffer buf):
+    cdef:
+        int32_t block
+        int16_t offset
+
+    block = hton.unpack_int32(buf.read(4))
+    offset = hton.unpack_int16(buf.read(2))
+
+    return (block, offset)
+
+
+cdef init_tid_codecs():
+    register_core_codec(TIDOID,
+                        <encode_func>&tid_encode,
+                        <decode_func>&tid_decode,
+                        PG_FORMAT_BINARY)
+
+
+init_tid_codecs()

asyncpg/protocol/protocol.pyx

@@ -66,6 +66,7 @@ include "codecs/json.pyx"
 include "codecs/money.pyx"
 include "codecs/network.pyx"
 include "codecs/numeric.pyx"
+include "codecs/tid.pyx"
 include "codecs/tsearch.pyx"
 include "codecs/txid.pyx"
 include "codecs/uuid.pyx"

docs/usage.rst

@@ -152,6 +152,8 @@ The table below shows the correspondence between PostgreSQL and Python types.
 +----------------------+-----------------------------------------------------+
 | ``uuid``             | :class:`uuid.UUID <python:uuid.UUID>`               |
 +----------------------+-----------------------------------------------------+
+| ``tid``              | :class:`tuple <python:tuple>`                       |
++----------------------+-----------------------------------------------------+
 
 All other types are encoded and decoded as text by default.
 

tests/test_codecs.py

@@ -379,6 +379,9 @@ def _timezone(offset):
     ('circle', 'circle', [
         asyncpg.Circle((0.0, 0.0), 100),
     ]),
+    ('tid', 'tid', [
+        (0, 2),
+    ]),
 ]
 
 
@@ -598,7 +601,31 @@ async def test_invalid_input(self):
             ]),
             ('text', TypeError, 'expected str, got list', [
                 [1]
-            ])
+            ]),
+            ('tid', TypeError, 'list or tuple expected', [
+                b'foo'
+            ]),
+            ('tid', ValueError, 'invalid number of elements in tid tuple', [
+                [],
+                (),
+                [1, 2, 3],
+                (4,),
+            ]),
+            ('tid', OverflowError, 'block too big to be encoded as INT4', [
+                (2**256, 0),
+                (decimal.Decimal("2000000000000000000000000000000"), 0),
+                (0xffffffff, 0),
+                (2**31, 0),
+                (-2**31 - 1, 0),
+            ]),
+            ('tid', OverflowError, 'offset too big to be encoded as INT2', [
+                (0, 2**256),
+                (0, decimal.Decimal("2000000000000000000000000000000")),
+                (0, 0xffff),
+                (0, 0xffffffff),
+                (0, 32768),
+                (0, -32769),
+            ]),
         ]
 
         for typname, errcls, errmsg, data in cases: