Repository: cogent3/cogent3
Pull request #1457 (2023-05-26 01:43:37)

Enhancements to annotation db


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index d3e693f06..81c32cb82 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -2111,14 +2111,9 @@ class SequenceCollection(_SequenceCollectionBase):
         else:
             seq_ids = self.names
 
-        if isinstance(self.annotation_db, GffAnnotationDb):
-            self.annotation_db = load_annotations(
-                path=f, seqids=seq_ids, db=self.annotation_db.db
-            )
-        else:
-            db = load_annotations(path=f, seqids=seq_ids)
-            db.update(self.annotation_db)
-            self.annotation_db = db
+        self.annotation_db = load_annotations(
+            path=f, seqids=seq_ids, db=self.annotation_db
+        )
 
     def make_feature(
         self,
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 021332a63..5fc844315 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -488,7 +488,22 @@ class SqliteAnnotationDbMixin:
     def table_names(self) -> tuple[str]:
         return self._table_names
 
-    def _init_tables(self):
+    def _setup_db(self, db: typing.Optional[SupportsFeatures]) -> None:
+        """initialises the db, using the db passed to the constructor"""
+        if isinstance(db, self.__class__):
+            self._db = db.db
+            return
+
+        if db and not self.compatible(db):
+            raise TypeError(f"cannot initialise annotation db from {type(db)}")
+
+        self._init_tables()
+
+        if db and len(db):
+            # update self with data from other
+            self.update(db)
+
+    def _init_tables(self) -> None:
         # bit of magic
         # assumes schema attributes named as `_<table name>_schema`
         for table_name in self.table_names:
@@ -581,7 +596,7 @@ class SqliteAnnotationDbMixin:
     def _get_feature_by_id(
         self,
         table_name: str,
-        columns: list[str],
+        columns: typing.Optional[list[str], tuple[str]],
         column: str,
         name: str,
         start: OptionalInt = None,
@@ -599,7 +614,7 @@ class SqliteAnnotationDbMixin:
         )
         for result in self._execute_sql(sql, values=vals):
             result = dict(zip(result.keys(), result))
-            result["on_alignment"] = result.get("on_alignment", None)
+            result["on_alignment"] = result.get("on_alignment")
             result["spans"] = [tuple(c) for c in result["spans"]]
             result["reversed"] = result.pop("strand", None) == "-"
             yield result
@@ -708,7 +723,7 @@ class SqliteAnnotationDbMixin:
                 table_name=table_name, columns=columns, **query_args
             ):
                 result = dict(zip(result.keys(), result))
-                result["on_alignment"] = result.get("on_alignment", None)
+                result["on_alignment"] = result.get("on_alignment")
                 result["spans"] = [tuple(c) for c in result["spans"]]
                 result["reversed"] = result.pop("strand", None) == "-"
                 yield result
@@ -746,28 +761,26 @@ class SqliteAnnotationDbMixin:
         if not columns:
             return
         placeholder = "{}"
-        if len(conditions) == 1:
-            sql_template = f"SELECT {columns}, COUNT(DISTINCT {columns}) FROM {placeholder} GROUP BY {columns};"
-        else:
-            sql_template = f"SELECT {columns}, COUNT(*) as count FROM {placeholder} GROUP BY {columns};"
-
+        sql_template = f"SELECT {columns}, COUNT(*) as count FROM {placeholder} GROUP BY {columns};"
         data = []
         for table in self.table_names:
             sql = sql_template.format(table)
             if result := self._execute_sql(sql).fetchall():
                 data.extend(tuple(r) for r in result)
-        return Table(header=header + ["count"], data=data)
+        return Table(
+            header=header + ["count"],
+            data=data,
+            column_templates=dict(count=lambda x: f"{x:,}"),
+        )
 
     @property
     def describe(self) -> Table:
         """top level description of the annotation db"""
-        from cogent3 import make_table
-
-        sql_template = "SELECT {}, COUNT(DISTINCT {}) FROM {} GROUP BY {};"
+        sql_template = "SELECT {}, COUNT(*) FROM {} GROUP BY {};"
         data = {}
         for column in ("seqid", "biotype"):
             for table in self.table_names:
-                sql = sql_template.format(column, column, table, column)
+                sql = sql_template.format(column, table, column)
                 if result := self._execute_sql(sql).fetchall():
                     counts = dict(tuple(r) for r in result)
                     for distinct, count in counts.items():
@@ -776,16 +789,18 @@ class SqliteAnnotationDbMixin:
 
         row_counts = []
         for table in self.table_names:
-            result = self._execute_sql(f"SELECT COUNT(*) FROM {table}").fetchone()
-            row_counts.append(result["COUNT(*)"])
+            result = self._execute_sql(
+                f"SELECT COUNT(*) as count FROM {table}"
+            ).fetchone()
+            row_counts.append(result["count"])
 
