Repository: cogent3/cogent3
Pull request #1869 (2024-05-15 22:32:19)

Improve performance of annotation db creation, querying


Patch:
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 54442a59e..5f8b65769 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -5,9 +5,7 @@ import copy
 import inspect
 import io
 import json
-import os
 import pathlib
-import re
 import sqlite3
 import sys
 import typing
@@ -19,7 +17,9 @@ import typing_extensions
 import cogent3.util.warning as c3warn
 
 from cogent3._version import __version__
+from cogent3.parse.gff import merged_gff_records
 from cogent3.util.deserialise import register_deserialiser
+from cogent3.util.io import PathType, iter_line_blocks
 from cogent3.util.misc import extend_docstring_from, get_object_provenance
 from cogent3.util.progress_display import display_wrap
 from cogent3.util.table import Table
@@ -28,6 +28,7 @@ from cogent3.util.table import Table
 OptionalInt = typing.Optional[int]
 OptionalStr = typing.Optional[str]
 OptionalStrList = typing.Optional[typing.Union[str, typing.List[str]]]
+OptionalStrContainer = typing.Optional[typing.Union[str, typing.Sequence]]
 OptionalBool = typing.Optional[bool]
 OptionalDbCursor = typing.Optional[sqlite3.Cursor]
 ReturnType = typing.Tuple[str, tuple]  # the sql statement and corresponding values
@@ -1124,7 +1125,7 @@ class SqliteAnnotationDbMixin:
     def to_json(self) -> str:
         return json.dumps(self.to_rich_dict())
 
-    def write(self, path: os.PathLike) -> None:
+    def write(self, path: PathType) -> None:
         """writes db as bytes to path"""
         backup = sqlite3.connect(path)
         with self.db:
