feat: add groupby.py

Signed-off-by: Charles Larivier <charles@dribbble.com>

Author: Charles Larivier

src/metabase/mbql/groupby.py

@@ -1,9 +1,7 @@
-from enum import Enum
+from metabase.mbql.base import Mbql, Option
 
-from metabase.mbql.base import Field
 
-
-class TemporalOption(Enum):
+class TemporalOption(Option):
     MINUTE = {"temporal-unit": "minute"}
     HOUR = {"temporal-unit": "hour"}
     DAY = {"temporal-unit": "day"}
@@ -21,24 +19,13 @@ class TemporalOption(Enum):
     QUARTER_OF_YEAR = {"temporal-unit": "quarter-of-year"}
 
 
-class BinOption(Enum):
+class BinOption(Option):
     AUTO = {"binning": {"strategy": "default"}}
     BINS_10 = {"binning": {"strategy": "num-bins", "num-bins": 10}}
     BINS_50 = {"binning": {"strategy": "num-bins", "num-bins": 50}}
     BINS_100 = {"binning": {"strategy": "num-bins", "num-bins": 100}}
     NONE = None
 
 
-class GroupBy(Field):
-    def __init__(self, field_id: int, option=None):
-        super(GroupBy, self).__init__(id=field_id, option=option)
-
-
-class TemporalGroupBy(GroupBy):
-    def __init__(self, field_id: int, option: TemporalOption):
-        super(TemporalGroupBy, self).__init__(field_id=field_id, option=option.value)
-
-
-class BinnedGroupBy(GroupBy):
-    def __init__(self, field_id: int, option: BinOption):
-        super(BinnedGroupBy, self).__init__(field_id=field_id, option=option.value)
+class GroupBy(Mbql):
+    pass

tests/mbql/test_groupby.py

@@ -0,0 +1,19 @@
+from unittest import TestCase
+
+from metabase.mbql.groupby import BinOption, GroupBy, TemporalOption
+
+
+class GroupbyTests(TestCase):
+    def test_groupby_compile(self):
+        """Ensure GroupBy.compile() returns ["field", self.id, self.option]"""
+        groupby = GroupBy(id=4)
+        self.assertEqual(["field", 4, None], groupby.compile())
+
+        groupby = GroupBy(id=4, option=TemporalOption.DAY)
+        self.assertEqual(["field", 4, {"temporal-unit": "day"}], groupby.compile())
+
+        groupby = GroupBy(id=4, option=BinOption.BINS_10)
+        self.assertEqual(
+            ["field", 4, {"binning": {"strategy": "num-bins", "num-bins": 10}}],
+            groupby.compile(),
+        )