Repository: cogent3/cogent3 (python)

Issue:
AnnotationDb's should support subsetting
If we have annotated features from entire genomes, then selecting a slice of some sequences will be more performant if we create a new AnnotationDb with only the subset of features that map to the selected sequences.

The `subset()` method should be added to the `SupportsQueryFeatures` protocol with the API matching `get_records_matching()`.  The implementations should return a new instance of the same db class, so separate implementations for Basic, Gff and Genbank.

This will support development of EnsemblLite

Pull request description:
ENH: AnnotationDb.subset returns a subset of rows matching condition, fixes #1641
[NEW] uses same querying interface as AnnotationDb.get_records_matching(), but
    returns a new instance of self.

[NEW] added cookbook example too

[NEW] added scriv changelog entries

Patch:
diff --git a/changelog.d/20231201_134621_Gavin.Huttley.md b/changelog.d/20231201_134621_Gavin.Huttley.md
new file mode 100644
index 000000000..42aad507a
--- /dev/null
+++ b/changelog.d/20231201_134621_Gavin.Huttley.md
@@ -0,0 +1,41 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+<!--
+### Contributors
+
+- A bullet item for the Contributors category.
+
+-->
+
+### ENH
+
+- The new AnnotationDb.subset() method creates a subset of the instance
+  with records matching the provided conditions.
+
+### BUG
+
+- The sequence returned from Alignment.get_seq() now reflects
+  slice operations applied to the Alignment.
+
+<!--
+### DOC
+
+- A bullet item for the DOC category.
+
+-->
+<!--
+### Deprecations
+
+- A bullet item for the Deprecations category.
+
+-->
+<!--
+### Discontinued
+
+- A bullet item for the Discontinued category.
+
+-->
diff --git a/doc/cookbook/annotation_db.rst b/doc/cookbook/annotation_db.rst
index d4678af51..e1183b1ad 100644
--- a/doc/cookbook/annotation_db.rst
+++ b/doc/cookbook/annotation_db.rst
@@ -10,7 +10,6 @@ Annotation Databases
 
 This guide shows you how to use ``cogent3``'s annotation databases, which are in-memory SQLite databases, to store, query and manipulate the :ref:`howto-features` (also known as annotations) of one or more biological sequences.
 
-
 What are the different types of ``AnnotationDb``?
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -28,7 +27,7 @@ Achieved by creating an ``BasicAnnotationDb`` instance. This is an empty databas
 
     anno_db = BasicAnnotationDb()
     anno_db
-    
+
 How to load an standalone ``AnnotationDb`` from a data file
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -84,12 +83,12 @@ We can add a feature to the empty ``BasicAnnotationDb`` we created above. Now th
     :raises:
 
     anno_db.add_feature(
-               seqid="NC_000908",
-               biotype="gene",
-               name="interesting_gene",
-               spans=[(1, 4)],
-               strand="+",
-                )
+        seqid="NC_000908",
+        biotype="gene",
+        name="interesting_gene",
+        spans=[(1, 4)],
+        strand="+",
+    )
     anno_db.describe
 
 We can also add a feature to our ``GffAnnotationDb`` or ``GenbankAnnotationDb``. Below, the previously empty "user" table now has a row count of one, indicating that our feature has been successfully added to the database.
@@ -104,7 +103,7 @@ We can also add a feature to our ``GffAnnotationDb`` or ``GenbankAnnotationDb``.
         spans=[(1, 4)],
         strand="+",
     )
-    gff_db.describe[-2:, :] # showing just last two rows
+    gff_db.describe[-2:, :]  # showing just last two rows
 
 How to write an ``AnnotationDb`` to disk for efficient re-loading
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -162,7 +161,7 @@ For example, querying the ``GffAnnotationDb`` for all pseudogenes:
     :raises:
 
     pseudogenes = list(gff_db.get_features_matching(biotype="pseudogene"))
-    pseudogenes[:2] # showing just the first two
+    pseudogenes[:2]  # showing just the first two
 
 Querying via region of interest
 """""""""""""""""""""""""""""""
@@ -192,7 +191,7 @@ For example, we can query for all CDS related to replication:
     replication_records = list(
         gff_db.get_records_matching(attributes="replication", biotype="CDS")
     )