@@ -1139,7 +1140,7 @@ class SqliteAnnotationDbMixin:
     def subset(
         self,
         *,
-        source: os.PathLike | str = ":memory:",
+        source: PathType = ":memory:",
         biotype: str = None,
         seqid: str = None,
         name: str = None,
@@ -1177,6 +1178,24 @@ class SqliteAnnotationDbMixin:
 
         return result
 
+    def close(self):
+        """closes the db"""
+        self.db.close()
+
+    def _make_index(self, *, table_name: str, col_names: tuple[str, ...]):
+        """index columns for faster search"""
+        sql = f"CREATE INDEX IF NOT EXISTS %s on {table_name}(%s)"
+        for col in col_names:
+            self._execute_sql(sql % (col, col))
+
+    def make_indexes(self):
+        """adds db indexes for core attributes"""
+        for table_name in self.table_names:
+            self._make_index(
+                table_name=table_name,
+                col_names=("biotype", "seqid", "name", "start", "stop", "parent_id"),
+            )
+
 
 class BasicAnnotationDb(SqliteAnnotationDbMixin):
     """Provides a user table for annotations. This can be merged with
@@ -1227,6 +1246,15 @@ class BasicAnnotationDb(SqliteAnnotationDbMixin):
             self._execute_sql(cmnd=cmnd, values=vals)
 
 
+def _merge_spans(old: numpy.ndarray, new: list[list[int]]) -> numpy.ndarray:
+    """returns sorted, merged old and new spans"""
+    if len(new) == old.shape[0] and (old == new).all():
+        return old
+
+    new = numpy.array(sorted(new), dtype=old.dtype)
+    return numpy.unique(numpy.concatenate([old, new]), axis=0)
+
+
 class GffAnnotationDb(SqliteAnnotationDbMixin):
     """Support for annotations from gff files. Records that span multiple
     rows in the gff are merged into a single record."""
@@ -1262,7 +1290,7 @@ class GffAnnotationDb(SqliteAnnotationDbMixin):
         self._serialisable["source"] = self.source
         self._db = None
         self._setup_db(db)
-        data = self._merged_data(data)
+        data, self._num_fakeids = merged_gff_records(data, self._num_fakeids)
         self.add_records(data)
 
     def add_records(self, reduced: dict) -> None:
@@ -1291,51 +1319,27 @@ class GffAnnotationDb(SqliteAnnotationDbMixin):
         self.db.commit()
         del reduced
 
-    def _merged_data(self, records) -> dict:
-        field_template = r"(?<={}=)[^;\s]+"
-        name = re.compile(field_template.format("ID"))
-        parent_id = re.compile(field_template.format("Parent"))
-
-        reduced = collections.OrderedDict()
-        # collapse records with ID's occurring in multiple rows into one
-        # row, converting their coordinates
-        # extract the name from ID and add this into the table
-        # 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")
-            record["stop"] = record.pop("End")
-
-            # we force all keys that map to table column names to be lower case
-            for key in tuple(record):
-                if key.lower() not in self._gff_schema:
-                    continue
-                record[key.lower()] = record.pop(key)
-
-            attrs = record["attributes"] or ""
-            if match := name.search(attrs):
-                record_id = match.group()
-            else:
-                record_id = f"unknown-{self._num_fakeids}"
-                self._num_fakeids += 1
-
-            record["name"] = record_id
-            if pid := parent_id.search(attrs):
-                record["parent_id"] = pid.group()
+    def update_record_spans(self, *, name: str, spans: list[list[int]]) -> None:
+        """updates spans attribute of a gff table record if present
 
-            if record_id not in reduced:
-                reduced[record_id] = record
-                reduced[record_id]["spans"] = []
+        Notes
+        -----
+        Has no effect if name is not present.
+        """
+        if not len(spans):
+            return
 
-            # should this just be an append?
-            reduced[record_id]["spans"].append((record["start"], record["stop"]))
+        result = self._execute_sql(
+            cmnd="SELECT spans from gff WHERE name = ?", values=(name,)
+        ).fetchone()
 
-        del records
+        if result is None:
+            return
 
-        return reduced
+        old_spans = _merge_spans(result["spans"], spans)
+        self._execute_sql(
+            cmnd="UPDATE gff SET spans = ? WHERE name = ?", values=(old_spans, name)
+        )
 
 
 # The GenBank format is less clear on the relationship between identifiers,
@@ -1608,7 +1612,9 @@ def convert_annotation_to_annotation_db(data: dict) -> SupportsFeatures:
 
 
 @display_wrap
-def _db_from_genbank(path: os.PathLike, db: SupportsFeatures, write_path, **kwargs):
+def _db_from_genbank(
+    path: PathType, db: typing.Optional[SupportsFeatures], write_path, **kwargs
+):
     from cogent3 import open_
     from cogent3.parse.genbank import MinimalGenbankParser
 
@@ -1630,6 +1636,8 @@ def _db_from_genbank(path: os.PathLike, db: SupportsFeatures, write_path, **kwar
 
     if not one_valid_path:
         raise IOError(f"{str(path)!r} not found")
+
+    db.make_indexes()
     return db
 
 
@@ -1637,45 +1645,55 @@ def _leave_attributes(*attrs):
     return attrs[0]
 
 
-OptionalStrContainer = typing.Optional[typing.Union[str, typing.Sequence]]
-
-
 @display_wrap
 def _db_from_gff(
-    path: os.PathLike,
+    path: PathType,
     seqids: OptionalStrContainer,
-    db: SupportsFeatures,
-    write_path,
+    db: typing.Optional[SupportsFeatures],
+    write_path: PathType,
+    num_lines: OptionalInt,
     **kwargs,
-):
-    from cogent3.parse.gff import gff_parser
+) -> SupportsFeatures:
+    from cogent3.parse.gff import gff_parser, is_gff3
 
     paths = pathlib.Path(path)
     paths = list(paths.parent.glob(paths.name))
 
     ui = kwargs.pop("ui")
     one_valid_path = False
+    seen_ids = set()
     for path in ui.series(paths):
-        data = list(
-            gff_parser(
-                path,
-                seqids=seqids,
-                attribute_parser=_leave_attributes,
+        num_fake_ids = 0
+        gff3 = is_gff3(path)
+        db = GffAnnotationDb(source=write_path, db=db)
+        for block in iter_line_blocks(path, num_lines=num_lines):
+            data = list(
+                gff_parser(
+                    block, seqids=seqids, attribute_parser=_leave_attributes, gff3=gff3
+                )
             )
-        )
-        db = GffAnnotationDb(source=write_path, data=data, db=db)
+            data, num_fake_ids = merged_gff_records(data, num_fake_ids)
+            if already_seen := seen_ids & data.keys():
+                for name in already_seen:
+                    db.update_record_spans(name=name, spans=data[name].spans)
+
+            seen_ids |= data.keys()
+            db.add_records(data)
         one_valid_path = True
     if not one_valid_path:
         raise IOError(f"{str(path)!r} not found")
+
+    db.make_indexes()
     return db
 
 
 def load_annotations(
     *,
-    path: os.PathLike,
+    path: PathType,
     seqids: OptionalStr = None,
-    db: SupportsFeatures = None,
-    write_path: os.PathLike = ":memory:",
+    db: typing.Optional[SupportsFeatures] = None,
+    write_path: PathType = ":memory:",
+    lines_per_block: OptionalInt = 500_000,
     show_progress: bool = False,
 ) -> SupportsFeatures:
     """loads annotations from flatfile into a db
@@ -1693,6 +1711,9 @@ def load_annotations(
     write_path
         where the constructed database should be written, defaults to
         memory only
+    lines_per_block
+        number of lines to insert into the db per iteration. This can help with
+        memory usage. Only applies to gff files.
     show_progress
         applied only if loading features from multiple files
 
@@ -1714,6 +1735,7 @@ def load_annotations(
             db=db,
             write_path=write_path,
             show_progress=show_progress,
+            num_lines=lines_per_block,
         )
     )
 
@@ -1743,7 +1765,7 @@ def _update_array_format(data: bytes) -> bytes:
 
 
 def update_file_format(
-    source_path: os.PathLike,
+    source_path: PathType,
     db_class: typing.Union[
         type[BasicAnnotationDb],
         type[GenbankAnnotationDb],
@@ -1768,7 +1790,7 @@ def update_file_format(
 
     Parameters
     ----------
-    source_path : os.PathLike
+    source_path : PathType
         The database file to reformat.
     db_class : typing.Union[ type[BasicAnnotationDb], type[GenbankAnnotationDb], type[GffAnnotationDb], ]
         The type of database the file is.
diff --git a/src/cogent3/parse/gff.py b/src/cogent3/parse/gff.py
index 01eacddb9..0aebed826 100644
--- a/src/cogent3/parse/gff.py
+++ b/src/cogent3/parse/gff.py
@@ -1,21 +1,166 @@
-import collections.abc as abc
+from __future__ import annotations
+
+import collections
 import functools
-import os
+import io
+import pathlib
+import re
 import typing
 
-from cogent3.util.io import open_
+from cogent3.util.io import PathType, open_
 
 
 OptionalCallable = typing.Optional[typing.Callable]
-OptionalStrContainer = typing.Optional[typing.Union[str, typing.Sequence]]
+OptionalStrContainer = typing.Optional[typing.Union[str, typing.Sequence[str]]]
+OptionalIntList = typing.Optional[list[list[int]]]
+OptionalStr = typing.Optional[str]
+OptionalInt = typing.Optional[str]
+OptionalStrDict = typing.Optional[typing.Union[str, dict[str, str]]]
+OptionalBool = typing.Optional[bool]
+
+
+@functools.singledispatch
+def is_gff3(f) -> bool:
+    """True if gff-version is 3"""
+    raise TypeError(f"unsupported type {type(f)}")
+
+
+@is_gff3.register
+def _(f: pathlib.PurePath) -> bool:
+    with open_(f) as f:
+        return is_gff3(f)
+
+
+@is_gff3.register
+def _(f: pathlib.Path) -> bool:
+    with open_(f) as f:
+        return is_gff3(f)
+
+
+@is_gff3.register
+def _(f: str) -> bool:
+    with open_(f) as f:
+        return is_gff3(f)
+
+
+@is_gff3.register
+def _(f: io.IOBase) -> bool:
+    gff3_header = "gff-version 3"
+    is_gff3 = gff3_header in f.readline()
+    f.seek(0)
+    return is_gff3
+
+
+@is_gff3.register
+def _(f: list) -> bool:
+    return is_gff3(tuple(f))
+
+
+@is_gff3.register
+def _(f: tuple) -> bool:
+    gff3_header = "gff-version 3"
+    return f and gff3_header in f[0]
+
+
+_gff_item_map = {
+    "type": "biotype",
+    "end": "stop",
+    "attributes": "attrs",
+}
+
+
+class GffRecord:
+    """Container for single GFF record. Elements can be indexed as a dict
+    or directly as attributes.
+
+    Notes
+    -----
+    Some fields defined in the GFF spec are aliased, or truncated.
+    """
+
+    # should be a dataclass, but in py 3.9 they don't support slots
+    # so until then ...
+    __slots__ = (
+        "seqid",
+        "source",
+        "biotype",
+        "name",
+        "parent_id",
+        "start",
+        "stop",
+        "spans",
+        "strand",
+        "score",
+        "phase",
+        "attrs",
+        "comments",
+    )
+
+    def __init__(
+        self,
+        seqid: OptionalStr = None,
+        source: OptionalStr = None,
+        biotype: OptionalStr = None,
+        name: OptionalStr = None,
+        parent_id: OptionalStr = None,
+        start: OptionalInt = None,
+        stop: OptionalInt = None,
+        spans: OptionalIntList = None,
+        score: OptionalStr = None,
+        strand: OptionalStr = None,
+        phase: OptionalInt = None,
+        comments: OptionalStr = None,
+        attrs: OptionalStrDict = None,
+    ):
+        self.seqid = seqid
+        self.source = source
+        self.biotype = biotype
+        self.name = name
+        self.parent_id = parent_id
+        self.start = start
+        self.stop = stop
+        self.spans = spans
+        self.score = score
+        self.strand = strand
+        self.phase = phase
+        self.comments = comments
+        self.attrs = attrs
+
+    def __repr__(self):
+        name = self.__class__.__name__
+        return (
+            f"{name}(seqid={self.seqid!r}, name={self.name!r}, "
+            f"start={self.start}, stop={self.stop}, strand={self.strand!r}, "
+            f"spans={self.spans})"
+        )
+
+    def __getitem__(self, item: str):
+        item = item.lower()
+        return getattr(self, _gff_item_map.get(item, item))
+
+    def __setitem__(self, key: str, value: typing.Any):
+        key = key.lower()
+        setattr(self, _gff_item_map.get(key, key), value)
+
+    def update(self, values: dict[str, typing.Any]):
+        for key, value in values.items():
+            self[key] = value
+
+    def get(self, item: str) -> typing.Any:
+        item = item.lower()
+        return getattr(self, _gff_item_map.get(item, item), None)
+
+    def to_dict(self) -> dict[str, typing.Any]:
+        return {k: getattr(self, k) for k in self.__slots__}
 
 
 @functools.singledispatch
 def gff_parser(
-    f: typing.Union[str, os.PathLike, abc.Sequence, typing.IO],
+    f: typing.Union[PathType, typing.IO, OptionalStrContainer],
     attribute_parser: OptionalCallable = None,
     seqids: OptionalStrContainer = None,
-) -> typing.Iterable[dict]:
+    gff3: OptionalBool = None,
+) -> typing.Iterable[GffRecord]:
     """parses a gff file
 
     Parameters
@@ -26,49 +171,83 @@ def gff_parser(
         callback function for custom handling of attributes field. Defaults to
         a function that converts Attributes field into a dict based on the gff
         version.
+    seqids
+        only gff records matching these sequence IDs are returned
+    gff3
+        True if the format is gff3
 
     Returns
     -------
     dict
         contains each of the 9 parameters specified by gff3, and comments.
     """
-    yield from _gff_parser(f, attribute_parser=attribute_parser, seqids=seqids)
+    gff3 = gff3 or is_gff3(f)
+    yield from _gff_parser(
+        f, attribute_parser=attribute_parser, seqids=seqids, gff3=gff3
+    )
 
 
 @gff_parser.register
 def _(
-    f: str,
+    f: pathlib.PurePath,
     attribute_parser: OptionalCallable = None,
     seqids: OptionalStrContainer = None,
-) -> typing.Iterable[dict]:
+    gff3: OptionalBool = None,
+) -> typing.Iterable[GffRecord]:
     with open_(f) as infile:
-        yield from gff_parser(infile, attribute_parser=attribute_parser, seqids=seqids)
+        yield from gff_parser(
+            infile, attribute_parser=attribute_parser, seqids=seqids, gff3=gff3
+        )
 
 
 @gff_parser.register
 def _(
-    f: os.PathLike,
+    f: pathlib.Path,
     attribute_parser: OptionalCallable = None,
     seqids: OptionalStrContainer = None,
-) -> typing.Iterable[dict]:
+    gff3: OptionalBool = None,
+) -> typing.Iterable[GffRecord]:
     with open_(f) as infile:
-        yield from gff_parser(infile, attribute_parser=attribute_parser, seqids=seqids)
+        yield from gff_parser(
+            infile, attribute_parser=attribute_parser, seqids=seqids, gff3=gff3
+        )
+
+
+@gff_parser.register
+def _(
+    f: str,
+    attribute_parser: OptionalCallable = None,
+    seqids: OptionalStrContainer = None,
+    gff3: OptionalBool = None,
+) -> typing.Iterable[GffRecord]:
+    with open_(f) as infile:
+        yield from gff_parser(
+            infile, attribute_parser=attribute_parser, seqids=seqids, gff3=gff3
+        )
+
+
+@gff_parser.register
+def _(
+    f: io.IOBase,
+    attribute_parser: OptionalCallable = None,
+    seqids: OptionalStrContainer = None,
+    gff3: OptionalBool = None,
+) -> typing.Iterable[GffRecord]:
+    yield from _gff_parser(
+        f, attribute_parser=attribute_parser, seqids=seqids, gff3=gff3
+    )
 
 
 def _gff_parser(
-    f, attribute_parser: OptionalCallable = None, seqids: OptionalStrContainer = None
-) -> typing.Iterable[dict]:
+    f: typing.Union[PathType, typing.IO, OptionalStrContainer],
+    attribute_parser: OptionalCallable = None,
+    seqids: OptionalStrContainer = None,
+    gff3: bool = True,
+) -> typing.Iterable[GffRecord]:
     """parses a gff file"""
     seqids = seqids or set()
     seqids = {seqids} if isinstance(seqids, str) else set(seqids)
 
