From 9c0d382ebc8aa5017f24508eb4396cbc28c0e8eb Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 11:57:07 +0200 Subject: [PATCH 01/11] chore: Some typing-related improvements --- augur/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/augur/utils.py b/augur/utils.py index a131af604..fb3d554c9 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -1,4 +1,5 @@ import argparse +from typing import Optional import Bio import Bio.Phylo import numpy as np @@ -406,7 +407,7 @@ def get_augur_version(): return __version__ -def read_bed_file(bed_file): +def read_bed_file(bed_file: str) -> list[int]: """Read a BED file and return a list of excluded sites. This function attempts to parse the given file as a BED file, based on @@ -441,7 +442,7 @@ def read_bed_file(bed_file): Sorted list of unique zero-indexed sites """ in_header = True - initial_chrom_value: str | None = None + initial_chrom_value: Optional[str] = None mask_sites: list[int] = [] bed_file_size = os.path.getsize(bed_file) From 8c56f85ecabdec33edf2c408484d534d9e0e5345 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 12:25:24 +0200 Subject: [PATCH 02/11] Add typing to node_data_read --- augur/util_support/node_data_reader.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/augur/util_support/node_data_reader.py b/augur/util_support/node_data_reader.py index 2b9105c73..7470639d7 100644 --- a/augur/util_support/node_data_reader.py +++ b/augur/util_support/node_data_reader.py @@ -1,9 +1,10 @@ -import Bio.Phylo import sys +from typing import Optional, Union + +import Bio.Phylo from augur.types import ValidationMode -from augur.util_support.node_data import DuplicatedNonDictAttributeError -from augur.util_support.node_data import NodeData +from augur.util_support.node_data import DuplicatedNonDictAttributeError, NodeData from augur.util_support.node_data_file import NodeDataFile @@ -19,21 +20,21 @@ class NodeDataReader: If validation_mode is set to :py:attr:`augur.types.ValidationMode.SKIP` no validation is performed. """ - def __init__(self, filenames, tree_file=None, validation_mode=ValidationMode.ERROR): + def __init__(self, filenames: Union[str, list[str]], tree_file: Optional[str] =None, validation_mode=ValidationMode.ERROR) -> None: if not isinstance(filenames, list): filenames = [filenames] self.filenames = filenames self.tree_file = tree_file self.validation_mode = validation_mode - def read(self): + def read(self) -> dict: node_data = self.build_node_data() self.check_against_tree_file(node_data) return node_data - def build_node_data(self): + def build_node_data(self) -> dict: node_data = NodeData() for node_data_file in self.node_data_files: @@ -54,7 +55,7 @@ def build_node_data(self): def node_data_files(self): return (NodeDataFile(fname, validation_mode = self.validation_mode) for fname in self.filenames) - def check_against_tree_file(self, node_data): + def check_against_tree_file(self, node_data: dict) -> None: if not self.tree_file: return @@ -67,7 +68,7 @@ def check_against_tree_file(self, node_data): sys.exit(2) @property - def node_names_from_tree_file(self): + def node_names_from_tree_file(self) -> set[str]: try: tree = Bio.Phylo.read(self.tree_file, "newick") except Exception as e: From 2434191d3b45f6daa71c88beec6923b4c345e66c Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:21:06 +0200 Subject: [PATCH 03/11] Add type hints for all utils functions and some more in other modules, start BioPython stubs --- augur/io/sequences.py | 47 ++++++---- augur/util_support/color_parser.py | 13 +-- augur/utils.py | 137 +++++++++++++++------------ mypy.ini | 3 + stubs/Bio/SeqFeature.pyi | 143 +++++++++++++++++++++++++++++ stubs/Bio/__init__.pyi | 1 + 6 files changed, 261 insertions(+), 83 deletions(-) create mode 100644 stubs/Bio/SeqFeature.pyi create mode 100644 stubs/Bio/__init__.pyi diff --git a/augur/io/sequences.py b/augur/io/sequences.py index 0e75922d8..5d7cc53ca 100644 --- a/augur/io/sequences.py +++ b/augur/io/sequences.py @@ -1,15 +1,19 @@ -import Bio.SeqIO import os - -from augur.errors import AugurError -from augur.utils import augur from importlib.metadata import version as installed_version -from packaging.version import Version from shlex import quote as shquote from shutil import which from tempfile import NamedTemporaryFile from textwrap import dedent -from typing import Iterator, Iterable, Union +from typing import Iterable, Iterator, Optional, Union + +import Bio +from Bio.SeqFeature import SeqFeature +from Bio.SeqRecord import SeqRecord +from packaging.version import Version + +from augur.errors import AugurError +from augur.utils import augur + from .file import open_file from .print import _n, indented_list from .shell_command_runner import run_shell_command @@ -186,12 +190,12 @@ def subset_fasta(input_filename: str, output_filename: str, ids_file: str, nthre if os.path.isfile(output_filename): # Remove the partial output file. os.remove(output_filename) - raise AugurError(f"Sequence output failed, see error(s) above.") + raise AugurError("Sequence output failed, see error(s) above.") else: - raise AugurError(f"Sequence output failed, see error(s) above. The command may have already written data to stdout. You may want to clean up any partial outputs.") + raise AugurError("Sequence output failed, see error(s) above. The command may have already written data to stdout. You may want to clean up any partial outputs.") -def load_features(reference, feature_names=None): +def load_features(reference: str, feature_names: Optional[Union[set[str], list[str]]] = None) -> dict: """ Parse a GFF/GenBank reference file. See the docstrings for _read_gff and _read_genbank for details. @@ -254,7 +258,7 @@ def _read_nuc_annotation_from_gff(record, reference): if len(sequence_regions)>1: raise AugurError(f"Reference {reference!r} contains multiple ##sequence-region pragma lines. Augur can only handle GFF files with a single one.") elif sequence_regions: - from Bio.SeqFeature import SeqFeature, FeatureLocation + from Bio.SeqFeature import FeatureLocation, SeqFeature (name, start, stop) = sequence_regions[0] nuc['pragma'] = SeqFeature( FeatureLocation(start, stop, strand=1), @@ -285,7 +289,7 @@ def _read_nuc_annotation_from_gff(record, reference): raise AugurError(f"Reference {reference!r} didn't define any information we can use to create the 'nuc' annotation. You can use a line with a 'record' or 'source' GFF type or a ##sequence-region pragma.") -def _read_gff(reference, feature_names): +def _read_gff(reference: str, feature_names: Optional[Union[set[str], list[str]]] = None) -> dict[str, SeqFeature]: """ Read a GFF file. We only read GFF IDs 'gene' or 'source' (the latter may not technically be a valid GFF field, but is used widely within the Nextstrain ecosystem). @@ -325,6 +329,7 @@ def _read_gff(reference, feature_names): # TODO: Remove warning suppression after it's addressed upstream: # import warnings + from Bio import BiopythonDeprecationWarning warnings.simplefilter("ignore", BiopythonDeprecationWarning) gff_entries = list(GFF.parse(in_handle, limit_info={'gff_type': valid_types})) @@ -377,7 +382,7 @@ def _read_gff(reference, feature_names): return features -def _read_nuc_annotation_from_genbank(record, reference): +def _read_nuc_annotation_from_genbank(record: SeqRecord, reference: str): """ Extracts the mandatory 'source' feature. If the sequence is present we check the length agrees with the source. (The 'ORIGIN' may be left blank, @@ -387,7 +392,8 @@ def _read_nuc_annotation_from_genbank(record, reference): Parameters ---------- - record : :py:class:`Bio.SeqRecord.SeqRecord` reference: string + record : :py:class:`Bio.SeqRecord.SeqRecord` + reference : string Returns ------- @@ -411,7 +417,7 @@ def _read_nuc_annotation_from_genbank(record, reference): return nuc -def _read_genbank(reference, feature_names): +def _read_genbank(reference: str, feature_names: Optional[Union[set[str], list[str]]] = None) -> dict[str, SeqFeature]: """ Read a GenBank file. We only read GenBank feature keys 'CDS' or 'source'. We create a "feature name" via: @@ -443,13 +449,14 @@ def _read_genbank(reference, feature_names): } features_skipped = 0 - for feat in gb.features: - if feat.type=='CDS': + for feature in gb.features: + feat = feature + if feat.type=='CDS': # type: ignore[has-type] fname = None - if "locus_tag" in feat.qualifiers: - fname = feat.qualifiers["locus_tag"][0] - elif "gene" in feat.qualifiers: - fname = feat.qualifiers["gene"][0] + if "locus_tag" in feat.qualifiers: # type: ignore[has-type] + fname = feat.qualifiers["locus_tag"][0] # type: ignore[has-type] + elif "gene" in feat.qualifiers: # type: ignore[has-type] + fname = feat.qualifiers["gene"][0] # type: ignore[has-type] else: features_skipped+=1 diff --git a/augur/util_support/color_parser.py b/augur/util_support/color_parser.py index 1f8c9360f..16008bec7 100644 --- a/augur/util_support/color_parser.py +++ b/augur/util_support/color_parser.py @@ -1,5 +1,6 @@ -from collections import defaultdict import functools +from collections import defaultdict +from typing import Dict, List, Optional, TextIO, Tuple from augur.data import as_file from augur.io.file import open_file @@ -7,14 +8,14 @@ class ColorParser: - def __init__(self, *, mapping_filename, use_defaults=True): + def __init__(self, *, mapping_filename: Optional[str], use_defaults: bool = True) -> None: self.mapping_filename = mapping_filename self.use_defaults = use_defaults @property @functools.lru_cache() - def mapping(self): - colors = {} + def mapping(self) -> Dict[str, List[Tuple[str, str]]]: + colors: Dict[str, List[Tuple[str, str]]] = {} if self.use_defaults: with as_file("colors.tsv") as file: @@ -27,8 +28,8 @@ def mapping(self): return colors - def parse_file(self, file): - file_mapping = defaultdict(list) + def parse_file(self, file: TextIO) -> Dict[str, List[Tuple[str, str]]]: + file_mapping: Dict[str, List[Tuple[str, str]]] = defaultdict(list) for pair in [ColorParserLine(line).pair() for line in file]: if pair is None: continue diff --git a/augur/utils.py b/augur/utils.py index fb3d554c9..f4d2d4d5f 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -1,27 +1,35 @@ import argparse -from typing import Optional +import json +import os +import sys +import warnings +from collections import Counter, OrderedDict +from io import RawIOBase +from shlex import quote as shquote +from typing import Any, Dict, List, Optional, Tuple, Union + import Bio import Bio.Phylo +import Bio.Phylo.BaseTree +import Bio.Phylo.Newick import numpy as np -import os, json, sys import pandas as pd -from collections import OrderedDict -from io import RawIOBase -from shlex import quote as shquote -from .__version__ import __version__ +from Bio.Phylo.BaseTree import Clade, Tree +from Bio.SeqFeature import CompoundLocation, FeatureLocation, SeqFeature +from typing_extensions import Buffer from augur.data import as_file +from augur.errors import AugurError from augur.io.file import open_file -from augur.io.print import print_err - from augur.types import ValidationMode -from augur.errors import AugurError - from augur.util_support.color_parser import ColorParser from augur.util_support.node_data_reader import NodeDataReader +from .__version__ import __version__ + +TreeLike = Union[Tree, Clade] -def augur(): +def augur() -> str: """ Locate how to re-invoke ourselves (_this_ specific Augur). """ @@ -30,15 +38,15 @@ def augur(): else: # A bit unusual we don't know our own Python executable, but assume we # can access ourselves as the ``augur`` command. - return f"augur" + return "augur" -def get_json_name(args, default=None): +def get_json_name(args: argparse.Namespace, default: Optional[str] = None) -> str: if args.output_node_data: return args.output_node_data else: if default: - print("WARNING: no name for the output file was specified. Writing results to %s."%default, file=sys.stderr) + print(f"WARNING: no name for the output file was specified. Writing results to {default}.", file=sys.stderr) return default else: raise ValueError("Please specify a name for the JSON file containing the results.") @@ -50,7 +58,7 @@ class InvalidTreeError(Exception): pass -def read_tree(fname, min_terminals=3): +def read_tree(fname: str, min_terminals: int = 3) -> Bio.Phylo.BaseTree.Tree: """Safely load a tree from a given filename or raise an error if the file does not contain a valid tree. @@ -101,11 +109,11 @@ def read_tree(fname, min_terminals=3): return T -def read_node_data(fnames, tree=None, validation_mode=ValidationMode.ERROR): +def read_node_data(fnames: Union[str, list[str]], tree: Optional[str] = None, validation_mode: ValidationMode = ValidationMode.ERROR) -> dict: return NodeDataReader(fnames, tree, validation_mode).read() -def write_json(data, file, indent=(None if os.environ.get("AUGUR_MINIFY_JSON") else 2), include_version=True): +def write_json(data: dict, file, indent: int = (None if os.environ.get("AUGUR_MINIFY_JSON") else 2), include_version: bool=True) -> None: """ Write ``data`` as JSON to the given ``file``, creating parent directories if necessary. The augur version is included as a top-level key "augur_version". @@ -145,17 +153,17 @@ def write_json(data, file, indent=(None if os.environ.get("AUGUR_MINIFY_JSON") e class BytesWrittenCounterIO(RawIOBase): """Binary stream to count the number of bytes sent via write().""" - def __init__(self): - self.written = 0 + def __init__(self) -> None: + self.written: int = 0 """Number of bytes written.""" - def write(self, b): - n = len(b) + def write(self, b: Buffer) -> int: + n = memoryview(b).nbytes self.written += n return n -def json_size(data): +def json_size(data: dict) -> int: """Return size in bytes of a Python object in JSON string form.""" with BytesWrittenCounterIO() as counter: write_json(data, counter, include_version=False) @@ -167,7 +175,7 @@ class AugurJSONEncoder(json.JSONEncoder): A custom JSONEncoder subclass to serialize data types used for various data stored in dictionary format. """ - def default(self, obj): + def default(self, obj: object) -> Union[int, float, list]: if isinstance(obj, np.integer): return int(obj) if isinstance(obj, np.floating): @@ -179,19 +187,25 @@ def default(self, obj): return super().default(obj) -def load_features(*args, **kwargs): - print_err("WARNING: augur.utils.load_features has been moved to augur.io.sequences.load_features.") +def load_features(reference: str, feature_names: Optional[Union[set[str], list[str]]] = None) -> dict: + """Deprecated function, moved to augur.io.sequences.load_features.""" + warnings.warn( + "augur.utils.load_features has been moved to augur.io.sequences.load_features. " + "Please update your imports. This function will be removed in a future version.", + FutureWarning, + stacklevel=2 + ) from augur.io.sequences import load_features as new_location - return new_location(*args, **kwargs) + return new_location(reference=reference, feature_names=feature_names) -def read_lat_longs(overrides=None, use_defaults=True): +def read_lat_longs(overrides: Optional[str] = None, use_defaults: bool = True) -> dict[Tuple[str,str], dict[str, float]]: coordinates = {} # TODO: make parsing of tsv files more robust while allow for whitespace delimiting for backwards compatibility def add_line_to_coordinates(line): if line.startswith('#') or line.strip() == "": return - fields = line.strip().split() if not '\t' in line else line.strip().split('\t') + fields = line.strip().split() if '\t' not in line else line.strip().split('\t') if len(fields) == 4: geo_field, loc = fields[0].lower(), fields[1].lower() lat, long = float(fields[2]), float(fields[3]) @@ -215,10 +229,10 @@ def add_line_to_coordinates(line): print("WARNING: input lat/long file %s not found." % overrides) return coordinates -def read_colors(overrides=None, use_defaults=True): +def read_colors(overrides: Optional[str] = None, use_defaults: bool = True) -> Dict[str, List[Tuple[str, str]]]: return ColorParser(mapping_filename=overrides, use_defaults=use_defaults).mapping -def first_line(text): +def first_line(text: str) -> str: """ Returns the first line of the given text, ignoring leading and trailing whitespace. @@ -251,7 +265,7 @@ def available_cpu_cores(fallback: int = 1) -> int: return os.cpu_count() or fallback -def nthreads_value(value): +def nthreads_value(value: str) -> int: """ Argument value validation and casting function for --nthreads. """ @@ -265,7 +279,7 @@ def nthreads_value(value): raise argparse.ArgumentTypeError("'%s' is not an integer or the word 'auto'" % value) from None -def get_parent_name_by_child_name_for_tree(tree): +def get_parent_name_by_child_name_for_tree(tree: Bio.Phylo.BaseTree.Tree) -> Dict[str, str]: ''' Return dictionary mapping child node names to parent node names ''' @@ -276,7 +290,7 @@ def get_parent_name_by_child_name_for_tree(tree): return parents -def annotate_parents_for_tree(tree): +def annotate_parents_for_tree(tree: Bio.Phylo.BaseTree.Tree) -> Bio.Phylo.BaseTree.Tree: """Annotate each node in the given tree with its parent. Examples @@ -300,7 +314,11 @@ def annotate_parents_for_tree(tree): return tree -def json_to_tree(json_dict, root=True, parent_cumulative_branch_length=None): +def json_to_tree( + json_dict: dict, + root: bool = True, + parent_cumulative_branch_length: Optional[Union[float, int]] = None + ) -> Bio.Phylo.BaseTree.Clade: """Returns a Bio.Phylo tree corresponding to the given JSON dictionary exported by `tree_to_json`. @@ -357,7 +375,7 @@ def json_to_tree(json_dict, root=True, parent_cumulative_branch_length=None): if root and "meta" in json_dict and "tree" in json_dict: json_dict = json_dict["tree"] - node = Bio.Phylo.Newick.Clade() + node: Bio.Phylo.Newick.Clade = Bio.Phylo.Newick.Clade() # v1 and v2 JSONs use different keys for strain names. if "name" in json_dict: @@ -372,13 +390,13 @@ def json_to_tree(json_dict, root=True, parent_cumulative_branch_length=None): # Only v1 JSONs support a single `attr` attribute. if hasattr(node, "attr"): - node.numdate = node.attr.get("num_date") - node.cumulative_branch_length = node.attr.get("div") + node.numdate = node.attr.get("num_date") # type: ignore + node.cumulative_branch_length = node.attr.get("div") # type: ignore if "translations" in node.attr: - node.translations = node.attr["translations"] + node.translations = node.attr["translations"] # type: ignore elif hasattr(node, "node_attrs"): - node.cumulative_branch_length = node.node_attrs.get("div") + node.cumulative_branch_length = node.node_attrs.get("div") # type: ignore[attr-defined] node.branch_length = 0.0 if parent_cumulative_branch_length is not None and hasattr(node, "cumulative_branch_length"): @@ -390,17 +408,17 @@ def json_to_tree(json_dict, root=True, parent_cumulative_branch_length=None): json_to_tree( child, root=False, - parent_cumulative_branch_length=node.cumulative_branch_length + parent_cumulative_branch_length=node.cumulative_branch_length # type: ignore[attr-defined] ) for child in json_dict["children"] ] if root: - node = annotate_parents_for_tree(node) + node = annotate_parents_for_tree(node) # type: ignore return node -def get_augur_version(): +def get_augur_version() -> str: """ Returns a string of the current augur version. """ @@ -497,7 +515,7 @@ def read_bed_file(bed_file: str) -> list[int]: return sorted(set(mask_sites)) -def read_mask_file(mask_file): +def read_mask_file(mask_file: str) -> list[int]: """Read a masking file and return a list of excluded sites. Masking files have a single masking site per line, either alone @@ -528,7 +546,7 @@ def read_mask_file(mask_file): raise return sorted(set(mask_sites)) -def load_mask_sites(mask_file): +def load_mask_sites(mask_file: str) -> list[int]: """Load masking sites from either a BED file or a masking file. Parameters @@ -554,7 +572,7 @@ def load_mask_sites(mask_file): } -def read_entries(*files, comment_char="#"): +def read_entries(*files: str, comment_char: str = "#") -> list[str]: """Reads entries (one per line) from one or more plain text files. Entries can be commented with full-line or inline comments. For example, the @@ -588,7 +606,7 @@ def read_entries(*files, comment_char="#"): return entries -def parse_genes_argument(input): +def parse_genes_argument(input: Optional[Union[str, List[str]]] = None) -> Union[str ,list[str], None]: if input is None: return None @@ -600,23 +618,30 @@ def parse_genes_argument(input): return input -def _get_genes_from_file(fname): +def _get_genes_from_file(fname: str) -> list[str]: if os.path.isfile(fname): genes = read_entries(fname) else: print("File with genes not found. Looking for", fname) genes = [] - unique_genes = np.unique(np.array(genes)) - if len(unique_genes) != len(genes): - print("You have duplicates in your genes file. They are being ignored.") + gene_counts = Counter(genes) + duplicates = [gene for gene, count in gene_counts.items() if count > 1] + + unique_genes = list(dict.fromkeys(genes)) # Preserves order while removing duplicates + if duplicates: + print(f"Warning: You have duplicates in your genes file. The following genes appear multiple times: {', '.join(duplicates)}. Duplicates are being ignored.") print("Read in {} specified genes to translate.".format(len(unique_genes))) return unique_genes -def genome_features_to_auspice_annotation(features, ref_seq_name=None, assert_nuc=False): +def genome_features_to_auspice_annotation( + features: Dict[str, SeqFeature], + ref_seq_name: Optional[str] = None, + assert_nuc: bool = False + ) -> Dict[str, Any]: """ Parameters ---------- @@ -633,19 +658,17 @@ def genome_features_to_auspice_annotation(features, ref_seq_name=None, assert_nu See schema-annotations.json for the schema this conforms to """ - from Bio.SeqFeature import SimpleLocation, CompoundLocation - if assert_nuc and 'nuc' not in features: raise AugurError("Genome features must include a feature for 'nuc'") - def _parse(feat): - a = {} + def _parse(feat: SeqFeature) -> Dict[str, Any]: + a: Dict[str, Any] = {} # Note that BioPython locations use "Pythonic" coordinates: [zero-origin, half-open) # Starting with augur v6 we use GFF coordinates: [one-origin, inclusive] - if type(feat.location)==SimpleLocation: + if isinstance(feat.location, FeatureLocation): a['start'] = int(feat.location.start)+1 a['end'] = int(feat.location.end) - elif type(feat.location)==CompoundLocation: + elif isinstance(feat.location, CompoundLocation): a['segments'] = [ {'start':int(segment.start)+1, 'end':int(segment.end)} for segment in feat.location.parts # segment: SimpleLocation diff --git a/mypy.ini b/mypy.ini index 873e5c085..b110ae6e3 100644 --- a/mypy.ini +++ b/mypy.ini @@ -10,6 +10,9 @@ files = augur/ # potentially returning None (via Optional[…]). strict_optional = False +# Directory for stubs +mypy_path = stubs + # In the future maybe we can contribute typing stubs for these modules (either # complete stubs in the python/typeshed repo or partial stubs just in # this repo), but for now that's more work than we want to invest. These diff --git a/stubs/Bio/SeqFeature.pyi b/stubs/Bio/SeqFeature.pyi new file mode 100644 index 000000000..43d04eca8 --- /dev/null +++ b/stubs/Bio/SeqFeature.pyi @@ -0,0 +1,143 @@ +# Generated by ChatGPT o3 based on code of SeqFeature.py +# https://github.com/biopython/biopython/blob/10bb67863892d23cfa611cabf9cd93c1e1e83ef9/Bio/SeqFeature.py + +from __future__ import annotations + +from typing import Dict, Iterator, List, Optional, Union + +from Bio.Seq import MutableSeq, Seq + +# --------------------------------------------------------------------------- # +# Convenience aliases # +# --------------------------------------------------------------------------- # + +Strand = Optional[int] # +1, -1, 0, or None +SeqLike = Union[str, Seq, MutableSeq] # DNA/RNA/protein or plain string + + +# --------------------------------------------------------------------------- # +# Position objects (minimal public surface – all behave like `int`) # +# --------------------------------------------------------------------------- # + +class Position(int): ... +class ExactPosition(Position): ... +class BeforePosition(Position): ... +class AfterPosition(Position): ... +class WithinPosition(Position): ... +class BetweenPosition(Position): ... +class OneOfPosition(Position): ... +class UnknownPosition(Position): ... +class UncertainPosition(Position): ... + + +# --------------------------------------------------------------------------- # +# Location hierarchy # +# --------------------------------------------------------------------------- # + +class Location: + start: Position + end: Position + strand: Strand + + def extract( + self, + parent_sequence: SeqLike, + references: Optional[Dict[str, Seq]] = ..., + ) -> SeqLike: ... + + # convenience dunders ---------------------------------------------------- # + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[int]: ... + def __contains__(self, value: int) -> bool: ... + + # internal helpers (return shifted / flipped copies) --------------------- # + def _shift(self, offset: int) -> "Location": ... + def _flip(self, length: int) -> "Location": ... + + +class SimpleLocation(Location): + parts: List["SimpleLocation"] + + def __init__( + self, + start: Union[int, Position], + end: Union[int, Position], + strand: Strand = ..., + *, + ref: Optional[str] = ..., + ref_db: Optional[str] = ..., + ) -> None: ... + + +class CompoundLocation(Location): + parts: List[SimpleLocation] + operator: str + + def __init__(self, parts: List[SimpleLocation], operator: str = ...) -> None: ... + + +# Retro-compat alias used throughout Biopython +FeatureLocation = SimpleLocation + + +# --------------------------------------------------------------------------- # +# SeqFeature – high-level wrapper around a Location # +# --------------------------------------------------------------------------- # + +class SeqFeature: + location: Union[SimpleLocation, CompoundLocation] + type: str + id: str + qualifiers: Dict[str, List[str]] + + def __init__( + self, + location: Optional[Union[SimpleLocation, CompoundLocation]] = ..., + *, + type: str = ..., + id: str = ..., + qualifiers: Optional[Dict[str, List[str]]] = ..., + sub_features: Optional[List["SeqFeature"]] = ..., + ) -> None: ... + + # convenience helpers ---------------------------------------------------- # + def extract( + self, + parent_sequence: SeqLike, + references: Optional[Dict[str, Seq]] = ..., + ) -> SeqLike: ... + + def translate( + self, + parent_sequence: SeqLike, + table: Union[str, int] = ..., + *, + start_offset: Optional[int] = ..., + stop_symbol: str = ..., + to_stop: bool = ..., + cds: Optional[bool] = ..., + gap: Optional[str] = ..., + ) -> Seq: ... + + # delegate dunders to underlying location -------------------------------- # + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[int]: ... + def __contains__(self, value: int) -> bool: ... + + +# --------------------------------------------------------------------------- # +# Miscellaneous public helpers # +# --------------------------------------------------------------------------- # + +class Reference: + location: List[Location] + authors: str + consrtm: str + title: str + journal: str + medline_id: str + pubmed_id: str + comment: str + + +class LocationParserError(ValueError): ... \ No newline at end of file diff --git a/stubs/Bio/__init__.pyi b/stubs/Bio/__init__.pyi new file mode 100644 index 000000000..41cf5475a --- /dev/null +++ b/stubs/Bio/__init__.pyi @@ -0,0 +1 @@ +class BiopythonDeprecationWarning(Warning): ... \ No newline at end of file From 862478d8c3a455420be654c6cba3c924623601a7 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:28:12 +0200 Subject: [PATCH 04/11] Fix bug --- augur/utils.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/augur/utils.py b/augur/utils.py index f4d2d4d5f..d97065b65 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -4,7 +4,7 @@ import sys import warnings from collections import Counter, OrderedDict -from io import RawIOBase +from io import TextIOBase from shlex import quote as shquote from typing import Any, Dict, List, Optional, Tuple, Union @@ -16,7 +16,6 @@ import pandas as pd from Bio.Phylo.BaseTree import Clade, Tree from Bio.SeqFeature import CompoundLocation, FeatureLocation, SeqFeature -from typing_extensions import Buffer from augur.data import as_file from augur.errors import AugurError @@ -151,16 +150,16 @@ def write_json(data: dict, file, indent: int = (None if os.environ.get("AUGUR_MI json.dump(data, handle, indent=indent, sort_keys=sort_keys, cls=AugurJSONEncoder) -class BytesWrittenCounterIO(RawIOBase): - """Binary stream to count the number of bytes sent via write().""" +class BytesWrittenCounterIO(TextIOBase): + """Text stream to count the number of bytes that would be written.""" def __init__(self) -> None: self.written: int = 0 """Number of bytes written.""" - def write(self, b: Buffer) -> int: - n = memoryview(b).nbytes + def write(self, s: str) -> int: + n = len(s.encode('utf-8')) self.written += n - return n + return len(s) def json_size(data: dict) -> int: From 3d2dd9c97a82c3e9ec31c20f3be525864242f9ed Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:37:26 +0200 Subject: [PATCH 05/11] Can remove some ignores now due to stub --- augur/io/sequences.py | 10 +++++----- augur/merge.py | 4 ++-- augur/utils.py | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/augur/io/sequences.py b/augur/io/sequences.py index 5d7cc53ca..c8f574dc8 100644 --- a/augur/io/sequences.py +++ b/augur/io/sequences.py @@ -451,12 +451,12 @@ def _read_genbank(reference: str, feature_names: Optional[Union[set[str], list[s features_skipped = 0 for feature in gb.features: feat = feature - if feat.type=='CDS': # type: ignore[has-type] + if feat.type=='CDS': fname = None - if "locus_tag" in feat.qualifiers: # type: ignore[has-type] - fname = feat.qualifiers["locus_tag"][0] # type: ignore[has-type] - elif "gene" in feat.qualifiers: # type: ignore[has-type] - fname = feat.qualifiers["gene"][0] # type: ignore[has-type] + if "locus_tag" in feat.qualifiers: + fname = feat.qualifiers["locus_tag"][0] + elif "gene" in feat.qualifiers: + fname = feat.qualifiers["gene"][0] else: features_skipped+=1 diff --git a/augur/merge.py b/augur/merge.py index 814633a64..bdcf017cb 100644 --- a/augur/merge.py +++ b/augur/merge.py @@ -499,7 +499,7 @@ def pairs(xs: Iterable[str]) -> Iterable[Tuple[str, str]]: >>> pairs(["abc=123=xyz", "=v=v"]) [('abc', '123=xyz'), ('', 'v=v')] """ - return [tuple(x.split("=", 1)) if "=" in x else ("", x) for x in xs] # type: ignore + return [tuple(x.split("=", 1)) if "=" in x else ("", x) for x in xs] # type: ignore[misc] def count_unique(xs: Iterable[T]) -> Iterable[Tuple[T, int]]: @@ -507,7 +507,7 @@ def count_unique(xs: Iterable[T]) -> Iterable[Tuple[T, int]]: # itertools.groupby(), which requires a sort. Preserving order is a nice # property for the user since we generate an error message with this. # -trs, 24 July 2024 - yield from reduce(lambda counts, x: {**counts, x: counts.get(x, 0) + 1}, xs, counts := {}).items() # type: ignore + yield from reduce(lambda counts, x: {**counts, x: counts.get(x, 0) + 1}, xs, {}).items() # type: ignore[arg-type,dict-item,return-value, call-overload] def shquote_humanized(x): diff --git a/augur/utils.py b/augur/utils.py index d97065b65..fb97031e1 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -413,7 +413,7 @@ def json_to_tree( ] if root: - node = annotate_parents_for_tree(node) # type: ignore + node = annotate_parents_for_tree(node) # type: ignore[arg-type,assignment] return node From f374b0941d53fee27e7c3702f7198da8b3dc5d45 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:39:45 +0200 Subject: [PATCH 06/11] Apply suggestions from code review --- augur/io/sequences.py | 3 +-- augur/utils.py | 1 - stubs/Bio/SeqFeature.pyi | 2 +- stubs/Bio/__init__.pyi | 2 +- 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/augur/io/sequences.py b/augur/io/sequences.py index c8f574dc8..2f7e2cb4a 100644 --- a/augur/io/sequences.py +++ b/augur/io/sequences.py @@ -7,7 +7,7 @@ from typing import Iterable, Iterator, Optional, Union import Bio -from Bio.SeqFeature import SeqFeature +from Bio.SeqFeature import SeqFeature, FeatureLocation from Bio.SeqRecord import SeqRecord from packaging.version import Version @@ -258,7 +258,6 @@ def _read_nuc_annotation_from_gff(record, reference): if len(sequence_regions)>1: raise AugurError(f"Reference {reference!r} contains multiple ##sequence-region pragma lines. Augur can only handle GFF files with a single one.") elif sequence_regions: - from Bio.SeqFeature import FeatureLocation, SeqFeature (name, start, stop) = sequence_regions[0] nuc['pragma'] = SeqFeature( FeatureLocation(start, stop, strand=1), diff --git a/augur/utils.py b/augur/utils.py index fb97031e1..57f24b297 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -26,7 +26,6 @@ from .__version__ import __version__ -TreeLike = Union[Tree, Clade] def augur() -> str: """ diff --git a/stubs/Bio/SeqFeature.pyi b/stubs/Bio/SeqFeature.pyi index 43d04eca8..a0ffd16d7 100644 --- a/stubs/Bio/SeqFeature.pyi +++ b/stubs/Bio/SeqFeature.pyi @@ -140,4 +140,4 @@ class Reference: comment: str -class LocationParserError(ValueError): ... \ No newline at end of file +class LocationParserError(ValueError): ... diff --git a/stubs/Bio/__init__.pyi b/stubs/Bio/__init__.pyi index 41cf5475a..fc428403c 100644 --- a/stubs/Bio/__init__.pyi +++ b/stubs/Bio/__init__.pyi @@ -1 +1 @@ -class BiopythonDeprecationWarning(Warning): ... \ No newline at end of file +class BiopythonDeprecationWarning(Warning): ... From 422369e327adc0945fadb4b3ba84b0a824a2c910 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:42:11 +0200 Subject: [PATCH 07/11] changes.md --- CHANGES.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 84b8d9d27..cc7df5928 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -11,9 +11,14 @@ * filter: Improved speed of using `--group-by month` on large datasets. [#1845][] (@victorlin) +### Internal changes + +* Improved type annotations for better type checking. [#1860][] (@corneliusroemer) + [#1819]: https://github.com/nextstrain/augur/pull/1819 [#1844]: https://github.com/nextstrain/augur/pull/1844 [#1845]: https://github.com/nextstrain/augur/pull/1845 +[#1860]: https://github.com/nextstrain/augur/pull/1860 ## 31.3.0 (3 July 2025) From f5240ce3278909050f3a14637a47c814a107f449 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:50:05 +0200 Subject: [PATCH 08/11] Get rid of unused ignores --- augur/filter/include_exclude_rules.py | 4 ++-- augur/io/shell_command_runner.py | 2 +- augur/parse.py | 6 ++++-- mypy.ini | 2 ++ 4 files changed, 9 insertions(+), 5 deletions(-) diff --git a/augur/filter/include_exclude_rules.py b/augur/filter/include_exclude_rules.py index 3c8cfc548..f8dd36e11 100644 --- a/augur/filter/include_exclude_rules.py +++ b/augur/filter/include_exclude_rules.py @@ -16,9 +16,9 @@ try: # pandas ≥1.5.0 only - PandasUndefinedVariableError = pd.errors.UndefinedVariableError # type: ignore[attr-defined] + PandasUndefinedVariableError = pd.errors.UndefinedVariableError # type: ignore[unused-ignore] except AttributeError: - PandasUndefinedVariableError = pd.core.computation.ops.UndefinedVariableError # type: ignore[attr-defined, misc] + PandasUndefinedVariableError = pd.core.computation.ops.UndefinedVariableError # type: ignore[misc] # The strains to keep as a result of applying a filter function. diff --git a/augur/io/shell_command_runner.py b/augur/io/shell_command_runner.py index b6df32f3b..a1538c4de 100644 --- a/augur/io/shell_command_runner.py +++ b/augur/io/shell_command_runner.py @@ -8,7 +8,7 @@ from signal import SIGKILL except ImportError: # A non-POSIX platform - SIGKILL = None # type: ignore[assignment] + SIGKILL = None def run_shell_command(cmd, raise_errors=False, extra_env=None): diff --git a/augur/parse.py b/augur/parse.py index 8cc76af2c..daf5228bd 100644 --- a/augur/parse.py +++ b/augur/parse.py @@ -36,7 +36,9 @@ def fix_dates(d: str, dayfirst: bool = True) -> str: try: try: # pandas <= 2.1 - from pandas.core.tools.datetimes import parsing # type: ignore[attr-defined, import-not-found] + from pandas.core.tools.datetimes import ( # type: ignore[attr-defined] + parsing, + ) except ImportError: # pandas >= 2.2 from pandas._libs.tslibs import parsing @@ -45,7 +47,7 @@ def fix_dates(d: str, dayfirst: bool = True) -> str: results = parsing.parse_datetime_string_with_reso(d, dayfirst=dayfirst) except AttributeError: # pandas 1.x - results = parsing.parse_time_string(d, dayfirst=dayfirst) # type: ignore[attr-defined] + results = parsing.parse_time_string(d, dayfirst=dayfirst) # type: ignore[attr-defined, unused-ignore] if len(results) == 2: dto, res = results else: diff --git a/mypy.ini b/mypy.ini index b110ae6e3..50d81719c 100644 --- a/mypy.ini +++ b/mypy.ini @@ -13,6 +13,8 @@ strict_optional = False # Directory for stubs mypy_path = stubs +warn_unused_ignores = True + # In the future maybe we can contribute typing stubs for these modules (either # complete stubs in the python/typeshed repo or partial stubs just in # this repo), but for now that's more work than we want to invest. These From 14c8351c06baac10a4bf1728c26bac1db9ee6ffd Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:58:40 +0200 Subject: [PATCH 09/11] unused import --- augur/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/augur/utils.py b/augur/utils.py index 57f24b297..a4a24bed9 100644 --- a/augur/utils.py +++ b/augur/utils.py @@ -14,7 +14,6 @@ import Bio.Phylo.Newick import numpy as np import pandas as pd -from Bio.Phylo.BaseTree import Clade, Tree from Bio.SeqFeature import CompoundLocation, FeatureLocation, SeqFeature from augur.data import as_file From f5854b88528ddde709d055ec0292e2ba5253acb8 Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 14:59:07 +0200 Subject: [PATCH 10/11] Actually log what failed when running flake8 --- tests/test_flake8.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_flake8.py b/tests/test_flake8.py index 683491758..f051cbbfb 100644 --- a/tests/test_flake8.py +++ b/tests/test_flake8.py @@ -5,5 +5,6 @@ def test_flake8(): # Check the exit status ourselves for nicer test output on failure - result = run(["flake8"], cwd = topdir) - assert result.returncode == 0, "flake8 exited with errors" + result = run(["flake8"], cwd=topdir, capture_output=True, text=True) + output = result.stderr + result.stdout + assert result.returncode == 0, f"flake8 exited with errors:\n{output}" From 68e7580762f8687dc0fb47530b2140fbffb5346d Mon Sep 17 00:00:00 2001 From: Cornelius Roemer Date: Sun, 20 Jul 2025 15:09:13 +0200 Subject: [PATCH 11/11] Remove confusing/misleading comment about python version --- mypy.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypy.ini b/mypy.ini index 50d81719c..6c9ec4ac5 100644 --- a/mypy.ini +++ b/mypy.ini @@ -2,8 +2,7 @@ # Check against lowest compatible Python version python_version = 3.9 -# Don't set python_version. Instead, use the default behavior of checking for -# compatibility against the version of Python used to run mypy. +# Which files to check files = augur/ # Require functions with an annotated return type to be explicit about @@ -13,6 +12,7 @@ strict_optional = False # Directory for stubs mypy_path = stubs +# Warn about unused ignores to help keep the code clean warn_unused_ignores = True # In the future maybe we can contribute typing stubs for these modules (either