Typing improvements

Author: jgrg

pyproject.toml

@@ -6,11 +6,11 @@ build-backend = "setuptools.build_meta"
 name = "tola-agp-tpf-utils"
 version = "1.2.1"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
     "click",
     "pyyaml",
-    'importlib-metadata; python_version>="3.10"',
+    'importlib-metadata; python_version>="3.11"',
 ]
 
 [dependency-groups]

src/tola/assembly/assembly.py

@@ -2,19 +2,20 @@
 import re
 import textwrap
 
+from tola.assembly.fragment import Junction
 from tola.assembly.scaffold import Scaffold
 
 
 class Assembly:
     def __init__(
         self, name, header=None, scaffolds=None, bp_per_texel=None, curated=False
     ):
-        self.name = str(name)
-        self.scaffolds = scaffolds if scaffolds else []
-        self.header = header if header else []
-        self.curated = curated
+        self.name: str = str(name)
+        self.scaffolds: list[Scaffold] = scaffolds if scaffolds else []
+        self.header: list[str] = header if header else []
+        self.curated: bool = curated
         if bp_per_texel:
-            self.bp_per_texel = bp_per_texel
+            self.bp_per_texel: float = bp_per_texel
 
     def __repr__(self):
         txt = io.StringIO()
@@ -98,13 +99,13 @@ def name_natural_key(obj):
             for i, x in enumerate(re.split(r"(I+V?|\d+)", obj.name))
         )
 
-    def fragment_junction_set(self):
-        junctions = set()
+    def fragment_junction_set(self) -> set[Junction]:
+        junctions: set[Junction] = set()
         for scffld in self.scaffolds:
             junctions |= scffld.fragment_junction_set()
         return junctions
 
-    def fragment_junctions_by_asm_prefix(self):
+    def fragment_junctions_by_asm_prefix(self) -> dict[str, set[Junction]]:
         prefix_junctions = {}
         for scffld in self.scaffolds:
             try:

src/tola/assembly/build_assembly.py

@@ -1,13 +1,12 @@
 import logging
 import math
-from collections.abc import Iterator
+from typing import TypeAlias
 
 from tola.assembly.assembly import Assembly
 from tola.assembly.assembly_stats import AssemblyStats
 from tola.assembly.build_utils import (
     ChrNamer,
     FoundFragment,
-    OverhangPremise,
     OverhangResolver,
     ScaffoldNamer,
 )
@@ -20,6 +19,9 @@
 log = logging.getLogger(__name__)
 
 
+AssemblyDict: TypeAlias = dict[str | None, Assembly]
+
+
 class BuildAssembly(Assembly):
     """
     Class for building an Assembly from a Pretext Assembly and the
@@ -42,8 +44,8 @@ def __init__(
         self.default_gap = default_gap
         self.found_fragments = {}
         self.fragments_found_more_than_once = {}
-        self.scaffold_namer = ScaffoldNamer()
-        self.assembly_stats = AssemblyStats()
+        self.scaffold_namer: ScaffoldNamer = ScaffoldNamer()
+        self.assembly_stats: AssemblyStats = AssemblyStats()
         if autosome_prefix:
             self.autosome_prefix = autosome_prefix
 
@@ -266,19 +268,17 @@ def add_missing_scaffolds_from_input(self, input_asm: Assembly) -> None:
                 new_scffld.haplotype = scaffold_namer.current_haplotype
                 self.add_scaffold(new_scffld)
 
-    def assembly_with_scaffolds_in_map_order(self) -> dict[None, Assembly]:
-        scaffolds, _ = self.__build_name_and_sort_assemblies()
-        return {
-            None: Assembly("Pretext", scaffolds=scaffolds)
-        }
+    def assembly_with_scaffolds_in_map_order(self) -> AssemblyDict:
+        scaffolds, assemblies = self.__build_name_and_sort_assemblies()
+        return {"map-order": Assembly("Pretext", scaffolds=scaffolds)}
 
-    def assemblies_with_scaffolds_fused(self) -> dict[str | None, Assembly]:
+    def assemblies_with_scaffolds_fused(self) -> AssemblyDict:
         _, assemblies = self.__build_name_and_sort_assemblies()
         return assemblies
 
     def __build_name_and_sort_assemblies(
         self,
-    ) -> tuple[list[Scaffold], dict[str | None, Assembly]]:
+    ) -> tuple[list[Scaffold], AssemblyDict]:
         chr_namer = ChrNamer(chr_prefix=self.autosome_prefix)
 
         scaffolds = self.scaffolds_fused_by_name()

src/tola/assembly/build_utils.py

@@ -5,6 +5,7 @@
 import logging
 import re
 import textwrap
+from abc import ABC, abstractmethod
 
 from tola.assembly.fragment import Fragment
 from tola.assembly.overlap_result import OverlapResult
@@ -476,7 +477,7 @@ def remove_scaffold(self, scaffold: Scaffold) -> None:
         self.scaffolds.remove(scaffold)
 
 
-class OverhangPremise:
+class OverhangPremise(ABC):
     """
     Stores a "what-if" for removal of a terminal (start or end) Fragment. Used
     to decide which OverlapResult to remove a Fragment from, where the
@@ -498,6 +499,32 @@ def __str__(self):
             + textwrap.indent(f"{self.scaffold}\n", "  ")
         )
 