-    gff3_header = "gff-version 3"
-    if isinstance(f, list):
-        gff3 = f and gff3_header in f[0]
-    else:
-        gff3 = gff3_header in f.readline()
-        f.seek(0)
-
     if attribute_parser is None:
         attribute_parser = parse_attributes_gff3 if gff3 else parse_attributes_gff2
 
@@ -86,7 +265,7 @@ def _gff_parser(
         # the final column (attributes) may be empty
         if len(cols) == 8:
             cols.append("")
-        assert len(cols) == 9, len(line)
+        assert len(cols) == 9, repr(cols)
         seqid, source, type_, start, end, score, strand, phase, attributes = cols
 
         if seqids and seqid not in seqids:
@@ -98,27 +277,24 @@ def _gff_parser(
         # features that extend beyond sequence have negative indices
         if start < 0 or end < 0:
             start, end = abs(start), abs(end)
-            if start > end:
-                start, end = end, start
-        # reverse indices when the feature is on the opposite strand
-        if strand == "-":
-            (start, end) = (end, start)
+        if start > end:
+            start, end = end, start
 
         # all attributes have an "ID" but this may not be unique
         attributes = attribute_parser(attributes, (start, end))
 
-        yield {
-            "SeqID": seqid,
-            "Source": source,
-            "Type": type_,
-            "Start": start,
-            "End": end,
-            "Score": score,
-            "Strand": strand,
-            "Phase": phase,
-            "Attributes": attributes,
-            "Comments": comments,
-        }
+        yield GffRecord(
+            seqid=seqid,
+            source=source,
+            biotype=type_,
+            start=start,
+            stop=end,
+            score=score,
+            strand=strand,
+            phase=phase,
+            attrs=attributes,
+            comments=comments,
+        )
 
 
 def parse_attributes_gff2(attributes: str, span: typing.Tuple[int, int]) -> dict:
@@ -142,3 +318,57 @@ def parse_attributes_gff3(attributes: str, span: typing.Tuple[int, int]) -> dict
             attributes["Parent"] = [attributes["Parent"]]
     attributes["ID"] = attributes.get("ID", "")
     return attributes
+
+
+def merged_gff_records(records: list[GffRecord], num_fake_ids: int) -> tuple[dict, int]:
+    """merges GFF records that have the same ID
+
+    Parameters
+    ----------
+    records
+        list of GFF data rows as dict
+    num_fake_ids
+        starting number for new fakeids
+
+    Returns
+    -------
+    a dictionary with IDs as keys and values being all attributes
+
+    Notes
+    -----
+    Merging means the coordinates are combined into a single "spans"
+    entry. Only records which have an ID field in the attributes get merged.
+    """
+    field_template = r"(?<={}=)[^;\s]+"
+    name = re.compile(field_template.format("ID"))
+    parent_id = re.compile(field_template.format("Parent"))
+
+    reduced = collections.OrderedDict()
+    # collapse records with ID's occurring in multiple rows into one
+    # row, converting their coordinates
+    # extract the name from ID and add this into the table
+    # 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
+    for record in records:
+        attrs = record.attrs or ""
+        if match := name.search(attrs):
+            record_id = match.group()
+        else:
+            record_id = f"unknown-{num_fake_ids}"
+            num_fake_ids += 1
+
+        record.name = record_id
+        if pid := parent_id.search(attrs):
+            record.parent_id = pid.group()
+
+        if record_id not in reduced:
+            reduced[record_id] = record
+            reduced[record_id].spans = []
+
+        # should this just be an append?
+        reduced[record_id].spans.append((record["start"], record["stop"]))
+
+    del records
+
+    return reduced, num_fake_ids
diff --git a/src/cogent3/util/io.py b/src/cogent3/util/io.py
index 8e4d20bf7..f53f65c21 100644
--- a/src/cogent3/util/io.py
+++ b/src/cogent3/util/io.py
@@ -398,7 +398,9 @@ def iter_splitlines(
 
 
 def iter_line_blocks(
-    path: PathType, num_lines: int = 1000, chunk_size: Optional[int] = 1_000_000
+    path: PathType,
+    num_lines: Optional[int] = 1000,
+    chunk_size: Optional[int] = 5_000_000,
 ) -> Iterator[list[str]]:
     """yields list with num_lines str from path
 
@@ -407,7 +409,7 @@ def iter_line_blocks(
     path
         data file
     num_lines
-        number of lines per block
+        number of lines per block. If None just returns all lines.
     chunk_size
         number of bytes to load in one go from path
     """


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index a0910bb13..e22b76a3d 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -288,6 +288,20 @@ def test_load_annotations_multi(DATA_DIR):
     assert len(got) == expect
 
 
+def test_load_annotations_chunked(gff_db, DATA_DIR):
+    path = DATA_DIR / "c_elegans_WS199_shortened_gff.gff3"
+
+    name = "CDS:B0019.1"
+    expect = list(gff_db.get_features_matching(name=name))[0]
+    assert len(expect["spans"]) == 3
+    # two lines splits a 3 line record into 2 and 1 line, so the
+    # update record code is invoked
+    db = load_annotations(path=path, write_path=":memory:", lines_per_block=2)
+    got = list(db.get_features_matching(name=name))[0]
+    assert got.pop("spans") == expect.pop("spans")
+    assert got == expect
+
+
 @pytest.mark.parametrize("parent_biotype, name", (("gene", "CNA00110"),))
 def test_gb_get_children(gb_db, parent_biotype, name):
     parent = list(gb_db.get_features_matching(biotype=parent_biotype, name=name))[0]
@@ -852,12 +866,12 @@ def test_writing_attributes_gff(gff_db):
     gff_db.add_feature(
         biotype="gene",
         seqid="blah",
-        name="cancer-gene",
+        name="agene",
         attributes="description=cancer",
         spans=[(0, 10)],
     )
     r = list(gff_db.get_records_matching(attributes="cancer"))[0]
-    assert r["name"] == "cancer-gene"
+    assert r["name"] == "agene"
 
 
 def test_equal():
@@ -1157,9 +1171,10 @@ def test_gbdb_get_parent_fails_no_coords(gb_db):
         _ = list(gb_db.get_feature_parent(name="CNA00110"))
 
 
-def test_load_annotations_invalid_path():
+@pytest.mark.parametrize("suffix", ("gff3", "gb"))
+def test_load_annotations_invalid_path(suffix):
     with pytest.raises(IOError):
-        load_annotations(path="invalidfile.gff3")
+        load_annotations(path=f"invalidfile.{suffix}")
 
 
 @pytest.mark.parametrize("integer", (int, numpy.int64))
@@ -1207,3 +1222,67 @@ def test_subset_gff3_db_source(gff_db, tmp_dir):
     # manual inspection of the original GFF3 file indicates 7 records
     # BUT the CDS records get merged into a single row
     assert len(subset) == 6
+
+
+@pytest.mark.parametrize(
+    "new_span", ([[22, 24]], [[9, 20]], [[0, 1]], [[9, 20], [29, 45], [59, 70]])
+)
+def test_gffdb_update_record(gff_db, new_span):
+    name = "CDS:B0019.1"
+    init_value = [[9, 20], [29, 45], [59, 70]]
+    combined = set(tuple(c) for c in new_span + init_value)
+    expect = [list(p) for p in sorted(combined)]
+    gff_db.update_record_spans(name=name, spans=new_span)
+    got = list(gff_db.get_records_matching(name=name))[0]
+    assert got["spans"].tolist() == expect
+
+
+def test_gffdb_update_record_empty(gff_db):
+    name = "CDS:B0019.1"
+    init_value = [[9, 20], [29, 45], [59, 70]]
+    gff_db.update_record_spans(name=name, spans=[])
+    got = list(gff_db.get_records_matching(name=name))[0]
+    assert got["spans"].tolist() == init_value
+
+
+def test_gffdb_update_absent_record(gff_db):
+    name = "qwerty"  # does not exist
+    init_value = [[9, 20], [29, 45], [59, 70]]
+    gff_db.update_record_spans(name=name, spans=init_value)
+    got = list(gff_db.get_records_matching(name=name))
+    assert not got
+
+
+@pytest.mark.parametrize("db", ("gff_db", "gb_db"))
+def test_db_close(request, db):
+    import sqlite3
+
+    db = request.getfixturevalue(db)
+    db.close()
+    with pytest.raises(sqlite3.ProgrammingError):
+        _ = list(db.get_features_matching(biotype="gene"))
+
+
+@pytest.mark.parametrize("db", ("gff_db", "gb_db"))
+@pytest.mark.parametrize(
+    "col", ("biotype", "seqid", "name", "start", "stop", "parent_id")
+)
+def test_db_make_index(request, db, col):
+    ann_db = request.getfixturevalue(db)
+    expect = {("index", col, tn) for tn in ann_db.table_names}
+    ann_db.make_indexes()
+    sql_template = (
+        "SELECT * FROM sqlite_master WHERE type = 'index' AND"  # nosec B608
+        f" tbl_name = '%s' and name = {col!r}"  # nosec B608
+    )
+
+    result = ann_db._execute_sql(sql_template % ann_db.table_names[0]).fetchone()
+    got = tuple(result)[:3]
+    assert got in expect
+
+
+@pytest.mark.parametrize("db", ("gff_db", "gb_db"))
+def test_db_repr(request, db):
+    ann_db = request.getfixturevalue(db)
+    got = repr(ann_db)
+    assert isinstance(got, str)
diff --git a/tests/test_parse/test_gff.py b/tests/test_parse/test_gff.py
index 6f52734f3..dac10325b 100644
--- a/tests/test_parse/test_gff.py
+++ b/tests/test_parse/test_gff.py
@@ -86,28 +86,6 @@ data_lines = [
 class GffTest(TestCase):
     """Setup data for all the GFF parsers."""
 
-    def testGffParserData(self):
-        """Test gff_parser with valid data lines"""
-        for line, canned_result in data_lines:
-            result = next(gff_parser(StringIO(line)))
-            canned_result = list(canned_result)
-            self.assertEqual(result.pop("Attributes")["Info"], canned_result.pop(8))
-            self.assertEqual(set(result.values()), set(canned_result))
-
-    def test_gff_parser_headers(self):
-        """Test gff_parser with valid data headers"""
-        data = "".join([x[0] for x in data_lines])
-        for header in headers:
-            result = list(gff_parser(StringIO(header + data)))
-            lines = [(x[0], list(x[1])) for x in data_lines]
-            self.assertEqual(
-                [l.pop("Attributes")["Info"] for l in result],
-                [x[1].pop(8) for x in lines],
-            )
-            self.assertEqual(
-                [set(l.values()) for l in result], [set(x[1]) for x in lines]
-            )
-
     def test_parse_attributes_gff2(self):
         """Test the parse_attributes_gff2 method"""
         self.assertEqual(
@@ -118,30 +96,6 @@ class GffTest(TestCase):
             ["HBA_HUMAN", "dJ102G20.C1.1", "", "BROADO5"],
         )
 
-    def test_gff2_parser_string(self):
-        """Test the gff_parser works with a string filepath"""
-        filepath = os.path.join("data/gff2_test.gff")
-        for i, result in enumerate(gff_parser(filepath)):
-            line = list(data_lines[i][1])
-            self.assertEqual(result.pop("Attributes")["Info"], line.pop(8))
-            self.assertEqual(set(result.values()), set(line))
-
-    def test_gff2_parser_path(self):
-        """Test the gff_parser works with a pathlib.Path filepath"""
-        filepath = Path("data/gff2_test.gff")
-        for i, result in enumerate(gff_parser(filepath)):
-            line = list(data_lines[i][1])
-            self.assertEqual(result.pop("Attributes")["Info"], line.pop(8))
-            self.assertEqual(set(result.values()), set(line))
-
-    def test_gff3_parser(self):
-        """Test the gff_parser works on a gff3 file"""
-        gff3_path = os.path.join("data/c_elegans_WS199_shortened_gff.gff3")
-        for i, result in enumerate(gff_parser(gff3_path)):
-            self.assertEqual(len(result), 10)
-        # 15 total lines, but 2 comments
-        self.assertEqual(i + 1, 15 - 2)
-
     def test_custom_attr_func(self):
         """user provided attr parser"""
         gff3_path = os.path.join("data/c_elegans_WS199_shortened_gff.gff3")
@@ -193,6 +147,66 @@ def test_parse_field_spaces(DATA_DIR):
     path = DATA_DIR / "simple.gff"
     got = list(gff_parser(path))
     for record in got:
-        for value in record.values():
+        for attr in dir(record):
+            if attr.startswith("_"):
+                continue
+            value = getattr(record, attr)
             if isinstance(value, str):
-                assert value.strip() == value, "should not have spaces!"
+                assert value.strip() == value, f"{attr} should not have spaces!"
+
+
+@pytest.mark.parametrize("line,canned_result", data_lines)
+def test_gff_parser_data(line, canned_result):
+    """Test gff_parser with valid data lines"""
+    result = next(gff_parser(StringIO(line))).to_dict()
+    canned_result = list(canned_result)
+    assert result["attrs"]["Info"] == canned_result.pop(8)
+    result.pop("attrs")
+    assert set(result.values()) == set(canned_result)
+
+
+def test_gff2_parser_path():
+    """Test the gff_parser works with a pathlib.Path filepath"""
+    filepath = Path("data/gff2_test.gff")
+    for i, result in enumerate(gff_parser(filepath)):
+        result = result.to_dict()
+        line = list(data_lines[i][1])
+        assert result.pop("attrs")["Info"] == line.pop(8)
+        assert set(result.values()) == set(line)
+
+
+def test_gff2_parser_string():
+    """Test the gff_parser works with a string filepath"""
+    filepath = os.path.join("data/gff2_test.gff")
+    for i, result in enumerate(gff_parser(filepath)):
+        result = result.to_dict()
+        line = list(data_lines[i][1])
+        assert result.pop("attrs")["Info"] == line.pop(8)
+        assert set(result.values()) == set(line)
+
+
+def test_gff3_parser():
+    """Test the gff_parser works on a gff3 file"""
+    gff3_path = os.path.join("data/c_elegans_WS199_shortened_gff.gff3")
+    records = list(gff_parser(gff3_path))
+    lengths = {len(result.to_dict()) for result in records}
+    # add 3 for name, parent_id and comments
+    assert lengths == {10 + 3}
+    # 15 total lines, but 2 comments
+    assert len(records) == 15 - 2
+
+
+def test_gff_parser_headers():
+    """Test gff_parser with valid data headers"""
+    data = "".join([x[0] for x in data_lines])
+    for header in headers:
+        result = list(gff_parser(StringIO(header + data)))
+        lines = [(x[0], list(x[1])) for x in data_lines]
+        assert [l.attrs["Info"] for l in result] == [x[1].pop(8) for x in lines]
+        expect = [set(x[1]) for x in lines]
+        got = []
+        for l in result:
+            l = l.to_dict()
+            l.pop("attrs")
+            got.append(set(l.values()))
+        assert got == expect
diff --git a/tests/test_util/test_io.py b/tests/test_util/test_io.py
index f4739f668..b952ee4a1 100644
--- a/tests/test_util/test_io.py
+++ b/tests/test_util/test_io.py
@@ -403,7 +403,7 @@ def test_iter_splitlines_one(tmp_path):
 def test_iter_splitlines_line_diff_newline(tmp_path, newline):
     path = tmp_path / "multi-line.txt"
     value = ["We have some", "text on different lines", "which load"]
-    with open(path, mode="w", newline=newline) as out:
+    with open_(path, mode="w", newline=newline) as out:
         out.write("\n".join(value))
     # we use a very small chunk size
     got = list(iter_splitlines(path, chunk_size=5))
@@ -414,8 +414,8 @@ def test_iter_splitlines_line_diff_newline(tmp_path, newline):
 def test_iter_splitlines_file_endswith_newline(tmp_path, newline):
     path = tmp_path / "multi-line.txt"
     value = ["We have some", "text on different lines", "which load"]
-    with open(path, mode="w", newline=newline) as out:
-        out.write("\n".join(value))
+    with open_(path, mode="w", newline=newline) as out:
+        out.write("\n".join(value) + "\n")
     # we use a very small chunk size
     got = list(iter_splitlines(path, chunk_size=5))
     assert got == value