From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 Mar 2022 00:32:25 +0000 (+0100)
Subject: fix: fix loading of text arrays with dimension information
X-Git-Tag: 3.1~166
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=52d0de41d45b6234a08c692807f33de03119d215;p=thirdparty%2Fpsycopg.git

fix: fix loading of text arrays with dimension information

The dimension information is a prefix such as ``[0:2]=`` in front of the
array. We just discard it when loading to lists, because for Python they
are always 0-based.

https://www.postgresql.org/docs/14/arrays.html#ARRAYS-IO

Close #253.
---

diff --git a/docs/news.rst b/docs/news.rst
index 4351e3d2a..549733cb9 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -7,7 +7,7 @@
 ``psycopg`` release notes
 =========================
 
-Current release
+Future releases
 ---------------
 
 Psycopg 3.1 (unreleased)
@@ -22,6 +22,15 @@ Psycopg 3.1 (unreleased)
 - Drop support for Python 3.6.
 
 
+Psycopg 3.0.11 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Fix `DataError` loading arrays with dimensions information (:ticket:`#253`).
+
+
+Current release
+---------------
+
 Psycopg 3.0.10
 ^^^^^^^^^^^^^^
 
diff --git a/psycopg/psycopg/types/array.py b/psycopg/psycopg/types/array.py
index 8d863545a..a4cf19b3d 100644
--- a/psycopg/psycopg/types/array.py
+++ b/psycopg/psycopg/types/array.py
@@ -323,6 +323,15 @@ class ArrayLoader(BaseArrayLoader):
         stack: List[Any] = []
         cast = self._tx.get_loader(self.base_oid, self.format).load
 
+        # Remove the dimensions information prefix (``[...]=``)
+        if data and data[0] == b"["[0]:
+            if isinstance(data, memoryview):
+                data = bytes(data)
+            idx = data.find(b"=")
+            if idx == -1:
+                raise e.DataError("malformed array, no '=' after dimension information")
+            data = data[idx + 1 :]
+
         re_parse = _get_array_parse_regexp(self.delimiter)
         for m in re_parse.finditer(data):
             t = m.group(1)
diff --git a/tests/types/test_array.py b/tests/types/test_array.py
index 855e2855d..eeacd1bc8 100644
--- a/tests/types/test_array.py
+++ b/tests/types/test_array.py
@@ -244,3 +244,34 @@ def test_load_array_no_comma_separator(conn):
     cur = conn.execute("select '{(2,2),(1,1);(5,6),(3,4)}'::box[]")
     # Not parsed at the moment, but split ok on ; separator
     assert cur.fetchone()[0] == ["(2,2),(1,1)", "(5,6),(3,4)"]
+
+
+@pytest.mark.parametrize("fmt_out", pq.Format)
+@pytest.mark.parametrize(
+    "obj, want",
+    [
+        ("'[0:1]={a,b}'::text[]", ["a", "b"]),
+        ("'[1:1][-2:-1][3:5]={{{1,2,3},{4,5,6}}}'::int[]", [[[1, 2, 3], [4, 5, 6]]]),
+    ],
+)
+def test_array_with_bounds(conn, obj, want, fmt_out):
+    got = conn.execute(f"select {obj}", binary=fmt_out).fetchone()[0]
+    assert got == want
+
+
+@pytest.mark.parametrize("fmt_out", pq.Format)
+def test_all_chars_with_bounds(conn, fmt_out):
+    cur = conn.cursor(binary=fmt_out)
+    for i in range(1, 256):
+        c = chr(i)
+        cur.execute("select '[0:1]={a,b}'::text[] || %s::text[]", ([c],))
+        assert cur.fetchone()[0] == ["a", "b", c]
+
+    a = list(map(chr, range(1, 256)))
+    a.append("\u20ac")
+    cur.execute("select '[0:1]={a,b}'::text[] || %s::text[]", (a,))
+    assert cur.fetchone()[0] == ["a", "b"] + a
+
+    s = "".join(a)
+    cur.execute("select '[0:1]={a,b}'::text[] || %s::text[]", ([s],))
+    assert cur.fetchone()[0] == ["a", "b", s]