+    @property
+    @abstractmethod
+    def bait_overlap(self) -> int:
+        """
+        The overlap between the bait and the `Fragment` on this end of the
+        `OverlapResult`.  (`0` if they don't overlap.)
+        """
+
+    @property
+    @abstractmethod
+    def overhang_if_applied(self) -> int:
+        """
+        Overhang of the `Fragment` at this end of the `OverlapResult` beyond
+        the `bait` that would be left if `Fragment` at this end was removed.
+        Value is negative if removing the `Fragment` would leave an
+        underhang.
+        """
+
+    @property
+    @abstractmethod
+    def overhang_error_delta_if_applied(self) -> int:
+        """
+        Change (positive or negative) in the absolute size of the overhang
+        (or underhang) if this `OverhangPremise` were applied.
+        """
+
     def improves(self, err_length) -> bool:
         if len(self.scaffold.rows) == 1:
             return False
@@ -557,7 +584,7 @@ class OverhangResolver:
     the OverlapPremises which were applied.
     """
 
-    def __init__(self, error_length=None):
+    def __init__(self, error_length: int):
         self.premises_by_fragment_key: dict[
             tuple[str, int, int], list[OverhangPremise]
         ] = {}

src/tola/assembly/fragment.py

@@ -1,3 +1,8 @@
+from typing import TypeAlias
+
+Junction: TypeAlias = tuple[str | int, str | int, str | int, str | int]
+
+
 class Fragment:
     __slots__ = "_name", "_start", "_end", "_strand", "_tags"
 
@@ -45,7 +50,7 @@ def length(self):
     def key_tuple(self) -> tuple[str, int, int]:
         return self._name, self._start, self._end
 
-    def junction_tuple(self, othr) -> tuple:
+    def junction_tuple(self, othr) -> Junction:
         """
         Encodes the positions of two adjacent Fragments in a Scaffold, with
         reverse strand ends encoded by flipping the order of the name and

src/tola/assembly/parser.py

@@ -129,7 +129,7 @@ def lowercase_and_dash_to_underscore():
     )
 
 
-def format_from_file_extn(pth, default=None):
+def format_from_file_extn(pth, default=None) -> str | None:
     """
     Guess the file format from the extension, or return the supplied default
     """

src/tola/assembly/scaffold.py

@@ -1,6 +1,8 @@
 import io
+from collections.abc import Generator
+from typing import Self
 
-from tola.assembly.fragment import Fragment
+from tola.assembly.fragment import Fragment, Junction
 from tola.assembly.gap import Gap
 
 
@@ -20,11 +22,11 @@ def __init__(
             self.rows = [*rows]
         else:
             self.rows = []
-        self.tag = tag
-        self.haplotype = haplotype
-        self.rank = rank
-        self.original_name = original_name
-        self.original_tags = original_tags
+        self.tag: str | None = tag
+        self.haplotype: str | None = haplotype
+        self.rank: int = rank
+        self.original_name: str | None = original_name
+        self.original_tags: set[str] | None = original_tags
 
     def __repr__(self):
         txt = io.StringIO()
@@ -59,54 +61,57 @@ def add_row(self, row):
         self.rows.append(row)
 
     @property
-    def length(self):
+    def length(self) -> int:
         return sum(r.length for r in self.rows)
 
     @property
-    def fragments_length(self):
+    def fragments_length(self) -> int:
         return sum(f.length for f in self.fragments())
 
     @property
-    def gaps_length(self):
+    def gaps_length(self) -> int:
         return sum(g.length for g in self.gaps())
 
     @property
-    def last_row_is_fragment(self):
+    def last_row_is_fragment(self) -> bool:
         if self.rows:
             return isinstance(self.rows[-1], Fragment)
         else:
             return False
 
-    def fragments(self):
+    def fragments(self) -> Generator[Fragment]:
         for row in self.rows:
             if isinstance(row, Fragment):
                 yield row
 
-    def idx_fragments(self):
+    def idx_fragments(self) -> Generator[tuple[int, Fragment]]:
         for i, row in enumerate(self.rows):
             if isinstance(row, Fragment):
                 yield i, row
 
-    def gaps(self):
+    def gaps(self) -> Generator[Fragment]:
         for row in self.rows:
             if isinstance(row, Gap):
                 yield row
 
-    def idx_gaps(self):
+    def idx_gaps(self) -> Generator[tuple[int, Gap]]:
         for i, row in enumerate(self.rows):
             if isinstance(row, Gap):
                 yield i, row
 
-    def fragment_tags(self):
+    def fragment_tags(self) -> set[str]:
         tag_set = set()
         for frag in self.fragments():
             for t in frag.tags:
                 tag_set.add(t)
         return tag_set
 
-    def reverse(self):
+    def reverse(self) -> Self:
         new = self.__class__(
             self.name,
+            tag=self.tag,
+            haplotype=self.haplotype,
+            rank=self.rank,
             original_name=self.original_name,
             original_tags=self.original_tags,
         )
@@ -121,7 +126,7 @@ def append_scaffold(self, othr, gap=None):
             self.add_row(gap)
         self.rows.extend(othr.rows)
 
-    def fragment_junction_set(self):
+    def fragment_junction_set(self) -> set[Junction]:
         junctions = set()
         itr = self.fragments()
 
@@ -130,12 +135,8 @@ def fragment_junction_set(self):
         except StopIteration:
             return junctions
 
-        while True:
-            try:
-                this = next(itr)
-                junctions.add(prev.junction_tuple(this))
-                prev = this
-            except StopIteration:
-                break
+        for this in itr:
+            junctions.add(prev.junction_tuple(this))
+            prev = this
 
         return junctions