From ddbe495d707afc141767da00fadf86cf9429b7c0 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 18 Jan 2019 15:47:01 +0000
Subject: [PATCH] Added BYTESARRAY typecaster

---
 lib/extensions.py           |  8 ++++----
 psycopg/typecast_array.c    |  1 +
 psycopg/typecast_builtins.c |  4 ++--
 tests/test_types_basic.py   | 26 ++++++++++++++++++++++++++
 4 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/lib/extensions.py b/lib/extensions.py
index bb907020..15635751 100644
--- a/lib/extensions.py
+++ b/lib/extensions.py
@@ -35,10 +35,10 @@ This module holds all the extensions to the DBAPI-2.0 provided by psycopg.
 import re as _re
 
 from psycopg2._psycopg import (                             # noqa
-    BINARYARRAY, BOOLEAN, BOOLEANARRAY, BYTES, DATE, DATEARRAY, DATETIMEARRAY,
-    DECIMAL, DECIMALARRAY, FLOAT, FLOATARRAY, INTEGER, INTEGERARRAY,
-    INTERVAL, INTERVALARRAY, LONGINTEGER, LONGINTEGERARRAY, ROWIDARRAY,
-    STRINGARRAY, TIME, TIMEARRAY, UNICODE, UNICODEARRAY,
+    BINARYARRAY, BOOLEAN, BOOLEANARRAY, BYTES, BYTESARRAY, DATE, DATEARRAY,
+    DATETIMEARRAY, DECIMAL, DECIMALARRAY, FLOAT, FLOATARRAY, INTEGER,
+    INTEGERARRAY, INTERVAL, INTERVALARRAY, LONGINTEGER, LONGINTEGERARRAY,
+    ROWIDARRAY, STRINGARRAY, TIME, TIMEARRAY, UNICODE, UNICODEARRAY,
     AsIs, Binary, Boolean, Float, Int, QuotedString, )
 
 try:
diff --git a/psycopg/typecast_array.c b/psycopg/typecast_array.c
index 10c2b374..9cec95c0 100644
--- a/psycopg/typecast_array.c
+++ b/psycopg/typecast_array.c
@@ -286,6 +286,7 @@ typecast_GENERIC_ARRAY_cast(const char *str, Py_ssize_t len, PyObject *curs)
 #define typecast_DECIMALARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_STRINGARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_UNICODEARRAY_cast typecast_GENERIC_ARRAY_cast
+#define typecast_BYTESARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_BOOLEANARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_DATETIMEARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_DATETIMETZARRAY_cast typecast_GENERIC_ARRAY_cast
diff --git a/psycopg/typecast_builtins.c b/psycopg/typecast_builtins.c
index 9d2633d2..0e4901d1 100644
--- a/psycopg/typecast_builtins.c
+++ b/psycopg/typecast_builtins.c
@@ -16,7 +16,6 @@ static long int typecast_LONGINTEGERARRAY_types[] = {1016, 0};
 static long int typecast_INTEGERARRAY_types[] = {1005, 1006, 1007, 0};
 static long int typecast_FLOATARRAY_types[] = {1021, 1022, 0};
 static long int typecast_DECIMALARRAY_types[] = {1231, 0};
-static long int typecast_UNICODEARRAY_types[] = {1002, 1003, 1009, 1014, 1015, 0};
 static long int typecast_STRINGARRAY_types[] = {1002, 1003, 1009, 1014, 1015, 0};
 static long int typecast_BOOLEANARRAY_types[] = {1000, 0};
 static long int typecast_DATETIMEARRAY_types[] = {1115, 0};
@@ -53,7 +52,8 @@ static typecastObject_initlist typecast_builtins[] = {
   {"INTEGERARRAY", typecast_INTEGERARRAY_types, typecast_INTEGERARRAY_cast, "INTEGER"},
   {"FLOATARRAY", typecast_FLOATARRAY_types, typecast_FLOATARRAY_cast, "FLOAT"},
   {"DECIMALARRAY", typecast_DECIMALARRAY_types, typecast_DECIMALARRAY_cast, "DECIMAL"},
-  {"UNICODEARRAY", typecast_UNICODEARRAY_types, typecast_UNICODEARRAY_cast, "UNICODE"},
+  {"UNICODEARRAY", typecast_STRINGARRAY_types, typecast_UNICODEARRAY_cast, "UNICODE"},
+  {"BYTESARRAY", typecast_STRINGARRAY_types, typecast_BYTESARRAY_cast, "BYTES"},
   {"STRINGARRAY", typecast_STRINGARRAY_types, typecast_STRINGARRAY_cast, "STRING"},
   {"BOOLEANARRAY", typecast_BOOLEANARRAY_types, typecast_BOOLEANARRAY_cast, "BOOLEAN"},
   {"DATETIMEARRAY", typecast_DATETIMEARRAY_types, typecast_DATETIMEARRAY_cast, "DATETIME"},
diff --git a/tests/test_types_basic.py b/tests/test_types_basic.py
index 9be4ac2b..b5660b67 100755
--- a/tests/test_types_basic.py
+++ b/tests/test_types_basic.py
@@ -32,6 +32,7 @@ import unittest
 from .testutils import ConnectingTestCase, long
 
 import psycopg2
+from psycopg2.compat import text_type
 
 
 class TypesBasicTests(ConnectingTestCase):
@@ -208,6 +209,31 @@ class TypesBasicTests(ConnectingTestCase):
             self.assertRaises(psycopg2.DataError,
                 psycopg2.extensions.STRINGARRAY, s.encode('utf8'), curs)
 
+    def testTextArray(self):
+        curs = self.conn.cursor()
+        curs.execute("select '{a,b,c}'::text[]")
+        x = curs.fetchone()[0]
+        self.assert_(isinstance(x[0], str))
+        self.assertEqual(x, ['a', 'b', 'c'])
+
+    def testUnicodeArray(self):
+        psycopg2.extensions.register_type(
+            psycopg2.extensions.UNICODEARRAY, self.conn)
+        curs = self.conn.cursor()
+        curs.execute("select '{a,b,c}'::text[]")
+        x = curs.fetchone()[0]
+        self.assert_(isinstance(x[0], text_type))
+        self.assertEqual(x, [u'a', u'b', u'c'])
+
+    def testBytesArray(self):
+        psycopg2.extensions.register_type(
+            psycopg2.extensions.BYTESARRAY, self.conn)
+        curs = self.conn.cursor()
+        curs.execute("select '{a,b,c}'::text[]")
+        x = curs.fetchone()[0]
+        self.assert_(isinstance(x[0], bytes))
+        self.assertEqual(x, [b'a', b'b', b'c'])
+
     @testutils.skip_before_postgres(8, 2)
     def testArrayOfNulls(self):
         curs = self.conn.cursor()