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 4351e3d2aabb61ecf6af4e03dbf526c92c284e92..549733cb9dc73c84b71359537c1b378542640870 100644 (file)
--- 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 8d863545a4aa4ac9c270ac29d4c40503c45affbf..a4cf19b3d5e668cfd479e691f27cb29ae1eece12 100644 (file)
--- 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 855e2855de8b83991d777e945e88214d4787dae5..eeacd1bc8550adc39a93b6306ef99b1563a5773f 100644 (file)
--- 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]