Merge pull request #108 from Maksim-Burtsev/106-map

Update MapType p_type method (closes issue #106)

Author: maximdanilchenko

aiochclient/_types.pyx

@@ -67,6 +67,12 @@ RE_MAP = re.compile(r"^Map\((.*)\)$")
 RE_REPLACE_QUOTE = re.compile(r"(?<!\\)'")
 
 
+cdef str remove_single_quotes(str string):
+    if string[0] == string[-1] == "'":
+        return string[1:-1]
+    return string
+
+
 cdef str decode(char* val):
     """
     Converting bytes from clickhouse with
@@ -169,7 +175,7 @@ cdef class StrType:
 
     cdef str _convert(self, str string):
         if self.container:
-            return string.strip("'")
+            return remove_single_quotes(string)
         return string
 
     cpdef str p_type(self, str string):
@@ -529,7 +535,7 @@ cdef class TupleType:
         return self._convert(value.decode())
 
 
-cdef class  MapType:
+cdef class MapType:
 
     cdef:
         str name
@@ -545,19 +551,10 @@ cdef class  MapType:
         self.key_type = what_py_type(tps[:comma_index], container=True)
         self.value_type = what_py_type(tps[comma_index + 1:], container=True)
 
-    cdef dict _convert(self, string):
-        if isinstance(string, str):
-            string = RE_REPLACE_QUOTE.sub('"', string)
-            string = string.replace('\\', '\\\\')
-            dict_from_string = json.loads(string)
-        else:
-            dict_from_string = string
+    cdef dict _convert(self, str string):
+        key, value = string[1:-1].split(':', 1)
         return {
-            self.key_type.p_type(
-                decode(key.encode()) if isinstance(key, str) else key
-            ): self.value_type.p_type(
-                decode(val.encode()) if isinstance(val, str) else val
-            ) for key, val in dict_from_string.items()
+            self.key_type.p_type(decode(key.encode())): self.value_type.p_type(decode(value.encode()))
         }
 
     cpdef dict p_type(self, string):

aiochclient/types.py

@@ -37,6 +37,12 @@ def datetime_parse_f(string):
 RE_REPLACE_QUOTE = re.compile(r"(?<!\\)'")
 
 
+def remove_single_quotes(string: str) -> str:
+    if string[0] == string[-1] == "'":
+        return string[1:-1]
+    return string
+
+
 class BaseType(ABC):
     __slots__ = ("name", "container")
 
@@ -141,9 +147,9 @@ def unconvert(value) -> bytes:
 
 
 class StrType(BaseType):
-    def p_type(self, string: str):
+    def p_type(self, string: str) -> str:
         if self.container:
-            return string.strip("'")
+            return remove_single_quotes(string)
         return string
 
     @staticmethod
@@ -316,18 +322,12 @@ def __init__(self, name: str, **kwargs):
         self.key_type = what_py_type(tps[:comma_index], container=True)
         self.value_type = what_py_type(tps[comma_index + 1 :], container=True)
 
-    def p_type(self, string: Any) -> dict:
-        if isinstance(string, str):
-            string = RE_REPLACE_QUOTE.sub('"', string)
-            string = string.replace('\\', '\\\\')
-            string = json.loads(string)
+    def p_type(self, string: str) -> dict[Any, Any]:
+        key, value = string[1:-1].split(':', 1)
         return {
-            self.key_type.p_type(
-                self.decode(key.encode()) if isinstance(key, str) else key
-            ): self.value_type.p_type(
-                self.decode(val.encode()) if isinstance(val, str) else val
+            self.key_type.p_type(self.decode(key.encode())): self.value_type.p_type(
+                self.decode(value.encode())
             )
-            for key, val in string.items()
         }
 
     def convert(self, value: bytes) -> dict:

tests.py

@@ -72,6 +72,7 @@ def rows(uuid):
             True,
             {"hello": "world {' and other things"},
             {"hello": {"inner": "world {' and other things"}},
+            {'key1': {'key2': [uuid]}},
             [(1, 2), (3, 4)],
             [('hello', dt.date(2018, 9, 21)), ('world', dt.date(2018, 9, 22))],
         ],
@@ -124,6 +125,7 @@ def rows(uuid):
             False,
             {"hello": "world {'"},
             {"hello": {"inner": "world {'"}},
+            {'key1': {'key2': [uuid, uuid, uuid]}},
             [(0, 1)],
             [
                 ('hello', dt.date(2018, 9, 21)),
@@ -215,6 +217,7 @@ async def all_types_db(chclient, rows):
                             bool Bool,
                             map Map(String, String),
                             map_map Map(String, Map(String, String)),
+                            map_map_array_uuid Map(String, Map(String, Array(UUID))),
                             nested_int Nested(value1 Integer, value2 Integer),
                             nested_str_date Nested(value1 String, value2 Date)
                             ) ENGINE = Memory
@@ -622,6 +625,21 @@ async def test_map_map(self):
         assert record[0] == result
         assert record["map_map"] == result
 
+    async def test_map_map_array_uuid(self, uuid):
+        result = {'key1': {'key2': [uuid]}}
+        print(await self.select_field("map_map_array_uuid"))
+        assert await self.select_field("map_map_array_uuid") == result
+        record = await self.select_record("map_map_array_uuid")
+        assert record[0] == result
+        assert record["map_map_array_uuid"] == result
+
+        result = ("{'key1':{'key2':" f"['{str(uuid)}']" "}}").encode()
+        print(await self.select_field_bytes("map_map_array_uuid"))
+        assert await self.select_field_bytes("map_map_array_uuid") == result
+        record = await self.select_record_bytes("map_map_array_uuid")
+        assert record[0] == result
+        assert record["map_map_array_uuid"] == result
+
     async def test_nullable(self):
         result = 0
         assert await self.select_field("nullable") == result
@@ -1090,6 +1108,14 @@ async def test_json_insert_select(self):
             }
         ]
 
+    async def test_map_map_array_uuid_json(self, uuid):
+        result = await self.ch.fetch(
+            "SELECT map_map_array_uuid FROM all_types WHERE has(nested_int.value1, 0) format JSONEachRow"
+        )
+        assert result[0]['map_map_array_uuid'] == {
+            'key1': {'key2': [str(uuid), str(uuid), str(uuid)]}
+        }
+
     async def test_select_nested_json(self):
         result = await self.ch.fetch(
             "SELECT nested_int, nested_str_date FROM all_types WHERE has(nested_int.value1, 0) format JSONEachRow"