-        table = make_table(
+        return Table(
             data={
                 "": list(data.keys()) + [f"num_rows({t!r})" for t in self.table_names],
-                "count": [v for v in data.values()] + row_counts,
-            }
+                "count": list(data.values()) + row_counts,
+            },
+            column_templates=dict(count=lambda x: f"{x:,}"),
         )
-        return table
 
     def biotype_counts(self) -> dict:
         """return counts of biological types across all tables and seqids"""
@@ -826,6 +841,8 @@ class SqliteAnnotationDbMixin:
             checks only that tables of other_db equal, or are a subset, of
             mine
         """
+        if not isinstance(other_db, SupportsFeatures):
+            raise TypeError(f"{type(other_db)} does not support features")
         mine = set(self.table_names)
         theirs = set(other_db.table_names)
         return mine <= theirs or mine > theirs if symmetric else mine >= theirs
@@ -919,7 +936,7 @@ class BasicAnnotationDb(SqliteAnnotationDbMixin):
     _table_names = ("user",)
 
     def __init__(
-        self, *, data: T = None, db: OptionalDbCursor = None, source=":memory:"
+        self, *, data: T = None, db: SupportsFeatures = None, source=":memory:"
     ):
         """
         Parameters
@@ -927,7 +944,8 @@ class BasicAnnotationDb(SqliteAnnotationDbMixin):
         data
             data for entry into the database
         db
-            an existing SQLite cursor
+            a compatible annotation db instance. If db is the same class, it's
+            db will be bound to self and directly modified.
         source
             location to store the db, defaults to in memory only
         """
@@ -935,9 +953,8 @@ class BasicAnnotationDb(SqliteAnnotationDbMixin):
         # note that data is destroyed
         self._num_fakeids = 0
         self.source = source
-        self._db = db
-        if db is None:
-            self._init_tables()
+        self._db = None
+        self._setup_db(db)
 
         self.add_records(data)
 
@@ -986,9 +1003,8 @@ class GffAnnotationDb(SqliteAnnotationDbMixin):
         # note that data is destroyed
         self._num_fakeids = 0
         self.source = source
-        self._db = db
-        if db is None:
-            self._init_tables()
+        self._db = None
+        self._setup_db(db)
         data = self._merged_data(data)
         self.add_records(data)
 
@@ -1094,24 +1110,31 @@ class GenbankAnnotationDb(SqliteAnnotationDbMixin):
         seqid: OptionalStr = None,
         db: OptionalDbCursor = None,
         source=":memory:",
+        namer: typing.Callable = None,
     ):
         """
         seqid
             name of the sequence data is associated with
+        namer
+            callable that takes a record dict and returns a
+            [name]
         """
         data = data or []
         # note that data is destroyed
-        self._db = db
         self._num_fakeids = 0
         self.source = source
-        self._init_tables()
+        self._db = None
+        self._setup_db(db)
+        if db:
+            # guessing there's multiple seqid's
+            self._serialisable["seqid"] = "<multiple seqids>"
+
+        self._namer = namer if callable(namer) else self._default_namer
         self.add_records(data, seqid)
 
     def add_records(self, records, seqid):
-        # need to capture genetic code from genbank, but what a
-
-        col_key_map = {"type": "biotype", "locus_tag": "name"}
-        exclude = {"translation", "location"}  # location is grabbed directly
+        # need to capture genetic code from genbank
+        exclude = {"translation", "location", "type"}  # location is grabbed directly
         for record in records:
             # our Feature code assumes start always < end,
             store = {"seqid": seqid}
@@ -1123,21 +1146,45 @@ class GenbankAnnotationDb(SqliteAnnotationDbMixin):
                 store["start"] = int(store["spans"].min())
                 store["end"] = int(store["spans"].max())
 
-            record_keys = col_key_map.keys() & record.keys()
-            attrs_keys = record.keys() - col_key_map.keys() - exclude
-            store.update({col_key_map[k]: record[k] for k in record_keys})
+            attrs_keys = record.keys() - exclude
+            store["biotype"] = record["type"]
+            name = self._namer(record)
+            if name is None:
+                name = self._make_fake_id(record)
             store["attributes"] = {k: record[k] for k in attrs_keys}
-            if "name" not in store:
-                store["name"] = [f"fakeid-{self._num_fakeids}"]
-                self._num_fakeids += 1
-
-            store["name"] = ",".join(store["name"])
+            store["name"] = ",".join(name)
             cmnd, vals = _add_record_sql(
                 self.table_names[0],
                 store,
             )
             self._execute_sql(cmnd=cmnd, values=vals)
 
+    def _default_namer(self, record: dict) -> typing.Union[typing.List[str], None]:
+        # we evaluate potential tokens in the genbank record in order of
+        # preference for naming. If none of these are found, a fake name
+        # will be generated
+        for key in (
+            "gene",
+            "locus_tag",
+            "strain",
+            "rpt_unit_seq",
+            "db_xref",
+            "bound_moiety",
+            "regulatory_class",
+        ):
+            if key in record:
+                return record[key]
+
+        if element := record.get("mobile_element_type"):
+            return element[0].split(":")[-1:]
+
+        return None
+
+    def _make_fake_id(self, record: dict) -> typing.List[str]:
+        name = [f"{record.get('type', 'fakeid')}-{self._num_fakeids}"]
+        self._num_fakeids += 1
+        return name
+
     def get_feature_children(
         self,
         name: str,
diff --git a/src/cogent3/core/sequence.py b/src/cogent3/core/sequence.py
index dc2ab9976..dce388ca3 100644
--- a/src/cogent3/core/sequence.py
+++ b/src/cogent3/core/sequence.py
@@ -1096,18 +1096,12 @@ class Sequence(SequenceI):
         f : path to gff annotation file.
         offset : Optional, the offset between annotation coordinates and sequence coordinates.
         """
