Repository: cogent3/cogent3
Pull request #1458 (2023-05-29 01:19:50)

API: extended behaviour of AnnotationDb.count_distinct()
[NEW] each argument can now take a string or bool value. The former is used
    as a condition for identifying the subset of rows to apply the boolean
    conditions.


Patch:
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 5fc844315..74f59f229 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -747,26 +747,56 @@ class SqliteAnnotationDbMixin:
             num += result
         return num
 
+    StrOrBool = typing.Union[str, bool]
+
     def count_distinct(
         self,
         *,
-        seqid: bool = False,
-        biotype: bool = False,
-        name: bool = False,
+        seqid: StrOrBool = False,
+        biotype: StrOrBool = False,
+        name: StrOrBool = False,
     ) -> typing.Optional[Table]:
-        """return table of counts of distinct values"""
-        conditions = {k for k, v in locals().items() if k != "self" and v}
-        header = list(conditions)
-        columns = ", ".join(conditions)
+        """return table of counts of distinct values
+
+        Parameters
+        ----------
+        seqid, biotype, name
+            if a string, selects the subset of rows matching the provided values
+            and counts distinct values for the other fields whose value is True.
+
+        Returns
+        -------
+        Table with columns corresponding to argument whose value was True
+
+        Examples
+        --------
+        To compute copy number by gene name within each genome
+
+        >>> counts_table = db.count_distinct(seqid=True, biotype='gene', name=True)
+        """
+        columns = {k for k, v in locals().items() if v is True}
         if not columns:
             return
+
+        if constraints := {k: v for k, v in locals().items() if isinstance(v, str)}:
+            where_clause, values = _matching_conditions(constraints)
+            where_clause = f"WHERE {where_clause}"
+        else:
+            where_clause = ""
+            values = ()
+
+        header = list(columns)
         placeholder = "{}"
-        sql_template = f"SELECT {columns}, COUNT(*) as count FROM {placeholder} GROUP BY {columns};"
+        sql_template = (
+            f"SELECT {', '.join(header)}, COUNT(*) as count FROM {placeholder}"
+            f" {where_clause} GROUP BY {', '.join(header)};"
+        )
         data = []
         for table in self.table_names:
             sql = sql_template.format(table)
-            if result := self._execute_sql(sql).fetchall():
+            if result := self._execute_sql(sql, values).fetchall():
                 data.extend(tuple(r) for r in result)
+
         return Table(
             header=header + ["count"],
             data=data,


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index b4cc35c88..e889c691a 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -109,10 +109,32 @@ def test_count_distinct_values(gb_db):
     # there are 8 biotypes in the c.elegans gff sample, 2 columns
     # all arguments returns, from our example, all the rows
     got = {tuple(r) for r in gb_db.count_distinct(name=True).tolist()}
-    expect = {("CNA00110", 4), ("CNA00120", 3), (f"cgg", 1), ("cat", 1), ("JEC21", 1)}
+    expect = {("CNA00110", 4), ("CNA00120", 3), ("cgg", 1), ("cat", 1), ("JEC21", 1)}
     assert got == expect
 
 
+def test_count_distinct_gene_name(gb_db):
+    expect = {("CNA00110", 1), ("CNA00120", 1)}
+    assert {
+        tuple(r) for r in gb_db.count_distinct(biotype="gene", name=True).tolist()
+    } == expect
+
+    assert {
+        tuple(r)
+        for r in gb_db.count_distinct(
+            seqid="AE017341", biotype="gene", name=True
+        ).tolist()
+    } == expect
+
+
+def test_count_distinct_no_match(gb_db):
+    # return a table with 0 rows, 2 columns
+    got = gb_db.count_distinct(biotype=True, name="blah")
+    assert got.shape == (0, 2)
+    got = gb_db.count_distinct(biotype="madeup", name=True)
+    assert got.shape == (0, 2)
+
+
 def test_gff_features_matching(gff_db):
     result = list(gff_db.get_features_matching(biotype="CDS"))
     assert len(result)