-    replication_records[0] # showing just the first match
+    replication_records[0]  # showing just the first match
 
 .. note:: Extended attribute querying only works for GFF databases!
 
@@ -320,10 +319,26 @@ You can assign a compatible database to the ``db`` argument in the ``AnnotationD
     :raises:
 
     from cogent3.core.annotation_db import GenbankAnnotationDb
-    
+
     new_gb_db = GenbankAnnotationDb(source="m-genitalium-database.gbdb", db=anno_db)
     new_gb_db
 
+How to get a subset of an ``AnnotationDb``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+If you want a subset of a db, use the same arguments as you would for ``db.get_records_matching()``.
+
+.. jupyter-execute::
+    :raises:
+
+    from cogent3 import load_annotations
+
+    gff_db = load_annotations(path="data/mycoplasma-genitalium.gff")
+    just_cds = gff_db.subset(biotype="CDS")
+    just_cds.describe
+
+.. note:: The result is an in-memory database by default. To have this written to disk, assign a path to the source argument, e.g. ``gff_db.subset(source="som/path/subset.gff3db", biotype="CDS")``.
+
 .. _assign_db_to_seq:
 
 How to assign an ``AnnotationDb`` to a sequence
@@ -393,9 +408,7 @@ In the above example, the sequence name in the fasta file does not match any rec
         label_to_name=lambda x: x.split()[0],
     )
     seq.annotation_db
-    
-    
-    
+
 .. jupyter-execute::
     :hide-code:
 
diff --git a/doc/examples/codon_models.rst b/doc/examples/codon_models.rst
index 175405bd0..7a0e020bd 100644
--- a/doc/examples/codon_models.rst
+++ b/doc/examples/codon_models.rst
@@ -217,11 +217,20 @@ The following implements a modification of the approach of Zhang, Nielsen and Ya
 
     from cogent3 import make_table
 
-    header = ['Site Class', 'Proportion', 'Background Edges', 'Foreground Edges']
-    data = {'Site Class': array(['0', '1', '2a', '2b'], dtype='<U2'), 'Proportion': array(['p0', 'p1', 'p2', 'p3'], dtype='<U2'), 'Background Edges': array(['0 < omega0 < 1', 'omega1 = 1', '0 < omega0 < 1', 'omega1 = 1'],
-      dtype='<U14'), 'Foreground Edges': array(['0 < omega0 < 1', 'omega1 = 1', '0 < omega2 > 1', '0 < omega0 < 1'],
-      dtype='<U14')}
-    data = {k: array(data[k], dtype='U') for k in data}
+    header = ["Site Class", "Proportion", "Background Edges", "Foreground Edges"]
+    data = {
+        "Site Class": array(["0", "1", "2a", "2b"], dtype="<U2"),
+        "Proportion": array(["p0", "p1", "p2", "p3"], dtype="<U2"),
+        "Background Edges": array(
+            ["0 < omega0 < 1", "omega1 = 1", "0 < omega0 < 1", "omega1 = 1"],
+            dtype="<U14",
+        ),
+        "Foreground Edges": array(
+            ["0 < omega0 < 1", "omega1 = 1", "0 < omega2 > 1", "0 < omega0 < 1"],
+            dtype="<U14",
+        ),
+    }
+    data = {k: array(data[k], dtype="U") for k in data}
     table = make_table(header, data=data)
     HTML(table.set_repr_policy(show_shape=False))
 
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index 38ebe29f9..373ef5448 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -4837,6 +4837,9 @@ class Alignment(AlignmentI, SequenceCollection):
         return Aligned(map, seq)
 
     def __getitem__(self, index):
+        # todo revisit this design Map's should probably not be used for
+        #  slicing, perhaps they could be used to create a slice object
+        #  instead?
         if hasattr(index, "get_slice"):
             if index.parent is not self:
                 raise ValueError(f"feature.parent {index.seqid!r} is not self")
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 5ac5a59d0..641e78ed0 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -12,6 +12,7 @@ import sys
 import typing
 
 import numpy
+import typing_extensions
 
 from cogent3._version import __version__
 from cogent3.util.deserialise import register_deserialiser
@@ -120,6 +121,19 @@ class SupportsQueryFeatures(typing.Protocol):  # pragma: no cover
     ) -> int:
         ...
 