-        if isinstance(self.annotation_db, GffAnnotationDb):
-            # the db is updated directly
-            self.annotation_db = load_annotations(
-                path=f, seqids=self.name, db=self.annotation_db.db
-            )
-        elif isinstance(self.annotation_db, GenbankAnnotationDb):
+        if isinstance(self.annotation_db, GenbankAnnotationDb):
             raise ValueError("GenbankAnnotationDb already attached")
-        else:
-            db = load_annotations(path=f, seqids=self.name)
-            self.annotation_db = (
-                db.update(self.annotation_db) if len(self.annotation_db) else db
-            )
+
+        self.annotation_db = load_annotations(
+            path=f, seqids=self.name, db=self.annotation_db
+        )
 
         if offset:
             self.annotation_offset = offset
diff --git a/src/cogent3/parse/genbank.py b/src/cogent3/parse/genbank.py
index 3c6603ac8..c1597a6f9 100644
--- a/src/cogent3/parse/genbank.py
+++ b/src/cogent3/parse/genbank.py
@@ -730,7 +730,6 @@ def RichGenbankParser(
                     yield rec["locus"], None
             continue
 
-        db = getattr(db, "db", None)
         seq.annotation_db = GenbankAnnotationDb(
             data=rec.pop("features", None), seqid=rec["locus"], db=db
         )


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index ab8cfef4b..b4cc35c88 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -13,6 +13,7 @@ from cogent3.core.annotation_db import (
     load_annotations,
 )
 from cogent3.core.sequence import Sequence
+from cogent3.parse.genbank import MinimalGenbankParser
 from cogent3.util import deserialise
 
 
@@ -59,6 +60,30 @@ def simple_seq_gff_db() -> Sequence:
     return seq
 
 
+@pytest.mark.parametrize(
+    "db_name,cls", (("gff_db", GenbankAnnotationDb), ("gb_db", GffAnnotationDb))
+)
+def test_constructor_db_fail(db_name, cls, request):
+    db = request.getfixturevalue(db_name)
+    with pytest.raises(TypeError):
+        cls(db=db)
+
+
+@pytest.mark.parametrize(
+    "db_name,cls",
+    (
+        ("gff_db", GffAnnotationDb),
+        ("anno_db", GffAnnotationDb),
+        ("gb_db", GenbankAnnotationDb),
+        ("anno_db", GenbankAnnotationDb),
+    ),
+)
+def test_constructor_db_works(db_name, cls, request):
+    # only compatible db's used to init
+    db = request.getfixturevalue(db_name)
+    cls(db=db)
+
+
 def test_gff_describe(gff_db):
     result = gff_db.describe
     assert isinstance(result, _Table)
@@ -80,6 +105,14 @@ def test_count_distinct(gff_db):
     assert got.shape == (11, 4)
 
 
+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)}
+    assert got == expect
+
+
 def test_gff_features_matching(gff_db):
     result = list(gff_db.get_features_matching(biotype="CDS"))
     assert len(result)
@@ -674,3 +707,26 @@ def test_incompatible():
     gb = GenbankAnnotationDb()
     assert not gff.compatible(gb)
     assert not gb.compatible(gff)
+
+
+@pytest.mark.parametrize("wrong_type", ({}, BasicAnnotationDb().db))
+def test_incompatible_invalid_type(wrong_type):
+    db = BasicAnnotationDb()
+    with pytest.raises(TypeError):
+        db.compatible(wrong_type)
+
+
+def _custom_namer(data):
+    for key in ("gene", "locus_tag", "strain"):
+        if key in data:
+            return data[key]
+    return ["default name"]
+
+
+def test_gb_namer():
+    path = DATA_DIR / "annotated_seq.gb"
+    got = list(MinimalGenbankParser(path.read_text().splitlines()))
+    data = got[0]["features"]
+    db = GenbankAnnotationDb(data=data, namer=_custom_namer, seqid=got[0]["locus"])
+    # there are 2 repeat regions, which we don't catch with our namer
+    assert db.num_matches(name="default name") == 2