From 8aa1f81b6ef8dffc4a39a79fb3876af7f089e7ce Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 7 Apr 2020 18:44:04 +1200
Subject: [PATCH] Added text cast of records

---
 psycopg3/types/__init__.py    |  2 +-
 psycopg3/types/composite.py   | 52 +++++++++++++++++++++++++++++++++++
 tests/types/test_composite.py | 42 ++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)
 create mode 100644 psycopg3/types/composite.py
 create mode 100644 tests/types/test_composite.py

diff --git a/psycopg3/types/__init__.py b/psycopg3/types/__init__.py
index 8379dfd22..b07d7f770 100644
--- a/psycopg3/types/__init__.py
+++ b/psycopg3/types/__init__.py
@@ -8,6 +8,6 @@ psycopg3 types package
 from .oids import builtins
 
 # Register default adapters
-from . import array, numeric, text  # noqa
+from . import array, composite, numeric, text  # noqa
 
 __all__ = ["builtins"]
diff --git a/psycopg3/types/composite.py b/psycopg3/types/composite.py
new file mode 100644
index 000000000..7a25595e9
--- /dev/null
+++ b/psycopg3/types/composite.py
@@ -0,0 +1,52 @@
+"""
+Support for composite types adaptation.
+"""
+
+import re
+from typing import Any, Generator, Optional, Tuple
+
+from ..pq import Format
+from ..adapt import TypeCaster, Transformer, AdaptContext
+from .oids import builtins
+
+
+TEXT_OID = builtins["text"].oid
+
+
+_re_tokenize = re.compile(
+    br"""(?x)
+      \(? ([,)])                        # an empty token, representing NULL
+    | \(? " ((?: [^"] | "")*) " [,)]    # or a quoted string
+    | \(? ([^",)]+) [,)]                # or an unquoted string
+    """
+)
+
+_re_undouble = re.compile(br'(["\\])\1')
+
+
+@TypeCaster.text(builtins["record"].oid)
+class RecordCaster(TypeCaster):
+    def __init__(self, oid: int, context: AdaptContext = None):
+        super().__init__(oid, context)
+        self.tx = Transformer(context)
+
+    def cast(self, data: bytes) -> Tuple[Any, ...]:
+        cast = self.tx.get_cast_function(TEXT_OID, format=Format.TEXT)
+        return tuple(
+            cast(item) if item is not None else None
+            for item in self.parse_record(data)
+        )
+
+    def parse_record(
+        self, data: bytes
+    ) -> Generator[Optional[bytes], None, None]:
+        if data == b"()":
+            return
+
+        for m in _re_tokenize.finditer(data):
+            if m.group(1) is not None:
+                yield None
+            elif m.group(2) is not None:
+                yield _re_undouble.sub(br"\1", m.group(2))
+            else:
+                yield m.group(3)
diff --git a/tests/types/test_composite.py b/tests/types/test_composite.py
new file mode 100644
index 000000000..e7d9bd29b
--- /dev/null
+++ b/tests/types/test_composite.py
@@ -0,0 +1,42 @@
+import pytest
+
+
+@pytest.mark.parametrize(
+    "rec, want",
+    [
+        ("", ()),
+        # Funnily enough there's no way to represent (None,) in Postgres
+        ("null", ()),
+        ("null,null", (None, None)),
+        ("null, ''", (None, "")),
+        (
+            "42,'foo','ba,r','ba''z','qu\"x'",
+            ("42", "foo", "ba,r", "ba'z", 'qu"x'),
+        ),
+        (
+            "'foo''', '''foo', '\"bar', 'bar\"' ",
+            ("foo'", "'foo", '"bar', 'bar"'),
+        ),
+    ],
+)
+def test_cast_record(conn, want, rec):
+    cur = conn.cursor()
+    res = cur.execute(f"select row({rec})").fetchone()[0]
+    assert res == want
+
+
+def test_cast_all_chars(conn):
+    cur = conn.cursor()
+    for i in range(1, 256):
+        res = cur.execute("select row(chr(%s::int))", (i,)).fetchone()[0]
+        assert res == (chr(i),)
+
+    cur.execute(
+        "select row(%s)" % ",".join(f"chr({i}::int)" for i in range(1, 256))
+    )
+    res = cur.fetchone()[0]
+    assert res == tuple(map(chr, range(1, 256)))
+
+    s = "".join(map(chr, range(1, 256)))
+    res = cur.execute("select row(%s)", [s]).fetchone()[0]
+    assert res == (s,)
-- 
2.47.3