Allow setting custom codecs on domains and enumerated types

Previously this was disallowed for no particular reason.

Author: elprans

asyncpg/connection.py

@@ -913,12 +913,12 @@ async def set_type_codec(self, typename, *,
         if not typeinfo:
             raise ValueError('unknown type: {}.{}'.format(schema, typename))
 
-        oid = typeinfo['oid']
-        if typeinfo['kind'] != b'b' or typeinfo['elemtype']:
+        if not introspection.is_scalar_type(typeinfo):
             raise ValueError(
                 'cannot use custom codec on non-scalar type {}.{}'.format(
                     schema, typename))
 
+        oid = typeinfo['oid']
         self._protocol.get_settings().add_python_codec(
             oid, typename, schema, 'scalar',
             encoder, decoder, format)

asyncpg/introspection.py

@@ -145,3 +145,14 @@
 WHERE
     t.typname = $1 AND ns.nspname = $2
 '''
+
+
+# 'b' for a base type, 'd' for a domain, 'e' for enum.
+SCALAR_TYPE_KINDS = (b'b', b'd', b'e')
+
+
+def is_scalar_type(typeinfo) -> bool:
+    return (
+        typeinfo['kind'] in SCALAR_TYPE_KINDS and
+        not typeinfo['elemtype']
+    )

tests/test_codecs.py

@@ -1082,6 +1082,40 @@ def hstore_encoder(obj):
                 DROP EXTENSION hstore
             ''')
 
+    async def test_custom_codec_on_domain(self):
+        """Test encoding/decoding using a custom codec on a domain."""
+        await self.con.execute('''
+            CREATE DOMAIN custom_codec_t AS int
+        ''')
+
+        try:
+            await self.con.set_type_codec(
+                'custom_codec_t',
+                encoder=lambda v: str(v),
+                decoder=lambda v: int(v))
+
+            v = await self.con.fetchval('SELECT $1::custom_codec_t', 10)
+            self.assertEqual(v, 10)
+        finally:
+            await self.con.execute('DROP DOMAIN custom_codec_t')
+
+    async def test_custom_codec_on_enum(self):
+        """Test encoding/decoding using a custom codec on an enum."""
+        await self.con.execute('''
+            CREATE TYPE custom_codec_t AS ENUM ('foo', 'bar', 'baz')
+        ''')
+
+        try:
+            await self.con.set_type_codec(
+                'custom_codec_t',
+                encoder=lambda v: str(v).lstrip('enum :'),
+                decoder=lambda v: 'enum: ' + str(v))
+
+            v = await self.con.fetchval('SELECT $1::custom_codec_t', 'foo')
+            self.assertEqual(v, 'enum: foo')
+        finally:
+            await self.con.execute('DROP TYPE custom_codec_t')
+
     async def test_custom_codec_override_binary(self):
         """Test overriding core codecs."""
         import json