Add function to parse arguments from a string

aronbierbaum · Mar 4, 2024 · 600db62 · 600db62
1 parent 8b60c24
commit 600db62
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 1 deletion.
diff --git a/clickhouse_sqlalchemy/drivers/util.py b/clickhouse_sqlalchemy/drivers/util.py
@@ -16,3 +16,30 @@ def get_inner_spec(spec):
             break
 
     return spec[offset + 1:i]
+
+
+def parse_arguments(param_string):
+    """
+    Given a string of function arguments, parse them into a tuple.
+    """
+    params = []
+    bracket_level = 0
+    current_param = ''
+
+    for char in param_string:
+        if char == '(':
+            bracket_level += 1
+        elif char == ')':
+            bracket_level -= 1
+        elif char == ',' and bracket_level == 0:
+            params.append(current_param.strip())
+            current_param = ''
+            continue
+
+        current_param += char
+
+    # Append the last parameter
+    if current_param:
+        params.append(current_param.strip())
+
+    return tuple(params)
diff --git a/tests/drivers/test_util.py b/tests/drivers/test_util.py
@@ -1,6 +1,6 @@
 from unittest import TestCase
 
-from clickhouse_sqlalchemy.drivers.util import get_inner_spec
+from clickhouse_sqlalchemy.drivers.util import get_inner_spec, parse_arguments
 
 
 class GetInnerSpecTestCase(TestCase):
@@ -10,3 +10,41 @@ def test_get_inner_spec(self):
         )
         self.assertEqual(get_inner_spec('Decimal(18, 2)'), "18, 2")
         self.assertEqual(get_inner_spec('DateTime64(3)'), "3")
+
+
+class ParseArgumentsTestCase(TestCase):
+    def test_parse_arguments(self):
+        self.assertEqual(
+            parse_arguments('uniq, UInt64'), ('uniq', 'UInt64')
+        )
+        self.assertEqual(
+            parse_arguments('anyIf, String, UInt8'),
+            ('anyIf', 'String', 'UInt8')
+        )
+        self.assertEqual(
+            parse_arguments('quantiles(0.5, 0.9), UInt64'),
+            ('quantiles(0.5, 0.9)', 'UInt64')
+        )
+        self.assertEqual(
+            parse_arguments('sum, Int64, Int64'), ('sum', 'Int64', 'Int64')
+        )
+        self.assertEqual(
+            parse_arguments('sum, Nullable(Int64), Int64'),
+            ('sum', 'Nullable(Int64)', 'Int64')
+        )
+        self.assertEqual(
+            parse_arguments('Float32, Decimal(18, 2)'),
+            ('Float32', 'Decimal(18, 2)')
+        )
+        self.assertEqual(
+            parse_arguments('sum, Float32, Decimal(18, 2)'),
+            ('sum', 'Float32', 'Decimal(18, 2)')
+        )
+        self.assertEqual(
+            parse_arguments('quantiles(0.5, 0.9), UInt64'),
+            ('quantiles(0.5, 0.9)', 'UInt64')
+        )
+        self.assertEqual(
+            parse_arguments("sumIf(total, status = 'accepted'), Float32"),
+            ("sumIf(total, status = 'accepted')", "Float32")
+        )