+    def subset(
+        self,
+        *,
+        seqid: OptionalStr = None,
+        biotype: OptionalStr = None,
+        name: OptionalStr = None,
+        start: OptionalInt = None,
+        end: OptionalInt = None,
+        strand: OptionalStr = None,
+        attributes: OptionalStr = None,
+    ):
+        ...
+
 
 @typing.runtime_checkable
 class SupportsWriteFeatures(typing.Protocol):  # pragma: no cover
@@ -1031,6 +1045,43 @@ class SqliteAnnotationDbMixin:
             self.db.backup(backup)
         backup.close()
 
+    def subset(
+        self,
+        *,
+        source: os.PathLike | str = ":memory:",
+        biotype: str = None,
+        seqid: str = None,
+        name: str = None,
+        start: int = None,
+        end: int = None,
+        strand: bool = None,
+        attributes: OptionalStr = None,
+        allow_partial: bool = False,
+    ) -> typing_extensions.Self:
+        """returns a new db instance with records matching the provided conditions"""
+        kwargs = {k: v for k, v in locals().items() if k not in {"self", "source"}}
+
+        result = self.__class__(source=source)
+        if not len(self):
+            return result
+
+        for table_name in self.table_names:
+            cols = [
+                r["name"]
+                for r in self.db.execute(f"PRAGMA table_info({table_name})").fetchall()
+            ]
+            pos = ", ".join("?" * len(cols))
+            sql = f"INSERT INTO {table_name} ({','.join(cols)}) VALUES ({pos});"
+            records = []
+            for record in self._get_records_matching(table_name=table_name, **kwargs):
+                records.append(tuple(record[c] for c in cols))
+
+            with result.db as cursor:
+                cursor.executemany(sql, records)
+                cursor.commit()
+
+        return result
+
 
 class BasicAnnotationDb(SqliteAnnotationDbMixin):
     """Provides a user table for annotations. This can be merged with
@@ -1156,6 +1207,8 @@ class GffAnnotationDb(SqliteAnnotationDbMixin):
         # I am not convinced we can rely on gff files to be ordered,
         # if we could, we could do this as one pass over the data
         # all keys need to be lower case
+        # NOTE only records which have an ID field get merged into a single
+        # record.
         for record in records:
             record["biotype"] = record.pop("Type")
 


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index c90509ba2..b2a106921 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -950,3 +950,47 @@ def test_gbdb_get_parent_fails_no_coords(gb_db):
 def test_load_annotations_invalid_path():
     with pytest.raises(IOError):
         load_annotations(path="invalidfile.gff3")
+
+
+def test_subset_gff3_db(gff_db):
+    subset = gff_db.subset(seqid="I", start=40, end=70, allow_partial=True)
+    # manual inspection of the original GFF3 file indicates 7 records
+    # BUT the CDS records get merged into a single row
+    assert len(subset) == 6
+
+
+def test_subset_empty_db(gff_db):
+    subset = gff_db.subset(seqid="X", start=40, end=70, allow_partial=True)
+    # no records
+    assert not len(subset)
+
+
+def test_subset_gff3_db_with_user(gff_db):
+    record = dict(
+        seqid="I", name="gene-01", biotype="gene", spans=[(23, 43)], strand="+"
+    )
+    gff_db.add_feature(**record)
+    subset = gff_db.subset(seqid="I", start=40, end=70, allow_partial=True)
+    # manual inspection of the original GFF3 file indicates 7 records
+    # BUT the CDS records get merged into a single row
+    assert len(subset) == 7
+
+
+def test_subset_gb_db(gb_db):
+    subset = gb_db.subset(biotype="gene")
+    # manual inspection of the original annotated gb file indicates 2 genes
+    assert len(subset) == 2
+
+
+def test_subset_gff3_db_source(gff_db, tmp_dir):
+    outpath = tmp_dir / "subset.gff3db"
+    subset = gff_db.subset(
+        seqid="I", start=40, end=70, allow_partial=True, source=outpath
+    )
+    subset.db.close()
+
+    # reload
+    subset = type(gff_db)(source=outpath)
+    # manual inspection of the original GFF3 file indicates 7 records
+    # BUT the CDS records get merged into a single row
+    assert len(subset) == 6