add alias and as_alias to Count method (piccolo-orm#150)

* add `alias` and `as_alias` to Count method

* move `as_alias` to `Selectable`

* add docs and split tests

* minor tweak

* Remove unused `as_alias` in AVG

Author: aminalaee

piccolo/columns/base.py

@@ -233,6 +233,8 @@ def __deepcopy__(self, memo) -> ColumnMeta:
 
 
 class Selectable(metaclass=ABCMeta):
+    alias: t.Optional[str]
+
     @abstractmethod
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         """
@@ -242,6 +244,13 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
         """
         pass
 
+    def as_alias(self, alias: str) -> Selectable:
+        """
+        Allows column names to be changed in the result of a select.
+        """
+        self.alias = alias
+        return self
+
 
 class Column(Selectable):
     """

piccolo/query/methods/select.py

@@ -45,10 +45,6 @@ def __init__(self, column: Column, alias: str = "avg"):
             raise ValueError("Column type must be numeric to run the query.")
         self.alias = alias
 
-    def as_alias(self, alias: str) -> Avg:
-        self.alias = alias
-        return self
-
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         column_name = self.column._meta.get_full_name(just_alias=just_alias)
         return f"AVG({column_name}) AS {self.alias}"
@@ -61,10 +57,17 @@ class Count(Selectable):
     If a column is specified, the count is for non-null values in that
     column. If no column is specified, the count is for all rows, whether
     they have null values or not.
+
+    Band.select(Band.name, Count()).group_by(Band.name).run()
+    Band.select(Band.name, Count(alias="total")).group_by(Band.name).run()
+    Band.select(Band.name, Count().as_alias("total")).group_by(Band.name).run()
     """
 
-    def __init__(self, column: t.Optional[Column] = None):
+    def __init__(
+        self, column: t.Optional[Column] = None, alias: str = "count"
+    ):
         self.column = column
+        self.alias = alias
 
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         if self.column is None:
@@ -73,7 +76,7 @@ def get_select_string(self, engine_type: str, just_alias=False) -> str:
             column_name = self.column._meta.get_full_name(
                 just_alias=just_alias
             )
-        return f"COUNT({column_name}) AS count"
+        return f"COUNT({column_name}) AS {self.alias}"
 
 
 class Max(Selectable):
@@ -89,10 +92,6 @@ def __init__(self, column: Column, alias: str = "max"):
         self.column = column
         self.alias = alias
 
-    def as_alias(self, alias: str) -> Max:
-        self.alias = alias
-        return self
-
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         column_name = self.column._meta.get_full_name(just_alias=just_alias)
         return f"MAX({column_name}) AS {self.alias}"
@@ -111,10 +110,6 @@ def __init__(self, column: Column, alias: str = "min"):
         self.column = column
         self.alias = alias
 
-    def as_alias(self, alias: str) -> Min:
-        self.alias = alias
-        return self
-
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         column_name = self.column._meta.get_full_name(just_alias=just_alias)
         return f"MIN({column_name}) AS {self.alias}"
@@ -136,10 +131,6 @@ def __init__(self, column: Column, alias: str = "sum"):
             raise ValueError("Column type must be numeric to run the query.")
         self.alias = alias
 
-    def as_alias(self, alias: str) -> Sum:
-        self.alias = alias
-        return self
-
     def get_select_string(self, engine_type: str, just_alias=False) -> str:
         column_name = self.column._meta.get_full_name(just_alias=just_alias)
         return f"SUM({column_name}) AS {self.alias}"

tests/table/test_select.py

@@ -448,6 +448,52 @@ def test_count_group_by(self):
             ]
         )
 
+    def test_count_with_alias_group_by(self):
+        """
+        Test grouping and counting all rows with alias.
+        """
+        self.insert_rows()
+        self.insert_rows()
+
+        response = (
+            Band.select(Band.name, Count(alias="total"))
+            .group_by(Band.name)
+            .order_by(Band.name)
+            .run_sync()
+        )
+
+        self.assertTrue(
+            response
+            == [
+                {"name": "CSharps", "total": 2},
+                {"name": "Pythonistas", "total": 2},
+                {"name": "Rustaceans", "total": 2},
+            ]
+        )
+
+    def test_count_with_as_alias_group_by(self):
+        """
+        Test grouping and counting all rows with as_alias.
+        """
+        self.insert_rows()
+        self.insert_rows()
+
+        response = (
+            Band.select(Band.name, Count().as_alias("total"))
+            .group_by(Band.name)
+            .order_by(Band.name)
+            .run_sync()
+        )
+
+        self.assertTrue(
+            response
+            == [
+                {"name": "CSharps", "total": 2},
+                {"name": "Pythonistas", "total": 2},
+                {"name": "Rustaceans", "total": 2},
+            ]
+        )
+
     def test_count_column_group_by(self):
         """
         Test grouping and counting a specific column. Any null values in the