diff --git a/.github/workflows/mypy.yaml b/.github/workflows/mypy.yaml index 9396103..ae69565 100644 --- a/.github/workflows/mypy.yaml +++ b/.github/workflows/mypy.yaml @@ -3,9 +3,11 @@ on: push: paths: - 'reditools/**py' + - 'test/**py' pull_request: paths: - 'reditools/**py' + - 'test/**py' jobs: mypy: runs-on: ubuntu-latest @@ -27,4 +29,4 @@ jobs: - name: Add mypy annotator uses: pr-annotators/mypy-pr-annotator@v1.0.0 - name: Run mypy - run: uv run mypy reditools/ + run: uv run mypy reditools/ test/ diff --git a/pyproject.toml b/pyproject.toml index fc553ca..859cce0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,3 +33,26 @@ issues = "https://github.com/BioinfoUNIBA/REDItools3/issues" [project.optional-dependencies] wps = ["flake8", "wemake-python-styleguide"] + +[tool.isort] +multi_line_output=3 +include_trailing_comma=true + +[tool.ruff] +line-length = 80 + +[tool.ruff.lint] +select = ["ALL"] +external = ["WPS"] +ignore = [ + "RUF100", "PYI034", "PLR0913", "FBT001", "FBT002", "D203", "D213", + "INP001", "S311", "I", +] + +[tool.ruff.lint.per-file-ignores] +"test/__main__.py" = ["F401"] +"**/__init__.py" = ["F401"] +"test/**py" = ["PT", "SLF001", "FLY002"] + +[tool.ruff.lint.pydocstyle] +convention = "numpy" diff --git a/reditools/__init__.py b/reditools/__init__.py index e69de29..6d26192 100644 --- a/reditools/__init__.py +++ b/reditools/__init__.py @@ -0,0 +1 @@ +"""REDItools: RNA Editing Package.""" diff --git a/reditools/__main__.py b/reditools/__main__.py index 92ea565..82563e6 100644 --- a/reditools/__main__.py +++ b/reditools/__main__.py @@ -1,13 +1,15 @@ +"""CLI entry point for REDItools.""" import sys -from reditools.tools import analyze, annotate, find_repeats, index +from reditools.tools.analyze import main as analyze +from reditools.tools.annotate import main as annotate +from reditools.tools.find_repeats import main as find_repeats +from reditools.tools.index import main as index def usage() -> None: - """ - Print the usage information for the REDItools3 toolkit. - """ + """Print the usage information for the REDItools3 toolkit.""" usage_str = """usage: reditools {analyze,find-repeats,index,annotate} REDItools3 @@ -22,15 +24,15 @@ def usage() -> None: annotate Annotate REDItools RNA output with DNA output """ - print(usage_str) # noqa: WPS421 + print(usage_str) # noqa: WPS421 T201 -if __name__ == '__main__': +if __name__ == "__main__": toolkit = { - 'analyze': analyze, - 'find-repeats': find_repeats, - 'index': index, - 'annotate': annotate, + "analyze": analyze, + "find-repeats": find_repeats, + "index": index, + "annotate": annotate, } if len(sys.argv) > 1: command = sys.argv.pop(1) diff --git a/reditools/alignment_file.py b/reditools/alignment_file.py index 2e2d065..7f30441 100644 --- a/reditools/alignment_file.py +++ b/reditools/alignment_file.py @@ -1,16 +1,21 @@ +"""A wrapper around pysam.AlignmentFile with integrated quality control.""" +from __future__ import annotations -from types import TracebackType -from typing import Any, Collection, Iterator +from typing import TYPE_CHECKING -from pysam import AlignedSegment from pysam.libcalignmentfile import AlignmentFile as PysamAlignmentFile -from reditools.region import Region +if TYPE_CHECKING: + from types import TracebackType + from typing import Any, Collection, Iterator + + from pysam import AlignedSegment + + from reditools.region import Region class ReadQC: - """ - Perform quality control checks on aligned reads. + """Perform quality control checks on aligned reads. Parameters ---------- @@ -21,16 +26,16 @@ class ReadQC: excluded_read_names : Collection[str] | None A collection of read names to be excluded. """ - _flags_to_keep = {0, 16, 83, 99, 147, 163} + + _flags_to_keep = frozenset((0, 16, 83, 99, 147, 163)) def __init__( self, min_quality: int, min_length: int, excluded_read_names: Collection[str] | None, - ): - """ - Initialize the ReadQC with quality and length thresholds. + ) -> None: + """Initialize the ReadQC with quality and length thresholds. Parameters ---------- @@ -43,20 +48,20 @@ def __init__( """ self.min_quality = min_quality self.min_length = min_length - self.excluded_read_names = excluded_read_names self.check_list = [self.check_baseline] if self.min_quality > 0: self.check_list.append(self.check_quality) if self.min_length > 0: self.check_list.append(self.check_length) - if self.excluded_read_names: - self.excluded_read_names = set(self.excluded_read_names) + if excluded_read_names: + self.excluded_read_names = set(excluded_read_names) self.check_list.append(self.check_excluded_read_names) + else: + self.excluded_read_names = set() def check_baseline(self, read: AlignedSegment) -> bool: - """ - Check if the read passes baseline flag and tag requirements. + """Check if the read passes baseline flag and tag requirements. Parameters ---------- @@ -68,11 +73,10 @@ def check_baseline(self, read: AlignedSegment) -> bool: bool True if the read passes, False otherwise. """ - return read.flag in self._flags_to_keep and not read.has_tag('SA') - + return read.flag in self._flags_to_keep and not read.has_tag("SA") + def check_quality(self, read: AlignedSegment) -> bool: - """ - Check if the read passes the minimum mapping quality threshold. + """Check if the read passes the minimum mapping quality threshold. Parameters ---------- @@ -87,8 +91,7 @@ def check_quality(self, read: AlignedSegment) -> bool: return read.mapping_quality >= self.min_quality def check_length(self, read: AlignedSegment) -> bool: - """ - Check if the read passes the minimum length threshold. + """Check if the read passes the minimum length threshold. Parameters ---------- @@ -103,8 +106,7 @@ def check_length(self, read: AlignedSegment) -> bool: return read.query_length >= self.min_length def check_excluded_read_names(self, read: AlignedSegment) -> bool: - """ - Check if the read name is not in the excluded list. + """Check if the read name is not in the excluded list. Parameters ---------- @@ -116,11 +118,10 @@ def check_excluded_read_names(self, read: AlignedSegment) -> bool: bool True if the read name is not excluded, False otherwise. """ - return read.query_name not in self.excluded_read_names # type: ignore + return read.query_name not in self.excluded_read_names def run_check(self, read: AlignedSegment) -> bool: - """ - Run all configured quality control checks on the read. + """Run all configured quality control checks on the read. Parameters ---------- @@ -136,8 +137,7 @@ def run_check(self, read: AlignedSegment) -> bool: class RTAlignmentFile: - """ - A wrapper around pysam.AlignmentFile with integrated quality control. + """A wrapper around pysam.AlignmentFile with integrated quality control. Parameters ---------- @@ -159,10 +159,9 @@ def __init__( min_quality: int=0, min_length: int=0, excluded_read_names: Collection[str] | None=None, - **kwargs: Any, + **kwargs: Any, # noqa: ANN401 ) -> None: - """ - Initialize the RTAlignmentFile. + """Initialize the RTAlignmentFile. Parameters ---------- @@ -177,14 +176,17 @@ def __init__( **kwargs Keyword arguments passed to pysam.AlignmentFile. """ - kwargs['ignore_truncation'] = True + kwargs["ignore_truncation"] = True self.alignment_file = PysamAlignmentFile(filename, **kwargs) self.alignment_file.check_index() - self.readqc = ReadQC(min_quality, min_length, excluded_read_names) + if excluded_read_names is None: + excluded_names: Collection[str] = [] + else: + excluded_names = excluded_read_names + self.readqc = ReadQC(min_quality, min_length, excluded_names) - def __enter__(self): # type: ignore - """ - Enter the runtime context related to this object. + def __enter__(self) -> RTAlignmentFile: + """Enter the runtime context related to this object. Returns ------- @@ -195,20 +197,19 @@ def __enter__(self): # type: ignore def __exit__( self, - exc_type: type, - exc_value: Exception, - traceback: TracebackType, + typ: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, ) -> None: - """ - Exit the runtime context related to this object. + """Exit the runtime context related to this object. Parameters ---------- - exc_type : type | None + typ : type[BaseException] | None The exception type. - exc_value : Exception | None + exc : BaseException | None The exception value. - traceback : TracebackType | None + tb : TracebackType | None The traceback. """ self.alignment_file.close() @@ -217,9 +218,7 @@ def fetch_by_position( self, region: Region | str, ) -> Iterator[list[AlignedSegment]]: - """ - Fetch reads from the alignment file grouped by their reference start - position. + """Fetch reads from the alignment file grouped by reference start. Parameters ---------- diff --git a/reditools/alignment_manager.py b/reditools/alignment_manager.py index a5f6baa..a6a8b94 100644 --- a/reditools/alignment_manager.py +++ b/reditools/alignment_manager.py @@ -1,26 +1,32 @@ -from itertools import chain -from typing import Collection, Iterable, Iterator +"""Fetch reads from multiple alignment files.""" +from __future__ import annotations -from pysam import AlignedSegment +from itertools import chain +from math import inf +from typing import TYPE_CHECKING from reditools.alignment_file import RTAlignmentFile -from reditools.region import Region +if TYPE_CHECKING: + from typing import Collection, Iterator + + from pysam import AlignedSegment + + from reditools.region import Region class ReadGroupIter: - """ - Iterator over groups of reads sharing the same reference start position. + """Iterator over groups of reads sharing the same reference start position. Parameters ---------- iterator : Iterator An iterator yielding lists of AlignedSegment objects. """ - __slots__ = ('iterator', 'reads', 'reference_start') - def __init__(self, iterator: Iterator): - """ - Initialize the ReadGroupIter. + __slots__ = ("iterator", "reads", "reference_start") + + def __init__(self, iterator: Iterator[list[AlignedSegment]]) -> None: + """Initialize the ReadGroupIter. Parameters ---------- @@ -31,8 +37,7 @@ def __init__(self, iterator: Iterator): next(self) def __bool__(self) -> bool: - """ - Check if there are more reads. + """Check if there are more reads. Returns ------- @@ -42,8 +47,7 @@ def __bool__(self) -> bool: return bool(self.reads) def __next__(self) -> list[AlignedSegment] | None: - """ - Get the next group of reads. + """Get the next group of reads. Returns ------- @@ -52,15 +56,13 @@ def __next__(self) -> list[AlignedSegment] | None: """ self.reads = next(self.iterator, None) if self.reads: - self.reference_start = self.reads[0].reference_start + self.reference_start: int | float = self.reads[0].reference_start else: - self.reference_start = None + self.reference_start = inf return self.reads class FetchGroupIter: - """ - Iterator that merges multiple ReadGroupIter objects, yielding reads grouped - by position. + """Iterator that merges multiple ReadGroupIter objects. Parameters ---------- @@ -68,9 +70,11 @@ class FetchGroupIter: A list of iterators, each yielding reads from an alignment file. """ - def __init__(self, fetch_iters: list[Iterator]): - """ - Initialize the FetchGroupIter. + def __init__( # noqa: WPS23 + self, + fetch_iters: list[Iterator[list[AlignedSegment]]], + ) -> None: + """Initialize the FetchGroupIter. Parameters ---------- @@ -84,8 +88,7 @@ def __init__(self, fetch_iters: list[Iterator]): self.read_groups.append(rgi) def __iter__(self) -> Iterator[list[AlignedSegment]]: - """ - Return the iterator object itself. + """Return the iterator object itself. Returns ------- @@ -96,8 +99,7 @@ def __iter__(self) -> Iterator[list[AlignedSegment]]: yield next(self) def __bool__(self) -> bool: - """ - Check if there are more read groups. + """Check if there are more read groups. Returns ------- @@ -107,9 +109,7 @@ def __bool__(self) -> bool: return bool(self.read_groups) def __next__(self) -> list[AlignedSegment]: - """ - Get the next group of reads from all alignment files for the same - position. + """Get the next group of reads from all alignment files. Returns ------- @@ -125,12 +125,10 @@ def __next__(self) -> list[AlignedSegment]: reads.append(rgi.reads) if next(rgi) is None: self.read_groups.pop(idx) - return list(chain(*reads)) # type: ignore + return list(chain(*reads)) # type: ignore[arg-type] class AlignmentManager: - """ - Manage multiple alignment files and provide unified access to reads by - position. + """Manage multiple alignment files. Parameters ---------- @@ -141,14 +139,14 @@ class AlignmentManager: min_length : int, optional Minimum read length (default is 0). """ + def __init__( self, excluded_read_names: Collection[str] | None=None, min_quality: int=0, min_length: int=0, - ): # noqa: WPS475 - """ - Initialize the AlignmentManager. + ) -> None: # noqa: WPS475 + """Initialize the AlignmentManager. Parameters ---------- @@ -162,13 +160,15 @@ def __init__( self._bams: list[RTAlignmentFile] = [] self.file_list: list[str] = [] self.next_read_start: int | None = None - self.excluded_read_names = excluded_read_names + if excluded_read_names is None: + self.excluded_read_names: Collection[str] = [] + else: + self.excluded_read_names = excluded_read_names self.min_quality = min_quality self.min_length = min_length def add_file(self, fname: str) -> None: - """ - Add an alignment file to the manager. + """Add an alignment file to the manager. Parameters ---------- @@ -187,9 +187,8 @@ def add_file(self, fname: str) -> None: def fetch_by_position( self, region: Region | str, - ) -> Iterable[list[AlignedSegment]]: - """ - Fetch reads from all managed files, grouped by position. + ) -> Iterator[list[AlignedSegment]]: + """Fetch reads from all managed files, grouped by position. Parameters ---------- diff --git a/reditools/compiled_position.py b/reditools/compiled_position.py index c23fbeb..1b476d9 100644 --- a/reditools/compiled_position.py +++ b/reditools/compiled_position.py @@ -1,6 +1,12 @@ +"""Class to store compiled information for a specific genomic position.""" +from __future__ import annotations + from dataclasses import dataclass, field from typing import Iterator +from reditools.constants import bases as base_order +from reditools.constants import comp_map + @dataclass class CompiledPosition: @@ -29,8 +35,6 @@ class CompiledPosition: strands: list[str] = field(default_factory=list) bases: list[str] = field(default_factory=list) - _comp = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C'} - def __len__(self) -> int: """Return the number of bases at this position. @@ -74,17 +78,17 @@ def calculate_strand(self, threshold: float = 0) -> str: pos_count = 0 neg_count = 0 for strand in self.strands: - if strand == '+': + if strand == "+": pos_count += 1 - elif strand == '-': + elif strand == "-": neg_count += 1 if pos_count == neg_count: - return '*' + return "*" if pos_count / (pos_count + neg_count) >= threshold: - return '+' + return "+" if neg_count / (pos_count + neg_count) >= threshold: - return '-' - return '*' + return "-" + return "*" def filter_by_strand(self, strand: str) -> None: """Filter observations to keep only those from a specific strand. @@ -94,7 +98,7 @@ def filter_by_strand(self, strand: str) -> None: strand : str The strand to keep ('+', '-', or '*'). If '*', no filtering is done. """ - if strand == '*': + if strand == "*": return keep = [ idx for idx in range(len(self.bases)) @@ -106,8 +110,8 @@ def filter_by_strand(self, strand: str) -> None: def complement(self) -> None: """Replace all bases and the reference with their complements.""" - self.bases = [self._comp[base] for base in self.bases] - self.ref = self._comp[self.ref] + self.bases = [comp_map[base] for base in self.bases] + self.ref = comp_map[self.ref] class RTResult: @@ -131,9 +135,11 @@ class RTResult: A list of observed variants (e.g., ['AG']). """ - _base_order = 'ACGT' - - def __init__(self, compiled_position: CompiledPosition, strand: str): + def __init__( + self, + compiled_position: CompiledPosition, + strand: str, + ) -> None: """Initialize RTResult. Parameters @@ -150,12 +156,12 @@ def __init__(self, compiled_position: CompiledPosition, strand: str): self.position = self.cp.position self.contig = self.cp.contig - self.counter = {_: 0 for _ in self._base_order} + self.counter = dict.fromkeys(base_order, 0) for base in self.cp.bases: self.counter[base] += 1 self.variants = [ - f'{self.reference}{_}' for _ in self._base_order + f"{self.reference}{_}" for _ in base_order if self[_] and _ != self.reference ] @@ -172,7 +178,7 @@ def __getitem__(self, base: str) -> int: int The count of the requested base. """ - if base.upper() == 'REF': + if base.upper() == "REF": return self.counter[self.reference] return self.counter[base] @@ -182,9 +188,9 @@ def __iter__(self) -> Iterator[int]: Yields ------ int - The count of each base in order: 'A, 'C', 'G', 'T'. + The count of each base in order: "A", "C", "G", "T". """ - return (self[base] for base in self._base_order) + return (self[base] for base in base_order) def __len__(self) -> int: """Return the total number of reads at this position. @@ -209,11 +215,11 @@ def edit_ratio(self) -> float: The editing ratio. """ max_edits = 0 - for base, count in zip(self._base_order, self): + for base, count in zip(base_order, self): if base != self.reference and count > max_edits: max_edits = count try: - return max_edits / (self['REF'] + max_edits) + return max_edits / (self["REF"] + max_edits) except ZeroDivisionError: return 0 diff --git a/reditools/compiled_reads.py b/reditools/compiled_reads.py index 8798663..b0f50d6 100644 --- a/reditools/compiled_reads.py +++ b/reditools/compiled_reads.py @@ -1,10 +1,16 @@ -from typing import Iterator, Optional +"""Aggregate reads from alignment file(s) that have the same start position.""" +from __future__ import annotations -from pysam import AlignedSegment +from typing import TYPE_CHECKING from reditools.compiled_position import CompiledPosition from reditools.fasta_file import RTFastaFile +if TYPE_CHECKING: + from typing import Iterator + + from pysam import AlignedSegment + class RefFetch: """Helper class to fetch reference sequences. @@ -14,7 +20,7 @@ class RefFetch: the AlignedSegment if MD tags are available. """ - def __init__(self, fasta_file_path: Optional[str] = None): + def __init__(self, fasta_file_path: str | None = None) -> None: """Initialize RefFetch. Parameters @@ -29,9 +35,10 @@ def __init__(self, fasta_file_path: Optional[str] = None): self._refseq_fn = self.get_ref_from_read def get_refseq(self, read: AlignedSegment) -> Iterator[str]: - """Fetch reference sequence. If a FASTA file was provided in the - constructor, this function calla get_ref_from_fasta. Otherwise it - calls get_ref_from_read. + """Fetch reference sequence. + + If a FASTA file was provided in the constructor, this function calls + get_ref_from_fasta. Otherwise it calls get_ref_from_read. Parameters ---------- @@ -80,7 +87,7 @@ def get_ref_from_fasta(self, read: AlignedSegment) -> Iterator[str]: pairs = read.get_aligned_pairs(matches_only=True) indices = [ref for _, ref in pairs] return self.fasta_file.get_base( - read.reference_name, # type: ignore + read.reference_name, # type: ignore[arg-type] *indices, ) @@ -92,7 +99,7 @@ class CompiledReads: and strand. """ - _strands = ('-', '+', '*') + _strands = ("-", "+", "*") def __init__( self, @@ -100,8 +107,8 @@ def __init__( min_base_position: int = 0, max_base_position: int = 0, min_base_quality: int = 0, - fasta_file: Optional[str] = None, - ): + fasta_file: str | None = None, + ) -> None: """Initialize CompiledReads. Parameters @@ -131,9 +138,9 @@ def __init__( self.reference = RefFetch(fasta_file) self._qc = { - 'min_base_quality': min_base_quality, - 'min_base_position': min_base_position, - 'max_base_position': max_base_position, + "min_base_quality": min_base_quality, + "min_base_position": min_base_position, + "max_base_position": max_base_position, } def add_reads(self, reads: list[AlignedSegment]) -> None: @@ -151,7 +158,7 @@ def add_reads(self, reads: list[AlignedSegment]) -> None: self._nucleotides[pos] = CompiledPosition( ref=ref, position=pos, - contig=read.reference_name, # type: ignore + contig=read.reference_name, # type: ignore[arg-type] ) self._nucleotides[pos].add_base(quality, strand, base) @@ -186,20 +193,20 @@ def _prep_read( # noqa: WPS231 self.reference.get_refseq(read), ): # Right end trim - if read_pos > read.query_length - self._qc['max_base_position']: + if read_pos > read.query_length - self._qc["max_base_position"]: break # Left end trim - if read_pos < self._qc['min_base_position']: + if read_pos < self._qc["min_base_position"]: continue - read_base = read.query_sequence[read_pos] # type: ignore - if ref_base == 'N' or read_base == 'N': + read_base = read.query_sequence[read_pos] # type: ignore[index] + if ref_base == "N" or read_base == "N": continue - phred = read.query_qualities[read_pos] # type: ignore - if phred < self._qc['min_base_quality']: + phred = read.query_qualities[read_pos] # type: ignore[index] + if phred < self._qc["min_base_quality"]: continue yield (ref_pos, read_base, phred, ref_base) - def _unstranded_strand(self, read: AlignedSegment) -> int: + def _unstranded_strand(self, read: AlignedSegment) -> int: # noqa: ARG002 return 2 def _stranded_strand(self, read: AlignedSegment) -> int: diff --git a/reditools/constants.py b/reditools/constants.py new file mode 100644 index 0000000..97df71d --- /dev/null +++ b/reditools/constants.py @@ -0,0 +1,10 @@ +"""Dictionary for base complements.""" +comp_map = { + "A": "T", + "T": "A", + "C": "G", + "G": "C", + "N": "N", + "-": "-", +} +bases = ("A", "C", "G", "T") diff --git a/reditools/fasta_file.py b/reditools/fasta_file.py index e4a6a71..8de35b3 100644 --- a/reditools/fasta_file.py +++ b/reditools/fasta_file.py @@ -1,53 +1,86 @@ -from types import TracebackType -from typing import Iterator +"""A wrapper around pysam.FastaFile for genomic sequence access.""" +from __future__ import annotations + +from typing import TYPE_CHECKING from pysam.libcfaidx import FastaFile as PysamFastaFile +if TYPE_CHECKING: + from types import TracebackType + from typing import Iterator + + +class MissingContigError(LookupError): + """Contig name is missing from the FASTA file.""" + + def __init__(self, contig_name: str) -> None: + """Initialize self. + + Parameters + ---------- + contig_name : str + Missing contig name. + """ + self.message = f"Reference name {contig_name} not found in FASTA file." + super().__init__(self.message) + +class PastContigEndError(LookupError): + """Genomic position is outside contig bounds.""" + + def __init__(self, contig_name: str, position: int) -> None: + """Initialize self. + + Parameters + ---------- + contig_name : str + Name of the contig. + position : int + Offending genomic position. + """ + self.message = ( + f"Base position {position} is outside the bounds of " + f"{contig_name}. Are you using the correct reference?" + ) + super().__init__(self.message) class RTFastaFile: - """ - A wrapper around pysam.FastaFile for genomic sequence access. - """ + """A wrapper around pysam.FastaFile for genomic sequence access.""" def __init__(self, filename: str) -> None: - """ - Initialize the RTFastaFile. + """Initialize the RTFastaFile. Parameters ---------- - *args - Arguments passed to pysam.FastaFile. - **kwargs - Keyword arguments passed to pysam.FastaFile. + filename : str + FASTA file path. """ self.pysam_fasta_file = PysamFastaFile(filename) - def __enter__(self): # type: ignore + def __enter__(self) -> RTFastaFile: + """Open RTFastaFile.""" return self def __exit__( self, - exc_type: type, - exc_value: Exception, - traceback: TracebackType, + typ: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, ) -> None: - """ - Exit the runtime context related to this object. + """Exit the runtime context related to this object. Parameters ---------- - exc_type : type | None + typ : type[BaseException] | None The exception type. - exc_value : Exception | None + exc : BaseException | None The exception value. - traceback : TracebackType | None + tb : TracebackType | None The traceback. """ self.pysam_fasta_file.close() def get_base(self, contig: str, *position: int) -> Iterator[str]: - """ - Retrieve bases at specified positions from a contig. + """Retrieve bases at specified positions from a contig. Parameters ---------- @@ -63,21 +96,18 @@ def get_base(self, contig: str, *position: int) -> Iterator[str]: Raises ------ - KeyError + MissingContigError If the contig is not found in the FASTA file. - IndexError + PastContigEndError If a position is outside the bounds of the contig. """ - - if contig not in self.pysam_fasta_file: - if contig.startswith('chr'): - new_contig = contig.replace('chr', '') + if contig not in self.pysam_fasta_file.references: + if contig.startswith("chr"): + new_contig = contig[3:] else: - new_contig = f'chr{contig}' - if new_contig not in self.pysam_fasta_file: - raise KeyError( - f'Reference name {contig} not found in FASTA file.', - ) + new_contig = f"chr{contig}" + if new_contig not in self.pysam_fasta_file.references: + raise MissingContigError(contig) contig = new_contig sorted_pos = sorted(position) seq = self.pysam_fasta_file.fetch( @@ -86,9 +116,7 @@ def get_base(self, contig: str, *position: int) -> Iterator[str]: sorted_pos[-1] + 1, ) try: - return (seq[_ - sorted_pos[0]].upper() for _ in position) + for pos in position: + yield seq[pos - sorted_pos[0]].upper() except IndexError as exc: - raise IndexError( - f'Base position {position} is outside the bounds of ' + - '{contig}. Are you using the correct reference?', - ) from exc + raise PastContigEndError(contig, max(position)) from exc diff --git a/reditools/file_utils.py b/reditools/file_utils.py index 78a52b5..32de3bb 100644 --- a/reditools/file_utils.py +++ b/reditools/file_utils.py @@ -1,20 +1,21 @@ +"""File handling utilities.""" +from __future__ import annotations import csv -import os import tempfile from gzip import open as gzip_open +from pathlib import Path from typing import IO, Iterator from reditools.region import Region -def open_stream( # type: ignore +def open_stream( # type: ignore[no-untyped-def] # noqa: ANN201 path: str, - mode: str='rt', - encoding: str='utf-8', + mode: str="rt", + encoding: str="utf-8", ): - """ - Open a file stream, handling both plain and gzipped files. + """Open a file stream, handling both plain and gzipped files. Parameters ---------- @@ -30,14 +31,12 @@ def open_stream( # type: ignore file-like object The opened file stream. """ - if path.endswith('gz'): + if path.endswith("gz"): return gzip_open(path, mode, encoding=encoding) - return open(path, mode, encoding=encoding) # noqa: WPS515 - + return Path(path).open(mode, encoding=encoding) # noqa: WPS515 SIM115 def read_bed_file(*path: str) -> Iterator[Region]: - """ - Read genomic regions from one or more BED files. + """Read genomic regions from one or more BED files. Parameters ---------- @@ -53,8 +52,8 @@ def read_bed_file(*path: str) -> Iterator[Region]: yield from read_bed_file(*path[1:]) with open_stream(path[0]) as stream: reader = csv.reader( - filter(lambda row: row[0] != '#', stream), - delimiter='\t', + filter(lambda row: row[0] != "#", stream), + delimiter="\t", ) for row in reader: yield Region( @@ -68,10 +67,9 @@ def concat( output: IO, *fnames: str, clean_up: bool=True, - encoding: str='utf-8', + encoding: str="utf-8", ) -> None: - """ - Concatenate multiple files into a single output stream. + """Concatenate multiple files into a single output stream. Parameters ---------- @@ -86,16 +84,14 @@ def concat( The encoding to use when reading files (default is 'utf-8'). """ for fname in fnames: - with open(fname, 'r', encoding=encoding) as stream: - for line in stream: - output.write(line) + with Path(fname).open("r", encoding=encoding) as stream: + output.writelines(stream) if clean_up: - os.remove(fname) + Path(fname).unlink() def load_text_file(file_name: str) -> list[str]: - """ - Load lines from a text file into a list, stripping whitespace. + """Load lines from a text file into a list, stripping whitespace. Parameters ---------- @@ -107,18 +103,17 @@ def load_text_file(file_name: str) -> list[str]: list[str] A list of stripped lines from the file. """ - with open_stream(file_name, 'r') as stream: + with open_stream(file_name, "r") as stream: return [line.strip() for line in stream] -def make_dir(prefix: str | None=None, dir: str | None=None) -> str: - """ - Creates a folder. +def make_dir(prefix: str | None=None, dirname: str | None=None) -> str: + """Create a folder. Parameters ---------- prefix : str Filename prefix. - dir : str + dirname : str Path to folder parent. Returns @@ -126,8 +121,8 @@ def make_dir(prefix: str | None=None, dir: str | None=None) -> str: str Path to the folder. """ - with tempfile.NamedTemporaryFile(prefix=prefix, dir=dir) as stream: + with tempfile.NamedTemporaryFile(prefix=prefix, dir=dirname) as stream: valid_name = stream.name - os.mkdir(valid_name) + Path(valid_name).mkdir() return valid_name diff --git a/reditools/logger.py b/reditools/logger.py index 4157ba4..229e657 100644 --- a/reditools/logger.py +++ b/reditools/logger.py @@ -1,3 +1,4 @@ +"""Handle logging operations with different severity levels.""" import os import socket import sys @@ -6,8 +7,7 @@ class Logger: - """ - Handle logging operations with different severity levels. + """Handle logging operations with different severity levels. Attriutes ---------- @@ -19,13 +19,12 @@ class Logger: Output all messages """ - silent_level = 'SILENT' - info_level = 'INFO' - debug_level = 'DEBUG' + silent_level = "SILENT" + info_level = "INFO" + debug_level = "DEBUG" - def __init__(self, level: str): - """ - Initialize the Logger with a specified logging level. + def __init__(self, level: str) -> None: + """Initialize the Logger with a specified logging level. Parameters ---------- @@ -33,9 +32,13 @@ def __init__(self, level: str): The logging level ('SILENT', 'INFO', or 'DEBUG'). """ hostname = socket.gethostname() - ip_addr = socket.gethostbyname(hostname) pid = os.getpid() - self.hostname_string = f'{hostname}|{ip_addr}|{pid}' + try: + ip_addr = socket.gethostbyname(hostname) + except socket.gaierror: + self.hostname_string = f"{hostname}|{pid}" + else: + self.hostname_string = f"{hostname}|{ip_addr}|{pid}" self._level = level.upper() if self._level == self.debug_level: @@ -45,7 +48,12 @@ def __init__(self, level: str): else: self._log_fn = self._log_silent - def log(self, level: str, message: str, *args: Any) -> None: + def log( + self, + level: str, + message: str, + *args: Any, # noqa: ANN401 + ) -> None: """Conditionally output a message to STDERR. Parameters @@ -61,8 +69,7 @@ def log(self, level: str, message: str, *args: Any) -> None: @property def level(self) -> str: - """ - Get the current logging level. + """Get the current logging level. Returns ------- @@ -71,17 +78,32 @@ def level(self) -> str: """ return self._level - def _log_all(self, level: str, message: str, *args: Any) -> None: - timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S') + def _log_all( + self, + level: str, + message: str, + *args: Any, # noqa: ANN401 + ) -> None: + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") # noqa: DTZ005 message = message.format(*args) sys.stderr.write( - f'{timestamp} [{self.hostname_string}] ' + - f'[{level}] {message}\n', + f"{timestamp} [{self.hostname_string}] " + f"[{level}] {message}\n", ) - def _log_info(self, level: str, message: str, *args: Any) -> None: + def _log_info( + self, + level: str, + message: str, + *args: Any, # noqa: ANN401 + ) -> None: if level == self.info_level: self._log_all(level, message, *args) - def _log_silent(self, level: str, message: str, *args: Any) -> None: + def _log_silent( + self, + level: str, + message: str, + *args: Any, # noqa: ANN401 + ) -> None: pass # noqa: WPS420 diff --git a/reditools/reditools.py b/reditools/reditools.py index edc04b6..a285847 100644 --- a/reditools/reditools.py +++ b/reditools/reditools.py @@ -1,11 +1,18 @@ +"""Main class for running REDItools analysis.""" +from __future__ import annotations -from typing import Iterator +from typing import TYPE_CHECKING -from reditools.alignment_manager import AlignmentManager -from reditools.compiled_position import CompiledPosition, RTResult +from reditools.compiled_position import RTResult from reditools.compiled_reads import CompiledReads from reditools.logger import Logger -from reditools.region import Region + +if TYPE_CHECKING: + from typing import Iterator + + from reditools.alignment_manager import AlignmentManager + from reditools.compiled_position import CompiledPosition + from reditools.region import Region """ Set the strand property to UNSTRANDED_MODE for unstranded analysis. @@ -26,17 +33,13 @@ class REDItools: - """ - Main class for running REDItools analysis. + """Main class for running REDItools analysis. Provides methods to set up analysis parameters and process alignment data. """ - def __init__(self) -> None: - """ - Initialize REDItools with default parameters. - """ + """Initialize REDItools with default parameters.""" self._min_column_length = 1 self._min_edits = 0 self._min_edits_per_nucleotide = 0 @@ -60,8 +63,7 @@ def __init__(self) -> None: @property def log_level(self) -> str: - """ - Get the current logging level. + """Get the current logging level. Returns ------- @@ -72,8 +74,7 @@ def log_level(self) -> str: @log_level.setter def log_level(self, level: str) -> None: - """ - Set the logging level. + """Set the logging level. Parameters ---------- @@ -88,8 +89,7 @@ def analyze( alignment_manager: AlignmentManager, region: Region, ) -> Iterator[RTResult]: - """ - Analyze a genomic region using alignment data. + """Analyze a genomic region using alignment data. Parameters ---------- @@ -114,7 +114,7 @@ def analyze( self.log( Logger.info_level, - 'Fetching reads [FILELIST={}] [REGION={}]', + "Fetching reads [FILELIST={}] [REGION={}]", alignment_manager.file_list, region, ) @@ -122,7 +122,7 @@ def analyze( for reads in alignment_manager.fetch_by_position(region=region): self.log( Logger.debug_level, - 'Adding {} reads starting from {}:{}', + "Adding {} reads starting from {}:{}", len(reads), region.contig, reads[0].reference_start, @@ -142,20 +142,19 @@ def analyze( if bases.position >= region.start: self.log( Logger.debug_level, - 'Yielding output for {} reads', + "Yielding output for {} reads", len(rtresult), ) yield rtresult self.log( Logger.info_level, - '[REGION={}] {} total reads', + "[REGION={}] {} total reads", region, total, ) def use_strand_correction(self) -> None: - """ - Enable strand correction during analysis. + """Enable strand correction during analysis. Strand correction will filter reads to only those of the consensus strand and also report the complement of the edits and reference base @@ -164,8 +163,7 @@ def use_strand_correction(self) -> None: self._use_strand_correction = True def add_reference(self, reference_fname: str) -> None: - """ - Add a reference FASTA file for genomic reference sequences. + """Add a reference FASTA file for genomic reference sequences. Parameters ---------- @@ -177,17 +175,17 @@ def add_reference(self, reference_fname: str) -> None: def _process_bases(self, bases: CompiledPosition) -> RTResult: self.log( Logger.debug_level, - 'Analyzing position {} {}', + "Analyzing position {} {}", bases.position, bases.contig, ) if self.strand == UNSTRANDED_MODE: - strand = '*' + strand = "*" else: strand = bases.calculate_strand( threshold=self.strand_confidence_threshold, ) bases.filter_by_strand(strand) - if self._use_strand_correction and strand == '-': + if self._use_strand_correction and strand == "-": bases.complement() return RTResult(bases, strand) diff --git a/reditools/region.py b/reditools/region.py index 8359da6..063b09a 100644 --- a/reditools/region.py +++ b/reditools/region.py @@ -1,3 +1,5 @@ +"""Represent a genomic region.""" +from __future__ import annotations import re from dataclasses import dataclass @@ -5,10 +7,77 @@ from pysam import AlignmentFile +class RegionSplitError(IndexError): + """Region cannot be split.""" + + def __init__(self) -> None: + """Initialize self.""" + self.message = "Can only split a region with a start and stop." + super().__init__(self.message) + +class RegionBadStartError(ValueError): + """Region start is less than one.""" + + def __init__(self, bad_start: int) -> None: + """Initialize self. + + Parameters + ---------- + bad_start : int + Offending start position. + """ + self.message = ( + f"Start position ({bad_start}) must be greater than or equal to one" + ) + super().__init__(self.message) + +class RegionNeedsAlignmentError(ValueError): + """Region needs an Alignment File for initalization.""" + + def __init__(self) -> None: + """Initialize self.""" + self.message = ( + "An alignment file must be provided if no stop position " + "is present in the region string." + ) + super().__init__(self.message) + +class RegionStartPastStopError(ValueError): + """Region start position is after the stop position.""" + + def __init__(self, start: int, stop: int) -> None: + """Initialize self. + + Parameters + ---------- + start : int + Region start position. + stop : int + Region stop position. + """ + self.message = ( + f"Stop position ({stop}) must be greater than or " + f"equal to start ({start}).", + ) + super().__init__(self.message) + +class RegionFormatError(ValueError): + """Region string is an unknown format.""" + + def __init__(self, region_str: str) -> None: + """Initialize self. + + Parameters + ---------- + region_str : str + Offending region string. + """ + self.message = f"Unrecognized format: {region_str}." + super().__init__(self.message) + @dataclass(slots=True, order=True, frozen=True) class Region: - """ - Represent a genomic region. + """Represent a genomic region. Parameters ---------- @@ -25,8 +94,7 @@ class Region: stop: int def __str__(self) -> str: - """ - Return a string representation of the region. + """Return a string representation of the region. Returns ------- @@ -36,13 +104,12 @@ def __str__(self) -> str: one_idx_start = self.start + 1 if self.stop is None: if self.start > 0: - return f'{self.contig}:{one_idx_start}' + return f"{self.contig}:{one_idx_start}" return self.contig - return f'{self.contig}:{one_idx_start}-{self.stop}' + return f"{self.contig}:{one_idx_start}-{self.stop}" - def split(self, window: int) -> list['Region']: - """ - Split the region into smaller sub-regions of a specified window size. + def split(self, window: int) -> list[Region]: + """Split the region into smaller sub-regions of a specified window size. Parameters ---------- @@ -56,27 +123,27 @@ def split(self, window: int) -> list['Region']: Raises ------ - IndexError + RegionSplitError If either start or stop is None. """ if self.stop is None or self.start is None: - raise IndexError('Can only split a region with a start and stop.') - sub_regions = [] - for new_start in range(self.start, self.stop, window): - sub_regions.append(Region( + raise RegionSplitError + return [ + Region( contig=self.contig, start=new_start, - stop=min(new_start + window, self.stop))) - return sub_regions + stop=min(new_start + window, self.stop), + ) + for new_start in range(self.start, self.stop, window) + ] @classmethod def from_string( cls, region_str: str, alignment_file: str | None=None, - ) -> 'Region': - """ - Create a Region object from a string and an alignment file. + ) -> Region: + """Create a Region object from a string and an alignment file. Parameters ---------- @@ -94,36 +161,30 @@ def from_string( Raises ------ - ValueError - If start is less than 0 or if stop is less than or equal to start. + RegionBadStartError + If region start is less than 0. + RegionNeedsAlignmentError + If region stop is None and alignment_file is None. + RegionStartPastStopError + if the region start is equal to or greater than region stop. """ contig, start, stop = Region.parse_string(region_str) if start is None: start = 0 elif start < 0: - raise ValueError( - f'Start position ({start}) must be greater than or ' - 'equal to one.', - ) + raise RegionBadStartError(start) if stop is None: if alignment_file is None: - raise ValueError( - 'An alignment file must be provided if no stop position ' - 'is present in the region string.' - ) + raise RegionNeedsAlignmentError with AlignmentFile(alignment_file, ignore_truncation=True) as bam: stop = bam.get_reference_length(contig) if stop <= start: - raise ValueError( - f'Stop position ({stop}) must be greater than or ' - f'equal to start ({start}).', - ) + raise RegionStartPastStopError(start, stop) return Region(contig, start, stop) @classmethod def parse_string(cls, region_str: str) -> tuple[str, int, int | None]: - """ - Parse a region string into its components. + """Parse a region string into its components. Parameters ---------- @@ -139,23 +200,20 @@ def parse_string(cls, region_str: str) -> tuple[str, int, int | None]: Raises ------ - ValueError + RegionFormatError If the region string format is unrecognized. """ if region_str is None: return None pa = re.compile( - '(?P[^:]+)(:(?P[0-9,]+)(-(?P[0-9,]+))?)?', + "(?P[^:]+)(:(?P[0-9,]+)(-(?P[0-9,]+))?)?", ) match = pa.fullmatch(region_str) if match is None: - raise ValueError(f'Unrecognized format: {region_str}.') - contig, start, stop = match.group('contig', 'start', 'stop') + raise RegionFormatError(region_str) + contig, start, stop = match.group("contig", "start", "stop") - if start is None: - start = 0 - else: - start = Region._to_int(start) - 1 + start = 0 if start is None else Region._to_int(start) - 1 if stop is not None: stop = Region._to_int(stop) @@ -163,4 +221,4 @@ def parse_string(cls, region_str: str) -> tuple[str, int, int | None]: @classmethod def _to_int(cls, number: str) -> int: - return int(re.sub(r'[\s,]', '', number)) + return int(re.sub(r"[\s,]", "", number)) diff --git a/reditools/region_collection.py b/reditools/region_collection.py index 51741a3..790973b 100644 --- a/reditools/region_collection.py +++ b/reditools/region_collection.py @@ -1,27 +1,27 @@ +"""A collection of genomic regions, providing efficient ordered lookup.""" +from __future__ import annotations + from collections import defaultdict -from typing import DefaultDict, Iterable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Iterable -from reditools.region import Region + from reditools.region import Region class RegionCollection: - """ - A collection of genomic regions, providing efficient ordered lookup. - """ + """A collection of genomic regions, providing efficient ordered lookup.""" def __init__(self) -> None: - """ - Initialize an empty RegionCollection. - """ - - self._regions: DefaultDict[str, list[Region]] = defaultdict(list) + """Initialize an empty RegionCollection.""" + self._regions: defaultdict[str, list[Region]] = defaultdict(list) self._index = 0 self._last_contig: str | None = None self._sorted = False def __bool__(self) -> bool: - """ - Check whether the collection is empty. + """Check whether the collection is empty. Returns ------- @@ -31,17 +31,13 @@ def __bool__(self) -> bool: return bool(self._regions) def sort(self) -> None: - """ - Sort all regions within each contig. - """ + """Sort all regions within each contig.""" for contig, regions in self._regions.items(): self._regions[contig] = sorted(regions) self._sorted = True def contains(self, contig: str, position: int) -> bool: - """ - Check if a given position is contained within any region of the - collection. + """Check if a given position is contained within the collection. This method only works if each subsequent call is done in sorted order. Otherwise the output will be inconsistent. @@ -61,17 +57,17 @@ def contains(self, contig: str, position: int) -> bool: if not self._sorted: self.sort() self._last_contig = contig - idx = 0 + start = 0 elif contig != self._last_contig: self._last_contig = contig - idx = 0 + start = 0 else: - idx = self._index + start = self._index for idx, region in enumerate( - self._regions[contig][idx:], - start=idx, + self._regions[contig][start:], + start=start, ): if position < region.start: self._index = idx @@ -83,8 +79,7 @@ def contains(self, contig: str, position: int) -> bool: return False def add_regions(self, regions: Iterable[Region]) -> None: - """ - Add multiple regions to the collection. + """Add multiple regions to the collection. Parameters ---------- @@ -96,8 +91,7 @@ def add_regions(self, regions: Iterable[Region]) -> None: self._regions[_.contig].append(_) def get_contig(self, contig: str) -> list[Region]: - """ - Retrieve the regions for a specific chromosome/contig. + """Retrieve the regions for a specific chromosome/contig. Parameters ---------- @@ -112,9 +106,9 @@ def get_contig(self, contig: str) -> list[Region]: return self._regions[contig] def reset(self) -> None: - """ - Restart the search parameters. RegionCollection requires checks be - done in order. This moves the checks back to the beginning of the - collections. + """Restart the search parameters. + + RegionCollection requires checks be done in order. This function moves + the checks back to the beginning of the collections. """ self._last_contig = None diff --git a/reditools/rtannotater.py b/reditools/rtannotater.py index 6567587..6e987ff 100644 --- a/reditools/rtannotater.py +++ b/reditools/rtannotater.py @@ -1,8 +1,20 @@ +"""Class to annotate RNA editing sites with DNA data.""" +from __future__ import annotations + import csv from typing import IO, Iterator from reditools import file_utils +from reditools.constants import comp_map + +class AnalyzeMismatchError(ValueError): + """Reference bases from two REDItools output files do not match.""" + + def __init__(self) -> None: + """Initialize self.""" + self.message = "Files do not appear to use the same reference." + super().__init__(self.message) class RTAnnotater: """Class to annotate RNA editing sites with DNA data. @@ -17,17 +29,20 @@ class RTAnnotater: """ legacy_map = ( - ('Coverage-q30', 'Coverage'), - ('gCoverage-q30', 'gCoverage'), + ("Coverage-q30", "Coverage"), + ("gCoverage-q30", "gCoverage"), ) - comp_map = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', '-': '-'} - ref_key = 'Reference' - sub_key = 'AllSubs' - bases_key = 'BaseCount[A,C,G,T]' + ref_key = "Reference" + sub_key = "AllSubs" + bases_key = "BaseCount[A,C,G,T]" - def __init__(self, contig_order: dict[str, int], do_complement: bool=False): + def __init__( + self, + contig_order: dict[str, int], + do_complement: bool=False, + ) -> None: """Initialize RTAnnotater. Parameters @@ -53,21 +68,21 @@ def annotate(self, rna_file: str, dna_file: str, stream: IO) -> None: stream : IO The output stream to write the annotated table. """ - writer = csv.DictWriter(stream, delimiter='\t', fieldnames=[ - 'Region', - 'Position', + writer = csv.DictWriter(stream, delimiter="\t", fieldnames=[ + "Region", + "Position", self.ref_key, - 'Strand', - 'Coverage', - 'MeanQ', + "Strand", + "Coverage", + "MeanQ", self.bases_key, self.sub_key, - 'Frequency', - 'gCoverage', - 'gMeanQ', - 'gBaseCount[A,C,G,T]', - 'gAllSubs', - 'gFrequency']) + "Frequency", + "gCoverage", + "gMeanQ", + "gBaseCount[A,C,G,T]", + "gAllSubs", + "gFrequency"]) writer.writeheader() writer.writerows(self.merge_files(rna_file, dna_file)) @@ -92,15 +107,15 @@ def cmp_position( """ if dna_entry is None: return -1 - rna_contig_idx = self.contig_order[rna_entry['Region']] + rna_contig_idx = self.contig_order[rna_entry["Region"]] # If the DNA contig is not in the RNA file, assume its position is # earlier than the current RNA contig to induce fast-forwarding. dna_contig_idx = self.contig_order.get( - dna_entry['Region'], - 0 + dna_entry["Region"], + 0, ) if rna_contig_idx == dna_contig_idx: - return int(rna_entry['Position']) - int(dna_entry['Position']) + return int(rna_entry["Position"]) - int(dna_entry["Position"]) return rna_contig_idx - dna_contig_idx def annotate_row( @@ -122,20 +137,20 @@ def annotate_row( dict[str, str] The annotated RNA row. """ - if rna_row[self.ref_key] == self.comp_map[dna_row[self.ref_key]]: + if rna_row[self.ref_key] == comp_map[dna_row[self.ref_key]]: if self.do_complement: self.complement(dna_row) elif rna_row[self.ref_key] != dna_row[self.ref_key]: - raise ValueError('Files do not appear to use the same reference.') - rna_row['gCoverage'] = dna_row['Coverage'] - rna_row['gMeanQ'] = dna_row['MeanQ'] - rna_row['gBaseCount[A,C,G,T]'] = dna_row[self.bases_key] - rna_row['gAllSubs'] = dna_row[self.sub_key] - rna_row['gFrequency'] = dna_row['Frequency'] + raise AnalyzeMismatchError + rna_row["gCoverage"] = dna_row["Coverage"] + rna_row["gMeanQ"] = dna_row["MeanQ"] + rna_row["gBaseCount[A,C,G,T]"] = dna_row[self.bases_key] + rna_row["gAllSubs"] = dna_row[self.sub_key] + rna_row["gFrequency"] = dna_row["Frequency"] return rna_row @classmethod - def legacy_translate(cls, row: dict[str, str]) -> dict[str, str]: + def legacy_translate(cls, row: dict[str, str]) -> None: """Translate legacy field names to current ones. Parameters @@ -151,7 +166,6 @@ def legacy_translate(cls, row: dict[str, str]) -> dict[str, str]: for old_key, new_key in cls.legacy_map: if old_key in row: row[new_key] = row.pop(old_key) - return row def merge_files( self, @@ -160,8 +174,8 @@ def merge_files( ) -> Iterator[dict[str, str]]: """Merge RNA and DNA files and yield annotated rows. - Parameters: - ----------- + Parameters + ---------- rna_file : str Path to the RNA editing file. dna_file : str @@ -172,10 +186,10 @@ def merge_files( dict[str, str] Annotated (or original if no match) RNA row. """ - with file_utils.open_stream(rna_file, 'r') as rna_stream, \ - file_utils.open_stream(dna_file, 'r') as dna_stream: - rna_reader = csv.DictReader(rna_stream, delimiter='\t') - dna_reader = csv.DictReader(dna_stream, delimiter='\t') + with file_utils.open_stream(rna_file, "r") as rna_stream, \ + file_utils.open_stream(dna_file, "r") as dna_stream: + rna_reader = csv.DictReader(rna_stream, delimiter="\t") + dna_reader = csv.DictReader(dna_stream, delimiter="\t") dna_entry = next(dna_reader, None) @@ -184,20 +198,35 @@ def merge_files( while self.cmp_position(rna_entry, dna_entry) > 0: dna_entry = next(dna_reader, None) - if self.cmp_position(rna_entry, dna_entry) == 0: - assert dna_entry is not None + if dna_entry is not None and \ + self.cmp_position(rna_entry, dna_entry) == 0: self.legacy_translate(dna_entry) yield self.annotate_row(rna_entry, dna_entry) else: yield rna_entry def complement(self, row: dict[str, str]) -> dict[str, str]: - row[self.ref_key] = self.comp_map[row[self.ref_key]] - row[self.sub_key] = ' '.join(sorted([ - ''.join([self.comp_map[_] for _ in sub]) - for sub in row[self.sub_key].split(' ') + """Compute complements of REDItools result. + + Speifically, complements are reported for the Reference, AllSubs, and + BaseCount columns. + + Parameters + ---------- + row : dict[str, str] + The data to complement. + + Returns + ------- + row : dict[str, str] + The complemented data. + """ + row[self.ref_key] = comp_map[row[self.ref_key]] + row[self.sub_key] = " ".join(sorted([ + "".join([comp_map[_] for _ in sub]) + for sub in row[self.sub_key].split(" ") ])) - base_counts = row[self.bases_key][1:-1].split(', ') + base_counts = row[self.bases_key][1:-1].split(", ") row[self.bases_key] = str(list(reversed([ int(_) for _ in base_counts ]))) diff --git a/reditools/rtindexer.py b/reditools/rtindexer.py index 5a8810d..ce67faf 100644 --- a/reditools/rtindexer.py +++ b/reditools/rtindexer.py @@ -1,35 +1,28 @@ +"""Calculate editing indices from REDItools output.""" +from __future__ import annotations + import csv from itertools import permutations from typing import Iterator +from reditools.constants import bases from reditools.file_utils import open_stream, read_bed_file from reditools.region_collection import RegionCollection -class RTIndexer(object): - _ref = 'Reference' - _position = 'Position' - _contig = 'Region' - _count = 'BaseCount[A,C,G,T]' - _nucs = 'ACGT' - - - """ - Calculate editing indices from REDItools output. +class RTIndexer: + """Calculate editing indices from REDItools output.""" - Parameters - ---------- - region : tuple[str, int, int | None] | None, optional - Genomic region (contig, start, stop) to limit analysis (default is - None). - """ + _ref = "Reference" + _position = "Position" + _contig = "Region" + _count = "BaseCount[A,C,G,T]" def __init__( self, region: tuple[str, int, int | None] | None=None, - ): - """ - Initialize the RTIndexer. + ) -> None: + """Initialize the RTIndexer. Parameters ---------- @@ -40,14 +33,13 @@ def __init__( self.targets = RegionCollection() self.exclusions = RegionCollection() self.counts = { - '-'.join(_): 0 - for _ in permutations(self._nucs, 2) + "-".join(_): 0 + for _ in permutations(bases, 2) } self.region = region def add_target_from_bed(self, fname: str) -> None: - """ - Add target regions from a BED file. + """Add target regions from a BED file. Parameters ---------- @@ -57,8 +49,7 @@ def add_target_from_bed(self, fname: str) -> None: self.targets.add_regions(read_bed_file(fname)) def add_exclusions_from_bed(self, fname: str) -> None: - """ - Exclude regions from a BED file. + """Exclude regions from a BED file. Parameters ---------- @@ -68,8 +59,7 @@ def add_exclusions_from_bed(self, fname: str) -> None: self.exclusions.add_regions(read_bed_file(fname)) def do_ignore(self, row: dict) -> bool: - """ - Check if a row from REDItools output should be ignored. + """Check if a row from REDItools output should be ignored. Parameters ---------- @@ -85,7 +75,7 @@ def do_ignore(self, row: dict) -> bool: position = int(row[self._position]) if self.region[0] != row[self._contig] or \ self.region[1] > position or \ - self.region[2] is not None and self.region[2] < position: + (self.region[2] is not None and self.region[2] < position): return True if self.exclusions and self.exclusions.contains( row[self._contig], @@ -101,8 +91,7 @@ def do_ignore(self, row: dict) -> bool: def add_rt_output(self, fname: str) -> None: - """ - Add base counts from a REDItools output file. + """Add base counts from a REDItools output file. Parameters ---------- @@ -112,20 +101,19 @@ def add_rt_output(self, fname: str) -> None: self.targets.reset() self.exclusions.reset() with open_stream(fname) as stream: - for row in csv.DictReader(stream, delimiter='\t'): + for row in csv.DictReader(stream, delimiter="\t"): if self.do_ignore(row): continue for nuc, count in zip( - self._nucs, + bases, self._counts_to_list(row[self._count]), ): ref = row[self._ref] - key = f'{ref}-{nuc}' + key = f"{ref}-{nuc}" self.counts[key] = self.counts.get(key, 0) + count def calc_index(self) -> dict[str, float]: - """ - Calculate editing indices for all base transitions. + """Calculate editing indices for all base transitions. Returns ------- @@ -134,10 +122,10 @@ def calc_index(self) -> dict[str, float]: indices. """ indices: dict[str, float] = {} - for idx in set(self.counts) - {f'{nuc}-{nuc}' for nuc in self._nucs}: + for idx in set(self.counts) - {f"{nuc}-{nuc}" for nuc in bases}: ref = idx[0] numerator = self.counts[idx] - denominator = self.counts.get(f'{ref}-{ref}', 0) + numerator + denominator = self.counts.get(f"{ref}-{ref}", 0) + numerator if denominator == 0: indices[idx] = 0 else: @@ -146,5 +134,5 @@ def calc_index(self) -> dict[str, float]: @classmethod def _counts_to_list(cls, counts_str: str) -> Iterator[int]: - pieces = counts_str[1:-1].split(', ') + pieces = counts_str[1:-1].split(", ") return (int(_) for _ in pieces) diff --git a/reditools/splicing_file.py b/reditools/splicing_file.py index 45dc183..dd985ec 100644 --- a/reditools/splicing_file.py +++ b/reditools/splicing_file.py @@ -1,3 +1,5 @@ +"""Load genomic regions around splice sites from a file.""" +from __future__ import annotations import csv from typing import IO, Iterator @@ -6,23 +8,40 @@ from reditools.region import Region +class SpliceFileFormatError(ValueError): + """Splice file is not in expected format.""" + + def __init__(self, file_name: str, line_number: int) -> None: + """Initialize self. + + Parameters + ---------- + file_name : str + The offending file. + line_number : int + The line number that does not match expected format. + """ + self.message = ( + f"Cannot parse splice file entry ({file_name}:{line_number})" + ) + super().__init__(self.message) + def _read_splice_sites( # noqa: WPS231 stream: IO, ) -> Iterator[tuple[str, int, str, str]]: - reader = csv.reader(stream, delimiter=' ') + reader = csv.reader(stream, delimiter=" ") for idx, row in enumerate(reader, start=1): - if row[0].startswith('#'): + if row[0].startswith("#"): continue - try: # noqa: WPS229 - assert len(row) == 5 - assert row[3] in ('A', 'D') - assert row[4] in ('+', '-') + if len(row) != 5 or \ + row[3] not in ("A", "D") or \ + row[4] not in ("+", "-"): # noqa: PLR2004 + raise SpliceFileFormatError(stream.name, idx) + try: position = int(row[1]) - yield (row[0], position, row[3], row[4]) - except (AssertionError, ValueError) as exc: - raise ValueError( - f'Cannot parse splice file entry ({stream.name}:{idx})' - ) from exc + except ValueError as exc: + raise SpliceFileFormatError(stream.name, idx) from exc + yield (row[0], position, row[3], row[4]) def _splice_site_to_region( contig: str, @@ -31,7 +50,7 @@ def _splice_site_to_region( strand: str, splicing_span: int, ) -> Region | None: - strand_map = {'-': 'D', '+': 'A'} + strand_map = {"-": "D", "+": "A"} position = position - 1 if strand_map[strand] == splice: start = max(position - splicing_span, 0) @@ -47,8 +66,7 @@ def load_splicing_file( splicing_file: str, splicing_span: int, ) -> Iterator[Region]: - """ - Load genomic regions around splice sites from a file. + """Load genomic regions around splice sites from a file. Splice site files are space delimited and have five columns: 1. Chromosome/contig @@ -69,7 +87,6 @@ def load_splicing_file( Region The genomic regions around the splice sites. """ - with open_stream(splicing_file) as stream: for splice_data in _read_splice_sites(stream): region = _splice_site_to_region(*splice_data, splicing_span) diff --git a/reditools/tools/analyze/__init__.py b/reditools/tools/analyze/__init__.py deleted file mode 100644 index 28de984..0000000 --- a/reditools/tools/analyze/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from reditools.tools.analyze.main import main diff --git a/reditools/tools/analyze/concat_output.py b/reditools/tools/analyze/concat_output.py index 0dbfdc1..1d120cb 100644 --- a/reditools/tools/analyze/concat_output.py +++ b/reditools/tools/analyze/concat_output.py @@ -1,3 +1,5 @@ +"""Concatenate temporary results files into the final output.""" +from __future__ import annotations import csv import sys @@ -5,28 +7,28 @@ from reditools import file_utils fieldnames = [ - 'Region', - 'Position', - 'Reference', - 'Strand', - 'Coverage', - 'MeanQ', - 'BaseCount[A,C,G,T]', - 'AllSubs', - 'Frequency', - 'gCoverage', - 'gMeanQ', - 'gBaseCount[A,C,G,T]', - 'gAllSubs', - 'gFrequency', + "Region", + "Position", + "Reference", + "Strand", + "Coverage", + "MeanQ", + "BaseCount[A,C,G,T]", + "AllSubs", + "Frequency", + "gCoverage", + "gMeanQ", + "gBaseCount[A,C,G,T]", + "gAllSubs", + "gFrequency", ] def concat_output( tfs: list[str], output_file: str | None = None, - mode: str = 'w', - encoding: str = 'utf-8', + mode: str = "w", + encoding: str = "utf-8", ) -> None: """Concatenate temporary results files into the final output. @@ -52,7 +54,7 @@ def concat_output( ) with stream: - writer = csv.writer(stream, delimiter='\t', lineterminator='\n') - if 'a' not in mode: + writer = csv.writer(stream, delimiter="\t", lineterminator="\n") + if "a" not in mode: writer.writerow(fieldnames) file_utils.concat(stream, *tfs, encoding=encoding) diff --git a/reditools/tools/analyze/main.py b/reditools/tools/analyze/main.py index 69fb1ed..5db8240 100644 --- a/reditools/tools/analyze/main.py +++ b/reditools/tools/analyze/main.py @@ -1,3 +1,4 @@ +"""REDItools analyze tool entry point form CLI.""" from __future__ import annotations import sys @@ -14,9 +15,7 @@ import argparse def main() -> None: - """ - The main entry point for the REDItools analyze command. - """ + """Begin REDItools analyze from CLI.""" options = parse_args.parse_args() logger = setup_logger(options) @@ -25,18 +24,18 @@ def main() -> None: logger.log( logger.info_level, ( - 'Resuming REDItools from directory "{}". Using parameters ' - 'from previous run. All other command line options will be ' - 'ignored.' + "Resuming REDItools from directory '{}'. Using parameters " + "from previous run. All other command line options will be " + "ignored." ), options.temp_dir, ) temp_dir = options.temp_dir else: - logger.log(logger.info_level, 'Starting REDItools') + logger.log(logger.info_level, "Starting REDItools") temp_dir = file_utils.make_dir( - prefix='reditools_', - dir=options.temp_dir, + prefix="reditools_", + dirname=options.temp_dir, ) json_args.args_to_json(options, temp_dir) @@ -53,7 +52,7 @@ def main() -> None: ) if analyze(options, temp_dir): - logger.log(Logger.info_level, 'Analyze Complete!') + logger.log(Logger.info_level, "Analyze Complete!") else: sys.exit(1) @@ -102,7 +101,7 @@ def analyze( sys.stderr.write( f"[WARNING] You have assigned {options.threads} threads, " f"But there are only {len(temp_file_manager)} genomic " - "range(s). Consider change the value of --window\n" + "range(s). Consider change the value of --window\n", ) options.threads = len(temp_file_manager) @@ -111,6 +110,6 @@ def analyze( temp_file_manager.concat( options.output_file, - 'a' if options.append_file else 'w', + "a" if options.append_file else "w", ) return True diff --git a/reditools/tools/analyze/parse_args/bounded_types.py b/reditools/tools/analyze/parse_args/bounded_types.py new file mode 100644 index 0000000..6b22a3e --- /dev/null +++ b/reditools/tools/analyze/parse_args/bounded_types.py @@ -0,0 +1,143 @@ +"""Validation tools for CLI options.""" +from __future__ import annotations + +import argparse +from typing import Callable + + +class ValueBelowMinimumError(argparse.ArgumentTypeError): + """CLI value is below minimum threshold.""" + + def __init__(self, min_value: float) -> None: + """Initialize self. + + Parameters + ---------- + min_value : float + The minimum threshold. + """ + self.message = f"Value must be at least {min_value}." + super().__init__(self.message) + +class ValueAboveMaximumError(argparse.ArgumentTypeError): + """CLI value is above maxmimum threshold.""" + + def __init__(self, max_value: float) -> None: + """Initialize self. + + Parameters + ---------- + max_value : float + The maxmimum threshold. + """ + self.message = f"Value cannot be larger than {max_value}." + super().__init__(self.message) + +class CastIntError(argparse.ArgumentTypeError): + """CLI value is not an integer.""" + + def __init__(self, cli_val: str) -> None: + """Initialize self. + + Parameters + ---------- + cli_val : str + Offending CLI value. + """ + self.message = f"Invalid int value: {cli_val}" + super().__init__(self.message) + +class CastFloatError(argparse.ArgumentTypeError): + """CLI value is not a float.""" + + def __init__(self, cli_val: str) -> None: + """Initialize self. + + Parameters + ---------- + cli_val : str + Offending CLI value. + """ + self.message = f"Invalid float value: {cli_val}" + super().__init__(self.message) + +def check_number_bounds( + number: float, + min_value: float | None = None, + max_value: float | None = None, +) -> None: + """Check if a number is within specified bounds. + + Parameters + ---------- + number : float + The number to check. + min_value : float | None, optional + The minimum allowed value, by default None. + max_value : float | None, optional + The maximum allowed value, by default None. + + Raises + ------ + ValueBelowMinimumError, ValueAboveMaximumError + If the number is outside the specified bounds. + """ + if min_value is not None and number < min_value: + raise ValueBelowMinimumError(min_value) + if max_value is not None and number > max_value: + raise ValueAboveMaximumError(max_value) + +def bounded_int( + min_value: int | None = None, + max_value: int | None = None, +) -> Callable: + """Create a function that parses a string to a bounded integer. + + Parameters + ---------- + min_value : int | None, optional + The minimum allowed value, by default None. + max_value : int | None, optional + The maximum allowed value, by default None. + + Returns + ------- + Callable + A function that takes a string and returns a bounded integer. + """ + def subfn(cli_value: str) -> int: # noqa: WPS430 + try: + int_value = int(cli_value) + except ValueError as exc: + raise CastIntError(cli_value) from exc + check_number_bounds(int_value, min_value, max_value) + return int_value + return subfn + + +def bounded_float( + min_value: float | None = None, + max_value: float | None = None, +) -> Callable: + """Create a function that parses a string to a bounded float. + + Parameters + ---------- + min_value : float | None, optional + The minimum allowed value, by default None. + max_value : float | None, optional + The maximum allowed value, by default None. + + Returns + ------- + Callable + A function that takes a string and returns a bounded float. + """ + def subfn(cli_value: str) -> float: # noqa: WPS430 + try: + float_value = float(cli_value) + except ValueError as exc: + raise CastFloatError(cli_value) from exc + check_number_bounds(float_value, min_value, max_value) + return float_value + return subfn diff --git a/reditools/tools/analyze/parse_args/json_args.py b/reditools/tools/analyze/parse_args/json_args.py index 5679e56..05384c3 100644 --- a/reditools/tools/analyze/parse_args/json_args.py +++ b/reditools/tools/analyze/parse_args/json_args.py @@ -1,34 +1,34 @@ +"""Save and load CLI options using JSON files.""" import argparse import json -import os +from pathlib import Path -json_args_filename = 'cli_args.json' +json_args_filename = "cli_args.json" def args_to_json( options: argparse.Namespace, dirname: str, filename: str=json_args_filename, ) -> None: - """ - Save commandline arguments to a JSON file. + """Save commandline arguments to a JSON file. - Parameters: - options : argparse.Namespace - The parsed commandline options. - dirname : str - Path to save file to. - filename : str - Name of the file (defaults to json_args_filename) + Parameters + ---------- + options : argparse.Namespace + The parsed commandline options. + dirname : str + Path to save file to. + filename : str + Name of the file (defaults to json_args_filename) """ - with open(os.path.join(dirname, filename), 'w') as stream: + with Path(dirname, filename).open("w") as stream: json.dump(vars(options), stream) # noqa: WPS421 def args_from_json( dirname: str, filename: str=json_args_filename, ) -> argparse.Namespace: - """ - Load commandline arguments from a JSON file. + """Load commandline arguments from a JSON file. Parameters ---------- @@ -42,6 +42,6 @@ def args_from_json( argparse.Namespace Commandline arguments for reditools analyze """ - with open(os.path.join(dirname, filename), 'r') as stream: + with Path(dirname, filename).open("r") as stream: json_args = json.load(stream) return argparse.Namespace(**json_args) diff --git a/reditools/tools/analyze/parse_args/parse_args.py b/reditools/tools/analyze/parse_args/parse_args.py index ae24eb1..dd56038 100644 --- a/reditools/tools/analyze/parse_args/parse_args.py +++ b/reditools/tools/analyze/parse_args/parse_args.py @@ -1,96 +1,33 @@ +"""CLI argument parsing.""" +from __future__ import annotations + import argparse +import json import tempfile -from typing import Callable from reditools import reditools +from reditools.tools.analyze.parse_args.bounded_types import ( + bounded_float, + bounded_int, +) from reditools.tools.analyze.parse_args.json_args import args_from_json -def check_number_bounds( - number: float, - min_value: float | None = None, - max_value: float | None = None, -) -> None: - """Check if a number is within specified bounds. - - Parameters - ---------- - number : float - The number to check. - min_value : float | None, optional - The minimum allowed value, by default None. - max_value : float | None, optional - The maximum allowed value, by default None. +class DNAStrandError(argparse.ArgumentTypeError): + """DNA mode requires strand set to 0.""" - Raises - ------ - argparse.ArgumentTypeError - If the number is outside the specified bounds. - """ - if min_value is not None and number < min_value: - raise argparse.ArgumentTypeError(f'Value must be at least {min_value}.') - if max_value is not None and number > max_value: - raise argparse.ArgumentTypeError( - f'Value cannot be larger than {max_value}.', - ) + def __init__(self) -> None: + """Initialize self.""" + self.message = "-N/--dna can only be used with -s/--strand 0." + super().__init__(self.message) -def bounded_int( - min_value: int | None = None, - max_value: int | None = None, -) -> Callable: - """Create a function that parses a string to a bounded integer. - - Parameters - ---------- - min_value : int | None, optional - The minimum allowed value, by default None. - max_value : int | None, optional - The maximum allowed value, by default None. - - Returns - ------- - Callable - A function that takes a string and returns a bounded integer. - """ - def subfn(cli_value: str) -> int: # noqa: WPS430 - try: - int_value = int(cli_value) - except ValueError: - raise argparse.ArgumentTypeError(f'invalid int value: {cli_value}') - check_number_bounds(int_value, min_value, max_value) - return int_value - return subfn - - -def bounded_float( - min_value: float | None = None, - max_value: float | None = None, -) -> Callable: - """Create a function that parses a string to a bounded float. - - Parameters - ---------- - min_value : float | None, optional - The minimum allowed value, by default None. - max_value : float | None, optional - The maximum allowed value, by default None. - - Returns - ------- - Callable - A function that takes a string and returns a bounded float. - """ - def subfn(cli_value: str) -> float: # noqa: WPS430 - try: - float_value = float(cli_value) - except ValueError: - raise argparse.ArgumentTypeError( - f'invalid float value: {cli_value}' - ) - check_number_bounds(float_value, min_value, max_value) - return float_value - return subfn +class StrictConflictError(argparse.ArgumentTypeError): + """Strict mode requires min-edits set to 1.""" + def __init__(self) -> None: + """Initialize self.""" + self.message = "-S/--strict can only be used with -me/--min-edits 1." + super().__init__(self.message) def build_argument_parser() -> argparse.ArgumentParser: # noqa: WPS213, WPS210 """Build the argument parser for reditools analyze. @@ -102,190 +39,190 @@ def build_argument_parser() -> argparse.ArgumentParser: # noqa: WPS213, WPS210 """ parser = argparse.ArgumentParser( prog="reditools analyze", - description='REDItools3', + description="REDItools3", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( - 'file', - nargs='+', + "file", + nargs="+", help=( - 'The BAM file(s) to be analyzed. BAM files must be sorted and ' - 'indexed.' + "The BAM file(s) to be analyzed. BAM files must be sorted and " + "indexed." ), ) parser.add_argument( - '-r', - '--reference', + "-r", + "--reference", help=( - 'Reference genome FASTA file. (Note: REDItools runs fastest when ' - 'BAM files have MD tags and -r is *not* used).' + "Reference genome FASTA file. (Note: REDItools runs fastest when " + "BAM files have MD tags and -r is *not* used)." ), ) parser.add_argument( - '-g', - '--region', + "-g", + "--region", help=( - 'Only analyzes the specified samtools formatted region. ' - '(1-index, start and end inclusive).' + "Only analyzes the specified samtools formatted region. " + "(1-index, start and end inclusive)." ), ) parser.add_argument( - '--resume', + "--resume", help=( - 'If REDItools crashed, will attempt to resume a stopped job from ' - 'existing temporary files. Note: --temp-dir is required.' + "If REDItools crashed, will attempt to resume a stopped job from " + "existing temporary files. Note: --temp-dir is required." ), - action='store_true', + action="store_true", ) output_group = parser.add_argument_group( - title='Output Options', + title="Output Options", ) output_group.add_argument( - '-o', - '--output-file', - help='Path to write output to.', - default='/dev/stdout', + "-o", + "--output-file", + help="Path to write output to.", + default="/dev/stdout", ) output_group.add_argument( - '-a', - '--append-file', - action='store_true', - help='Appends results to file (and creates if not existing).', + "-a", + "--append-file", + action="store_true", + help="Appends results to file (and creates if not existing).", ) bqf_group = parser.add_argument_group( - title='Base/Read Quality Controls', + title="Base/Read Quality Controls", ) bqf_group.add_argument( - '-mrl', - '--min-read-length', + "-mrl", + "--min-read-length", type=int, default=30, - help='Reads shorter than -mrl will be discarded.', + help="Reads shorter than -mrl will be discarded.", ) bqf_group.add_argument( - '-q', - '--min-read-quality', + "-q", + "--min-read-quality", type=int, default=20, - help='Reads with MAPQ below -q will be discarded.', + help="Reads with MAPQ below -q will be discarded.", ) bqf_group.add_argument( - '-bq', - '--min-base-quality', + "-bq", + "--min-base-quality", type=int, default=30, - help='Bases with a Phred quality score below -bq will bed discarded.', + help="Bases with a Phred quality score below -bq will bed discarded.", ) bqf_group.add_argument( - '-mbp', - '--min-base-position', + "-mbp", + "--min-base-position", type=bounded_int(min_value=0), default=0, - help='Ignores the first -mbp bases in each read.', + help="Ignores the first -mbp bases in each read.", ) bqf_group.add_argument( - '-Mbp', - '--max-base-position', + "-Mbp", + "--max-base-position", type=bounded_int(min_value=0), default=0, - help='Ignores the last -Mpb bases in each read.', + help="Ignores the last -Mpb bases in each read.", ) bqf_group.add_argument( - '-E', - '--exclude-reads', - help='Text file listing read names to exclude from analysis.', + "-E", + "--exclude-reads", + help="Text file listing read names to exclude from analysis.", ) bqf_group.add_argument( - '--exclude_reads', + "--exclude_reads", help=argparse.SUPPRESS, ) gr_group = parser.add_argument_group( - title='Genomic Region Filters', + title="Genomic Region Filters", ) gr_group.add_argument( - '-k', - '--exclude-regions', - nargs='+', - help='Do not report on regions in the provided BED file(s).', + "-k", + "--exclude-regions", + nargs="+", + help="Do not report on regions in the provided BED file(s).", ) gr_group.add_argument( - '--exclude_regions', - nargs='+', + "--exclude_regions", + nargs="+", help=argparse.SUPPRESS, ) gr_group.add_argument( - '-B', - '--bed-file', - nargs='+', - help='Only reports on regions in the provided BED file(s).', + "-B", + "--bed-file", + nargs="+", + help="Only reports on regions in the provided BED file(s).", ) gr_group.add_argument( - '--bed_file', - nargs='+', + "--bed_file", + nargs="+", help=argparse.SUPPRESS, ) rf_group = parser.add_argument_group( - title='Result Filters', + title="Result Filters", ) rf_group.add_argument( - '-men', - '--min-edits-per-nucleotide', + "-men", + "--min-edits-per-nucleotide", type=int, default=0, help=( - 'Position where any variant has a frequency less than -men but ' - 'more than zero will not be reported. (Corresponds to the ' - 'BaseCount column.)' + "Position where any variant has a frequency less than -men but " + "more than zero will not be reported. (Corresponds to the " + "BaseCount column.)" ), ) rf_group.add_argument( - '-me', - '--min-edits', + "-me", + "--min-edits", type=bounded_int(0, 4), default=1, help=( - 'Positions with fewer than -me unique variants (listed in the ' - 'AllSubs column) will be excluded from the results.' + "Positions with fewer than -me unique variants (listed in the " + "AllSubs column) will be excluded from the results." ), ) rf_group.add_argument( - '-Men', - '--max-editing-nucleotides', + "-Men", + "--max-editing-nucleotides", type=bounded_int(min_value=0, max_value=4), default=4, help=( - 'Positions with more than -Men unique variants (listed in the ' - 'AllSubs column) will be excluded from the results.' + "Positions with more than -Men unique variants (listed in the " + "AllSubs column) will be excluded from the results." ), ) rf_group.add_argument( - '-v', - '--variants', - nargs='*', - default=['all'], + "-v", + "--variants", + nargs="*", + default=["all"], help=( - 'Which editing events to report. Each edit should be two ' - 'characters and separated by spaces (e.g. AG CT). Use "all" to ' - 'report all variants. (Corresponds to the AllSubs column)' + "Which editing events to report. Each edit should be two " + "characters and separated by spaces (e.g. AG CT). Use 'all' to " + "report all variants. (Corresponds to the AllSubs column)" ), ) rf_group.add_argument( - '-l', - '--min-read-depth', + "-l", + "--min-read-depth", type=bounded_int(min_value=1), default=1, help=( - 'Only report on positions with at least -l reads (corresponds to ' - 'the Coverage column.)' + "Only report on positions with at least -l reads (corresponds to " + "the Coverage column.)" ), ) strand_group = parser.add_argument_group( - title='Strandedness Options', + title="Strandedness Options", ) strand_group.add_argument( - '-s', - '--strand', + "-s", + "--strand", choices=( reditools.UNSTRANDED_MODE, reditools.FORWARD_STRAND_MODE, @@ -294,150 +231,150 @@ def build_argument_parser() -> argparse.ArgumentParser: # noqa: WPS213, WPS210 type=int, default=reditools.UNSTRANDED_MODE, help=( - f'Infer RNA strand and filter reads not of the same strand. ' - f'This option may be {reditools.UNSTRANDED_MODE} (unstranded), ' - f'{reditools.FORWARD_STRAND_MODE} (read1 is original RNA), or ' - f'{reditools.REVERSE_STRAND_MODE} (read2 is original RNA). ' - 'From RSeQC infer_experiment.py, 1++,1--,2+-,2-+ should be run ' - f'as --strand {reditools.FORWARD_STRAND_MODE} and 1+-,1-+,2++,2-- ' - f'should be run as --strand {reditools.REVERSE_STRAND_MODE}. ' - 'From Salmon, forward libraries (ISF, MSF, OSF) should be run as ' - f'--strand {reditools.FORWARD_STRAND_MODE} and reverse libraries ' - f'(ISR, MSR, OSR) as --strand {reditools.REVERSE_STRAND_MODE}. ' - 'All DNA sequencing experiments and non-stranded experiments ' - f'should be run with --strand {reditools.UNSTRANDED_MODE}.' + f"Infer RNA strand and filter reads not of the same strand. " + f"This option may be {reditools.UNSTRANDED_MODE} (unstranded), " + f"{reditools.FORWARD_STRAND_MODE} (read1 is original RNA), or " + f"{reditools.REVERSE_STRAND_MODE} (read2 is original RNA). " + "From RSeQC infer_experiment.py, 1++,1--,2+-,2-+ should be run " + f"as --strand {reditools.FORWARD_STRAND_MODE} and 1+-,1-+,2++,2-- " + f"should be run as --strand {reditools.REVERSE_STRAND_MODE}. " + "From Salmon, forward libraries (ISF, MSF, OSF) should be run as " + f"--strand {reditools.FORWARD_STRAND_MODE} and reverse libraries " + f"(ISR, MSR, OSR) as --strand {reditools.REVERSE_STRAND_MODE}. " + "All DNA sequencing experiments and non-stranded experiments " + f"should be run with --strand {reditools.UNSTRANDED_MODE}." ), ) strand_group.add_argument( - '-T', - '--strand-confidence-threshold', + "-T", + "--strand-confidence-threshold", type=bounded_float(max_value=1), default=0.7, help=( - 'Only report the strandedness if at least -T proportion of ' - 'reads are of a given strand. This option is only applicable ' - 'if -s/--strand is not zero.' + "Only report the strandedness if at least -T proportion of " + "reads are of a given strand. This option is only applicable " + "if -s/--strand is not zero." ), ) strand_group.add_argument( - '-C', - '--strand-correction', + "-C", + "--strand-correction", default=False, help=( - 'Report the base complements for the Reference, AllSubs, and ' - 'BaseCount columns in the output if the detected edit is on ' - 'the minus strand. ' - 'This option is only applicable if -s/--strand is not zero.' + "Report the base complements for the Reference, AllSubs, and " + "BaseCount columns in the output if the detected edit is on " + "the minus strand. " + "This option is only applicable if -s/--strand is not zero." ), - action='store_true', + action="store_true", ) para_group = parser.add_argument_group( - title='Parallel Processing Options', + title="Parallel Processing Options", ) para_group.add_argument( - '-t', - '--threads', + "-t", + "--threads", help=( - 'Number of threads for parallel processing. Note that the ' - 'maximum number of usable threads is equivalent to the number of ' - 'chromosomes in your alignment genome unless you use the --window ' - 'option.' + "Number of threads for parallel processing. Note that the " + "maximum number of usable threads is equivalent to the number of " + "chromosomes in your alignment genome unless you use the --window " + "option." ), type=bounded_int(min_value=1), default=1, ) para_group.add_argument( - '-w', - '--window', + "-w", + "--window", help=( - 'How many bp should be processed by each thread at a time. ' - 'Zero uses the full contig.' + "How many bp should be processed by each thread at a time. " + "Zero uses the full contig." ), type=bounded_int(min_value=0), default=0, ) tech_group = parser.add_argument_group( - title='Technical Options', + title="Technical Options", ) tech_group.add_argument( - '-V', - '--verbose', + "-V", + "--verbose", default=False, - help='Run in verbose mode.', - action='store_true', + help="Run in verbose mode.", + action="store_true", ) tech_group.add_argument( - '-d', - '--debug', + "-d", + "--debug", default=False, help=( - 'Run in debug mode. Every step of REDItools logic will be printed ' - 'to STDERR. If REDItools crashes, debug mode will print the stack ' - 'trace as well.' + "Run in debug mode. Every step of REDItools logic will be printed " + "to STDERR. If REDItools crashes, debug mode will print the stack " + "trace as well." ), - action='store_true', + action="store_true", ) tech_group.add_argument( - '--temp-dir', - help='Location to save temporary files', + "--temp-dir", + help="Location to save temporary files", default=tempfile.gettempdir(), ) leg_group = parser.add_argument_group( - title='Legacy Options', + title="Legacy Options", ) leg_group.add_argument( - '-e', - '--exclude-multis', + "-e", + "--exclude-multis", default=False, help=( - 'Do not report any position with more than one alternate base. ' - '(Equivalent to -Men/--max-editing-nucleotides 1)' + "Do not report any position with more than one alternate base. " + "(Equivalent to -Men/--max-editing-nucleotides 1)" ), - action='store_true', + action="store_true", ) leg_group.add_argument( - '-N', - '--dna', + "-N", + "--dna", default=False, - help='Run REDItools on DNA-Seq data. (Equivalent to -s/--strand 0)', - action='store_true', + help="Run REDItools on DNA-Seq data. (Equivalent to -s/--strand 0)", + action="store_true", ) leg_group.add_argument( - '-m', - '--load-omopolymeric-file', + "-m", + "--load-omopolymeric-file", help=( - 'BED file of homopolymeric positions. Regions in the BED file ' - 'will be excluded from the analysis. (Same effect as providing ' - 'the BED file with the -k/--exclude-regions option.' + "BED file of homopolymeric positions. Regions in the BED file " + "will be excluded from the analysis. (Same effect as providing " + "the BED file with the -k/--exclude-regions option." ), ) leg_group.add_argument( - '-S', - '--strict', + "-S", + "--strict", help=( - 'Activate strict mode: only sites with edits will be included in ' - 'the output. (Equivalent to -me/--min-edits 1)' + "Activate strict mode: only sites with edits will be included in " + "the output. (Equivalent to -me/--min-edits 1)" ), - action='store_true', + action="store_true", ) leg_group.add_argument( - '-sf', - '--splicing-file', + "-sf", + "--splicing-file", help=( - 'The splicing file is a space delimited file with five columns: ' - 'chromosome, start position (one-index inclusive), stop ' - '(ignored), splice (either A for acceptor or D for donor), and ' - 'strand (either + or -). A header is optional, but must start ' - 'with #. Used in conjunctions with -ss/--splicing-span.' + "The splicing file is a space delimited file with five columns: " + "chromosome, start position (one-index inclusive), stop " + "(ignored), splice (either A for acceptor or D for donor), and " + "strand (either + or -). A header is optional, but must start " + "with #. Used in conjunctions with -ss/--splicing-span." ), ) leg_group.add_argument( - '-ss', - '--splicing-span', + "-ss", + "--splicing-span", type=bounded_int(min_value=1), default=4, help=( - 'The splicing span. Used in conjunction with -sf/--splicing-file.' + "The splicing span. Used in conjunction with -sf/--splicing-file." ), ) @@ -459,29 +396,26 @@ def fix_legacy_options(args: argparse.Namespace) -> None: Raises ------ - Exception + DNAStrandError, StrictConflictError If mutually exclusive options are provided. """ if args.strand != 0 and args.dna: - raise Exception('-N/--dna can only be used with -s/--strand 0.') - delattr(args, 'dna') # noqa: WPS421 + raise DNAStrandError + delattr(args, "dna") # noqa: WPS421 if args.exclude_multis: - setattr(args, 'max_editing_nucleotides', 1) - delattr(args, 'exclude_multis') # noqa: WPS421 + args.max_editing_nucleotides = 1 + delattr(args, "exclude_multis") # noqa: WPS421 - if args.strict: - if args.min_edits != 1: - raise Exception( - '-S/--strict can only be used with -me/--min-edits 1.' - ) - delattr(args, 'strict') # noqa: WPS421 + if args.strict and args.min_edits != 1: + raise StrictConflictError + delattr(args, "strict") # noqa: WPS421 if args.load_omopolymeric_file: if args.exclude_regions is None: args.exclude_regions = [] args.exclude_regions.append(args.load_omopolymeric_file) - delattr(args, 'load_omopolymeric_file') # noqa: WPS421 + delattr(args, "load_omopolymeric_file") # noqa: WPS421 def parse_args(sys_args: list[str] | None = None) -> argparse.Namespace: """Parse command-line arguments for reditools analyze. @@ -503,33 +437,32 @@ def parse_args(sys_args: list[str] | None = None) -> argparse.Namespace: temp_dir = args.temp_dir try: args = args_from_json(temp_dir) - except Exception as exc: - parser.error(f'Unable to resume analysis.\n{exc}') + except (json.JSONDecodeError, OSError) as exc: + parser.error(f"Unable to resume analysis.\n{exc}") args.resume = True args.temp_dir = temp_dir return args try: fix_legacy_options(args) - except Exception as exc: + except argparse.ArgumentTypeError as exc: parser.error(message=str(exc)) if args.max_editing_nucleotides < args.min_edits: parser.error( - '-Men/--max-editing-nucleotides cannot be smaller than ' - '-me/--min-edits.', + "-Men/--max-editing-nucleotides cannot be smaller than " + "-me/--min-edits.", ) if args.strand == 0 and args.strand_correction: parser.error( - '-s/--strand 0 and -C/--strand-correction are mutually exclusive.' + "-s/--strand 0 and -C/--strand-correction are mutually exclusive.", ) return args def args_to_string(args: argparse.Namespace) -> str: - """ - Convert argparse options to a comma-separated string of key:value pairs. + """Convert argparse options to a comma-separated string of key:value pairs. Parameters ---------- diff --git a/reditools/tools/analyze/redi_pool.py b/reditools/tools/analyze/redi_pool.py index 1663c91..ecbe033 100644 --- a/reditools/tools/analyze/redi_pool.py +++ b/reditools/tools/analyze/redi_pool.py @@ -1,8 +1,9 @@ +"""Manage multirprocessing Pool for REDItools analysis.""" import argparse import sys import traceback from functools import partial -from multiprocessing.context import TimeoutError +from multiprocessing.context import TimeoutError as MPTimeoutError from multiprocessing.pool import Pool from reditools.tools.analyze.redi_thread import REDIThreadManager @@ -13,8 +14,7 @@ def run_pool( options: argparse.Namespace, temp_filemanager: TempFileManager, ) -> bool: - """ - Create a pool of threads and analyze the data. + """Create a pool of threads and analyze the data. Parameters ---------- @@ -44,9 +44,9 @@ def run_pool( pool.close() pool.join() [_.get(1) for _ in imap_iter] - except TimeoutError: + except MPTimeoutError: return False - except Exception: + except Exception: # noqa: BLE001 if options.debug: traceback.print_exception(*sys.exc_info()) return False @@ -57,8 +57,7 @@ def terminate_pool( debug: bool, exc: Exception, ) -> None: - """ - Terminates a multiprocessing Pool. + """Terminates a multiprocessing Pool. Parameters ---------- @@ -72,6 +71,6 @@ def terminate_pool( pool.terminate() if debug: raise exc.__cause__ # type: ignore[misc] - sys.stderr.write(f'[ERROR] ({type(exc)}) {exc}\n') + sys.stderr.write(f"[ERROR] ({type(exc)}) {exc}\n") diff --git a/reditools/tools/analyze/redi_thread.py b/reditools/tools/analyze/redi_thread.py index 11d78a6..913b3ab 100644 --- a/reditools/tools/analyze/redi_thread.py +++ b/reditools/tools/analyze/redi_thread.py @@ -1,15 +1,32 @@ -import argparse +"""Create and manage threads for REDItools analyze tool.""" +from __future__ import annotations + from pathlib import Path +from typing import TYPE_CHECKING -from reditools.region import Region from reditools.tools.analyze.rtchecks import RTChecks -from reditools.tools.analyze.setup_alignment_manager import \ - setup_alignment_manager +from reditools.tools.analyze.setup_alignment_manager import ( + setup_alignment_manager, +) from reditools.tools.analyze.setup_rtools import setup_rtools from reditools.tools.analyze.write_results import write_results +if TYPE_CHECKING: + import argparse + + from reditools.region import Region + +class UninitializedError(AttributeError): + """REDIThread.init has not been called.""" + + def __init__(self) -> None: + """Initialize self.""" + self.message = "REDIThreadManager not initialized." + super().__init__(self.message) class REDIThread: + """Worker thread for parallel REDItools analysis.""" + def __init__(self, options: argparse.Namespace) -> None: """Worker thread function for parallel REDItools analysis. @@ -64,7 +81,6 @@ def init_thread(cls, options: argparse.Namespace) -> None: options : argparse.Namespace The command-line options. """ - cls.thread = REDIThread(options) @classmethod @@ -78,10 +94,9 @@ def analyze(cls, region: Region, filename: str) -> None: filename : str Path to save output to. """ - if cls.thread is None: - raise AttributeError('REDIThreadManager not initialized.') - done_file = f'{filename}.done' + raise UninitializedError + done_file = f"{filename}.done" if not Path(done_file).exists(): cls.thread.analyze(region, filename) Path(done_file).touch() diff --git a/reditools/tools/analyze/region_args.py b/reditools/tools/analyze/region_args.py index e24450e..9831cfc 100644 --- a/reditools/tools/analyze/region_args.py +++ b/reditools/tools/analyze/region_args.py @@ -1,9 +1,15 @@ -import argparse +"""Parse region-related arguments and return a list of Regions.""" +from __future__ import annotations + +from typing import TYPE_CHECKING from pysam import AlignmentFile from reditools.region import Region +if TYPE_CHECKING: + import argparse + def region_args(options: argparse.Namespace) -> list[Region]: """Parse region-related arguments and return a list of Regions. diff --git a/reditools/tools/analyze/rtchecks/__init__.py b/reditools/tools/analyze/rtchecks/__init__.py index 7cf8e3d..651ceca 100644 --- a/reditools/tools/analyze/rtchecks/__init__.py +++ b/reditools/tools/analyze/rtchecks/__init__.py @@ -1,13 +1,20 @@ -from reditools.tools.analyze.rtchecks.check_column_edit_frequency import \ - CheckColumnEditFrequency -from reditools.tools.analyze.rtchecks.check_column_min_edits import \ - CheckColumnMinEdits +"""Manage and execute a suites of checks and filters on RNA editing results.""" + +from reditools.tools.analyze.rtchecks.check_column_edit_frequency import ( + CheckColumnEditFrequency, +) +from reditools.tools.analyze.rtchecks.check_column_min_edits import ( + CheckColumnMinEdits, +) from reditools.tools.analyze.rtchecks.check_exclusions import CheckExclusions -from reditools.tools.analyze.rtchecks.check_max_editing_nucleotides import \ - CheckMaxEditingNucleotides -from reditools.tools.analyze.rtchecks.check_min_read_depth import \ - CheckMinReadDepth -from reditools.tools.analyze.rtchecks.check_target_positions import \ - CheckTargetPositions +from reditools.tools.analyze.rtchecks.check_max_editing_nucleotides import ( + CheckMaxEditingNucleotides, +) +from reditools.tools.analyze.rtchecks.check_min_read_depth import ( + CheckMinReadDepth, +) +from reditools.tools.analyze.rtchecks.check_target_positions import ( + CheckTargetPositions, +) from reditools.tools.analyze.rtchecks.check_variants import CheckVariants from reditools.tools.analyze.rtchecks.rtchecks import RTChecks diff --git a/reditools/tools/analyze/rtchecks/check_column_edit_frequency.py b/reditools/tools/analyze/rtchecks/check_column_edit_frequency.py index a9bf783..5eafb31 100644 --- a/reditools/tools/analyze/rtchecks/check_column_edit_frequency.py +++ b/reditools/tools/analyze/rtchecks/check_column_edit_frequency.py @@ -1,6 +1,12 @@ -import argparse +"""Check if a position has a minimum number of total edits.""" +from __future__ import annotations -from reditools.compiled_position import RTResult +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckColumnEditFrequency: @@ -12,7 +18,7 @@ class CheckColumnEditFrequency: The minimum required total edits. """ - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckColumnEditFrequency. Parameters @@ -52,10 +58,10 @@ def run_check(self, rtresult: RTResult) -> None | tuple: None if total edits are sufficient, a tuple with error message otherwise. """ - edits_no = len(rtresult) - rtresult['REF'] + edits_no = len(rtresult) - rtresult["REF"] if edits_no < self.min_edits: return ( - 'DISCARDING COLUMN edits={} < {}', + "DISCARDING COLUMN edits={} < {}", edits_no, self.min_edits, ) diff --git a/reditools/tools/analyze/rtchecks/check_column_min_edits.py b/reditools/tools/analyze/rtchecks/check_column_min_edits.py index de6119b..de4b098 100644 --- a/reditools/tools/analyze/rtchecks/check_column_min_edits.py +++ b/reditools/tools/analyze/rtchecks/check_column_min_edits.py @@ -1,10 +1,19 @@ -import argparse +"""Check if a position has a minimum number of edits per nucleotide.""" +from __future__ import annotations -from reditools.compiled_position import RTResult +from typing import TYPE_CHECKING + +from reditools.constants import bases + +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckColumnMinEdits: """Check if a position has a minimum number of edits per nucleotide. + Specifically, checks that all non-zero, non-reference bases pass a given threshold. @@ -14,9 +23,7 @@ class CheckColumnMinEdits: The minimum required edits per nucleotide. """ - _bases = ('A', 'T', 'C', 'G') - - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckColumnMinEdits. Parameters @@ -56,11 +63,11 @@ def run_check(self, rtresult: RTResult) -> tuple | None: None if all nucleotide edits are sufficient, a tuple with error message otherwise. """ - for base in self._bases: + for base in bases: if base != rtresult.reference and \ 0 < rtresult[base] < self.min_edits_per_nucleotide: return ( - 'DISCARDING COLUMN edits={} < {}', + "DISCARDING COLUMN edits={} < {}", rtresult[base], self.min_edits_per_nucleotide, ) diff --git a/reditools/tools/analyze/rtchecks/check_exclusions.py b/reditools/tools/analyze/rtchecks/check_exclusions.py index 4f2e5f0..74fa1da 100644 --- a/reditools/tools/analyze/rtchecks/check_exclusions.py +++ b/reditools/tools/analyze/rtchecks/check_exclusions.py @@ -1,10 +1,16 @@ -import argparse +"""Check if a position is within excluded regions.""" +from __future__ import annotations + +from typing import TYPE_CHECKING from reditools import file_utils -from reditools.compiled_position import RTResult from reditools.region_collection import RegionCollection from reditools.splicing_file import load_splicing_file +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckExclusions: """Check if a position is within excluded regions. @@ -15,7 +21,7 @@ class CheckExclusions: The collection of excluded regions. """ - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckExclusions. Parameters @@ -68,5 +74,5 @@ def run_check(self, rtresult: RTResult) -> None | tuple: error message otherwise. """ if self.regions.contains(rtresult.contig, rtresult.position): - return ('DISCARD COLUMN in excluded region',) + return ("DISCARD COLUMN in excluded region",) return None diff --git a/reditools/tools/analyze/rtchecks/check_max_editing_nucleotides.py b/reditools/tools/analyze/rtchecks/check_max_editing_nucleotides.py index 19b6508..df7fbe2 100644 --- a/reditools/tools/analyze/rtchecks/check_max_editing_nucleotides.py +++ b/reditools/tools/analyze/rtchecks/check_max_editing_nucleotides.py @@ -1,7 +1,12 @@ -import argparse +"""Check if a position has at most a certain number of editing nucleotides.""" +from __future__ import annotations -from reditools.compiled_position import RTResult +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckMaxEditingNucleotides: """Check if a position has at most a certain number of editing nucleotides. @@ -12,7 +17,7 @@ class CheckMaxEditingNucleotides: The maximum allowed number of editing nucleotides. """ - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckMaxEditingNucleotides. Parameters @@ -36,7 +41,7 @@ def is_needed(cls, options: argparse.Namespace) -> bool: bool True if max_editing_nucleotides < 3, False otherwise. """ - return options.max_editing_nucleotides < 3 + return options.max_editing_nucleotides < 3 # noqa: PLR2004 def run_check(self, rtresult: RTResult) -> None | tuple: """Run the check on a specific position. @@ -55,7 +60,7 @@ def run_check(self, rtresult: RTResult) -> None | tuple: variants = rtresult.variants if len(variants) > self.max_editing_nucleotides: return ( - 'DISCARD COLUMN variants={} > {}', + "DISCARD COLUMN variants={} > {}", len(variants), self.max_editing_nucleotides, ) diff --git a/reditools/tools/analyze/rtchecks/check_min_read_depth.py b/reditools/tools/analyze/rtchecks/check_min_read_depth.py index 0f3cbbd..ab8f181 100644 --- a/reditools/tools/analyze/rtchecks/check_min_read_depth.py +++ b/reditools/tools/analyze/rtchecks/check_min_read_depth.py @@ -1,7 +1,12 @@ -import argparse +"""Check if a position has minimum read depth.""" +from __future__ import annotations -from reditools.compiled_position import RTResult +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckMinReadDepth: """Check if a position has minimum read depth. @@ -12,7 +17,7 @@ class CheckMinReadDepth: The minimum required read depth. """ - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckMinReadDepth. Parameters @@ -54,7 +59,7 @@ def run_check(self, rtresult: RTResult) -> None | tuple: """ if len(rtresult) < self.min_read_depth: return ( - 'DISCARDING COLUMN {} [MIN_READ_DEPTH={}]', + "DISCARDING COLUMN {} [MIN_READ_DEPTH={}]", len(rtresult), self.min_read_depth, ) diff --git a/reditools/tools/analyze/rtchecks/check_target_positions.py b/reditools/tools/analyze/rtchecks/check_target_positions.py index 2c5cb9f..3e91795 100644 --- a/reditools/tools/analyze/rtchecks/check_target_positions.py +++ b/reditools/tools/analyze/rtchecks/check_target_positions.py @@ -1,9 +1,15 @@ -import argparse +"""Check if a position is within target regions.""" +from __future__ import annotations + +from typing import TYPE_CHECKING from reditools import file_utils -from reditools.compiled_position import RTResult from reditools.region_collection import RegionCollection +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult class CheckTargetPositions: """Check if a position is within target regions. @@ -14,7 +20,7 @@ class CheckTargetPositions: The collection of target regions. """ - def __init__(self, options: argparse.Namespace): + def __init__(self, options: argparse.Namespace) -> None: """Initialize CheckTargetPositions. Parameters @@ -56,5 +62,5 @@ def run_check(self, rtresult: RTResult) -> None | tuple: error message otherwise. """ if not self.regions.contains(rtresult.contig, rtresult.position): - return ('DISCARD COLUMN not in target regions',) + return ("DISCARD COLUMN not in target regions",) return None diff --git a/reditools/tools/analyze/rtchecks/check_variants.py b/reditools/tools/analyze/rtchecks/check_variants.py index 6a1d2c3..1eaf3fa 100644 --- a/reditools/tools/analyze/rtchecks/check_variants.py +++ b/reditools/tools/analyze/rtchecks/check_variants.py @@ -1,12 +1,31 @@ -import argparse +"""Check if detected variants match specified allowed variants.""" +from __future__ import annotations + import re +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import argparse + + from reditools.compiled_position import RTResult -from reditools.compiled_position import RTResult +class BadVariantError(ValueError): + """Variant string is improperly formatted.""" + + def __init__(self, bad_alt: str) -> None: + """Initialize self. + + Parameters + ---------- + bad_alt : str + The offending variant string. + """ + self.message = f"Bad variant ({bad_alt}). Must be two bases (e.g. AG)." + super().__init__(self.message) class CheckVariants: - """ - Check if detected variants match specified allowed variants. + """Check if detected variants match specified allowed variants. Parameters ---------- @@ -14,9 +33,8 @@ class CheckVariants: Command-line options containing allowed variants. """ - def __init__(self, options: argparse.Namespace): - """ - Initialize CheckVariants with allowed variants. + def __init__(self, options: argparse.Namespace) -> None: + """Initialize CheckVariants with allowed variants. Parameters ---------- @@ -25,24 +43,21 @@ def __init__(self, options: argparse.Namespace): Raises ------ - ValueError + BadVariantError If a variant is not exactly two bases (e.g., 'AG'). """ - pa = re.compile('[ATCG]{2}', re.IGNORECASE) + pa = re.compile("[ATCG]{2}", re.IGNORECASE) bad_alt = next( (_ for _ in options.variants if not pa.fullmatch(_)), None, ) if bad_alt is not None: - raise ValueError( - f'Bad variant ({bad_alt}). Must be two bases (e.g. AG).' - ) + raise BadVariantError(bad_alt) self.variants = {_.upper() for _ in options.variants} @classmethod def is_needed(cls, options: argparse.Namespace) -> bool: - """ - Determine if the variant check is required. + """Determine if the variant check is required. Parameters ---------- @@ -55,11 +70,10 @@ def is_needed(cls, options: argparse.Namespace) -> bool: True if specific variants are required, False if 'ALL' is in the variant list. """ - return 'ALL' not in [_.upper() for _ in options.variants] + return "ALL" not in [_.upper() for _ in options.variants] def run_check(self, rtresult: RTResult) -> None | tuple: - """ - Verify that detected variants are among the allowed ones. + """Verify that detected variants are among the allowed ones. Parameters ---------- @@ -75,7 +89,7 @@ def run_check(self, rtresult: RTResult) -> None | tuple: if any(_ in self.variants for _ in rtresult.variants): return None return ( - 'DISCARD COLUMN Edits {} not in requested alts {}', + "DISCARD COLUMN Edits {} not in requested alts {}", rtresult.variants, self.variants, ) diff --git a/reditools/tools/analyze/rtchecks/rtchecks.py b/reditools/tools/analyze/rtchecks/rtchecks.py index 19ae123..d9be9b3 100644 --- a/reditools/tools/analyze/rtchecks/rtchecks.py +++ b/reditools/tools/analyze/rtchecks/rtchecks.py @@ -1,12 +1,27 @@ -import argparse +"""Manage and execute a suites of checks and filters on RNA editing results.""" +from __future__ import annotations + +from typing import TYPE_CHECKING -from reditools.compiled_position import RTResult from reditools.tools.analyze import rtchecks +if TYPE_CHECKING: + import argparse -class RTChecks(object): - """ - Manage and execute a suite of checks on RNA editing results. + from reditools.compiled_position import RTResult + +all_checks = ( + rtchecks.CheckColumnEditFrequency, + rtchecks.CheckColumnMinEdits, + rtchecks.CheckMinReadDepth, + rtchecks.CheckExclusions, + rtchecks.CheckMaxEditingNucleotides, + rtchecks.CheckTargetPositions, + rtchecks.CheckVariants, +) + +class RTChecks: + """Manage and execute a suite of checks on RNA editing results. Parameters ---------- @@ -14,32 +29,20 @@ class RTChecks(object): Command-line options that determine which checks are enabled. """ - def __init__(self, options: argparse.Namespace): - """ - Initialize RTChecks with enabled check instances. + def __init__(self, options: argparse.Namespace) -> None: + """Initialize RTChecks with enabled check instances. Parameters ---------- options : argparse.Namespace Command-line options used to filter and configure checks. """ - self.check_list = [] - - for check in ( - rtchecks.CheckColumnEditFrequency, - rtchecks.CheckColumnMinEdits, - rtchecks.CheckMinReadDepth, - rtchecks.CheckExclusions, - rtchecks.CheckMaxEditingNucleotides, - rtchecks.CheckTargetPositions, - rtchecks.CheckVariants, - ): - if check.is_needed(options): - self.check_list.append(check(options)) + self.check_list = [ + check(options) for check in all_checks if check.is_needed(options) + ] def check(self, rtresult: RTResult) -> None | tuple: - """ - Run all enabled checks against a set of base results. + """Run all enabled checks against a set of base results. Parameters ---------- diff --git a/reditools/tools/analyze/setup_alignment_manager.py b/reditools/tools/analyze/setup_alignment_manager.py index aa03afe..94d7144 100644 --- a/reditools/tools/analyze/setup_alignment_manager.py +++ b/reditools/tools/analyze/setup_alignment_manager.py @@ -1,3 +1,6 @@ +"""Initalized and configure ALignmentManager objects for the analyze tool.""" +from __future__ import annotations + from reditools import file_utils from reditools.alignment_manager import AlignmentManager diff --git a/reditools/tools/analyze/setup_rtools.py b/reditools/tools/analyze/setup_rtools.py index e18f76b..932f1fc 100644 --- a/reditools/tools/analyze/setup_rtools.py +++ b/reditools/tools/analyze/setup_rtools.py @@ -1,3 +1,4 @@ +"""Create REDItools objects for parallel processing with the analyze tool.""" import argparse from reditools import reditools diff --git a/reditools/tools/analyze/temp_file_manager.py b/reditools/tools/analyze/temp_file_manager.py index 4d03ec5..6811fb5 100644 --- a/reditools/tools/analyze/temp_file_manager.py +++ b/reditools/tools/analyze/temp_file_manager.py @@ -1,23 +1,27 @@ +"""Manages temporary output files for the analyze tool.""" from __future__ import annotations import csv -import os import sys import tempfile -from types import TracebackType -from typing import Iterator +from pathlib import Path +from typing import TYPE_CHECKING from reditools.region import Region from reditools.tools.analyze.concat_output import concat_output from reditools.tools.analyze.parse_args import json_args -save_file = 'region_file_list.csv' +if TYPE_CHECKING: + from types import TracebackType + from typing import Iterator + +save_file = "region_file_list.csv" class TempFileManager: """Manages the temporary output files for REDItools.""" def __init__(self, dirpath: str, regions: list[Region] | None=None) -> None: - """Create a new TempFileManager + """Create a new TempFileManager. Parameters ---------- @@ -28,40 +32,71 @@ def __init__(self, dirpath: str, regions: list[Region] | None=None) -> None: """ self.dirpath = dirpath if regions: - temp_files = [ - tempfile.NamedTemporaryFile(dir=self.dirpath, delete=False).name - for _ in regions - ] + temp_files = [] + for _ in regions: + with tempfile.NamedTemporaryFile( + dir=self.dirpath, + delete=False, + ) as tf: + temp_files.append(tf.name) self.region_file_list = list(zip(regions, temp_files)) - with open(os.path.join( - self.dirpath, - save_file, - ), 'w') as stream: - writer = csv.writer(stream) - writer.writerow(['Region', 'Filename']) - for region, filename in self.region_file_list: - writer.writerow([region, os.path.basename(filename)]) + self.save_to_file() else: - with open(os.path.join(self.dirpath, save_file), 'r') as stream: + with Path(self.dirpath, save_file).open("r") as stream: reader = csv.DictReader(stream) self.region_file_list = [ ( - Region.from_string(row['Region']), - os.path.join(self.dirpath, row['Filename']), + Region.from_string(row["Region"]), + str(Path(self.dirpath, row["Filename"])), ) for row in reader ] def __enter__(self) -> TempFileManager: + """Open TempFileManager.""" return self - def __iter__(self) -> Iterator: + def __exit__( + self, + typ: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + """If no error occurred, remove all temporary files. + + Deletes the *.done files, cli JSON, and region CSV files. + """ + if typ is not None: + return + + for _, filename in self.region_file_list: + Path(f"{filename}.done").unlink() + + for temp_file in (json_args.json_args_filename, save_file): + Path(self.dirpath, temp_file).unlink() + try: + Path(self.dirpath).rmdir() + except OSError as os_exc: + sys.stderr.write( + "[WARNING] Could not delete temporary files directory " + f"{self.dirpath}. {os_exc}\n", + ) + + def __iter__(self) -> Iterator[tuple[Region, str]]: + """Iterate over region-filename pairs. + + Yields + ------ + tuple[Region, str] + Genomic region and path to save results to. + """ yield from self.region_file_list def __len__(self) -> int: + """List the number of temporary output files.""" return len(self.region_file_list) - def concat(self, filepath: str, mode: str='w') -> None: + def concat(self, filepath: str, mode: str="w") -> None: """Concat all temporary files, then delete them. Parameters @@ -77,32 +112,13 @@ def concat(self, filepath: str, mode: str='w') -> None: mode, ) - def cleanup(self) -> None: - """Delete the *.done files, cli JSON, and region CSV files. - - Raises - ------ - OSError - If the temporary directory cannot be emptied. - """ - for _, filename in self.region_file_list: - os.remove(f'{filename}.done') - - for temp_file in (json_args.json_args_filename, save_file): - os.remove(os.path.join(self.dirpath, temp_file)) - try: - os.rmdir(self.dirpath) - except OSError as exc: - sys.stderr.write( - '[WARNING] Could not delete temporary files directory ' - f'{self.dirpath}. {exc}\n' - ) - - def __exit__( - self, - exc_type: type, - exc_value: Exception, - traceback: TracebackType, - ) -> None: - if exc_type is None: - self.cleanup() + def save_to_file(self) -> None: + """Save list of region files to CSV.""" + with Path(self.dirpath, save_file).open("w") as stream: + writer = csv.writer(stream) + writer.writerow(["Region", "Filename"]) + for region, filename in self.region_file_list: + writer.writerow([ + region, + Path(filename).name, + ]) diff --git a/reditools/tools/analyze/write_results.py b/reditools/tools/analyze/write_results.py index 1096536..77772d8 100644 --- a/reditools/tools/analyze/write_results.py +++ b/reditools/tools/analyze/write_results.py @@ -1,11 +1,14 @@ +"""Write analysis results.""" + import csv +from pathlib import Path from typing import Callable, Iterator from reditools.compiled_position import RTResult from reditools.logger import Logger from reditools.tools.analyze.rtchecks import RTChecks -_empty = '-' +_empty = "-" def write_results( rtresults: Iterator[RTResult], @@ -13,7 +16,7 @@ def write_results( filters: RTChecks, logger: Callable, ) -> None: - """Write analysis results to a temporary file. + """Write analysis results to a file. Parameters ---------- @@ -26,8 +29,8 @@ def write_results( logger : Callable The logger function for debug messages. """ - with open(filename, 'w') as stream: - writer = csv.writer(stream, delimiter='\t', lineterminator='\n') + with Path(filename).open("w") as stream: + writer = csv.writer(stream, delimiter="\t", lineterminator="\n") for rt_result in rtresults: msg = filters.check(rt_result) if msg: @@ -40,9 +43,9 @@ def write_results( rt_result.reference, rt_result.strand, len(rt_result), - f'{rt_result.mean_quality:.2f}', + f"{rt_result.mean_quality:.2f}", list(rt_result), - ' '.join(sorted(variants)) if variants else _empty, - f'{rt_result.edit_ratio:.2f}', + " ".join(sorted(variants)) if variants else _empty, + f"{rt_result.edit_ratio:.2f}", _empty, _empty, _empty, _empty, _empty, ]) diff --git a/reditools/tools/annotate/__init__.py b/reditools/tools/annotate/__init__.py deleted file mode 100644 index 776e742..0000000 --- a/reditools/tools/annotate/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from reditools.tools.annotate.main import main diff --git a/reditools/tools/annotate/main.py b/reditools/tools/annotate/main.py index 9185593..c0d35c3 100644 --- a/reditools/tools/annotate/main.py +++ b/reditools/tools/annotate/main.py @@ -1,3 +1,7 @@ +"""Entry point for annotate tool.""" + +from __future__ import annotations + import csv import sys import traceback @@ -8,7 +12,21 @@ from reditools.rtannotater import RTAnnotater from reditools.tools.annotate.parse_args import parse_args -_contig = 'Region' +_contig = "Region" + +class UnsortedInputError(ValueError): + """REDItools output file is unsorted.""" + + def __init__(self, file_name: str) -> None: + """Initialize self. + + Parameters + ---------- + file_name : str + The name of the unsorted file. + """ + self.message = f"File {file_name} does not appear to be in sorted order" + super().__init__(self.message) def contig_order_from_bam(bam_fname: str) -> dict[str, int]: """Get contig order from a BAM file. @@ -23,11 +41,11 @@ def contig_order_from_bam(bam_fname: str) -> dict[str, int]: dict[str, int] A dictionary mapping contig names to their 1-based order. """ - contigs = {} with pysam.AlignmentFile(bam_fname, ignore_truncation=True) as bam: - for idx, contig in enumerate(bam.references, start=1): - contigs[contig] = idx - return contigs + return { + contig: idx + for idx, contig in enumerate(bam.references, start=1) + } def contig_order_from_fai(fai_fname: str) -> dict[str, int]: """Get contig order from a FASTA index file. @@ -43,9 +61,9 @@ def contig_order_from_fai(fai_fname: str) -> dict[str, int]: A dictionary mapping contig names to their 1-based order. """ contigs = {} - with file_utils.open_stream(fai_fname, 'r') as stream: + with file_utils.open_stream(fai_fname, "r") as stream: for idx, line in enumerate(stream, start=1): - contig = line.split('\t')[0] + contig = line.split("\t")[0] contigs[contig] = idx return contigs @@ -64,20 +82,17 @@ def contig_order_from_out(out_fname: str) -> dict[str, int]: Raises ------ - ValueError + UnsortedInputError If the file does not appear to be in sorted order. """ contigs: dict[str, int] = {} - with file_utils.open_stream(out_fname, 'r') as stream: - reader = csv.DictReader(stream, delimiter='\t') + with file_utils.open_stream(out_fname, "r") as stream: + reader = csv.DictReader(stream, delimiter="\t") last_contig = None for row in reader: if row[_contig] != last_contig: if row[_contig] in contigs: - raise ValueError( - f'File {out_fname} does not appear to be in sorted ' - 'order.' - ) + raise UnsortedInputError(out_fname) contigs[row[_contig]] = len(contigs) + 1 last_contig = row[_contig] return contigs @@ -98,20 +113,20 @@ def main() -> None: else: order_fname = options.rna_file contig_order = contig_order_from_out(options.rna_file) - except Exception as exc: + except Exception as exc: # noqa: BLE001 if options.debug: traceback.print_exception(*sys.exc_info()) sys.stderr.write( - '[ERROR] There was an error getting the contig order from ' - f'{order_fname}: ({type(exc)}) {exc}\n' + "[ERROR] There was an error getting the contig order from " + f"{order_fname}: ({type(exc)}) {exc}\n", ) sys.exit(1) rta = RTAnnotater(contig_order, options.strand_correction) try: rta.annotate(options.rna_file, options.dna_file, sys.stdout) - except Exception as exc: + except Exception as exc: # noqa: BLE001 if options.debug: traceback.print_exception(*sys.exc_info()) - sys.stderr.write(f'[ERROR] ({type(exc)}) {exc}\n') + sys.stderr.write(f"[ERROR] ({type(exc)}) {exc}\n") sys.exit(1) diff --git a/reditools/tools/annotate/parse_args.py b/reditools/tools/annotate/parse_args.py index 747cfed..fc53d02 100644 --- a/reditools/tools/annotate/parse_args.py +++ b/reditools/tools/annotate/parse_args.py @@ -1,3 +1,5 @@ +"""Parse CLI options for the annotate tool.""" + import argparse @@ -10,58 +12,58 @@ def parse_args() -> argparse.Namespace: The parsed command-line arguments. """ parser = argparse.ArgumentParser( - prog='reditools annotate', - description='Annotates RNA REDItools output with DNA output.', + prog="reditools annotate", + description="Annotates RNA REDItools output with DNA output.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( - 'rna_file', - help='The REDItools output from RNA data', + "rna_file", + help="The REDItools output from RNA data", ) parser.add_argument( - 'dna_file', - help='The REDItools output from corresponding DNA data', + "dna_file", + help="The REDItools output from corresponding DNA data", ) parser.add_argument( - '-d', - '--debug', - help='Report stack trace on crash.', - action='store_true', + "-d", + "--debug", + help="Report stack trace on crash.", + action="store_true", ) parser.add_argument( - '-C', - '--strand-correction', + "-C", + "--strand-correction", help=( - 'Report the DNA base complement if the RNA data comes from the ' - 'minus strand.' + "Report the DNA base complement if the RNA data comes from the " + "minus strand." ), - action='store_true' + action="store_true", ) order_group = parser.add_argument_group( - title='Contig order options', + title="Contig order options", description=( - 'By default, the annotate tool will determine the contig order ' - 'by reading through the rna_file once before performing the ' - 'annotation. Depending on the size of the rna_file, this could ' - 'take a while. To skip this time consuming step, you can provide ' - 'the contig order from another source.' + "By default, the annotate tool will determine the contig order " + "by reading through the rna_file once before performing the " + "annotation. Depending on the size of the rna_file, this could " + "take a while. To skip this time consuming step, you can provide " + "the contig order from another source." ), ) order_group.add_argument( - '-b', - '--bam', - help='BAM file to get contig order from.', + "-b", + "--bam", + help="BAM file to get contig order from.", ) order_group.add_argument( - '-f', - '--fai', - help='FASTA Index file to get contig order from.', + "-f", + "--fai", + help="FASTA Index file to get contig order from.", ) options = parser.parse_args() if options.bam and options.fai: parser.error( - message='Options -b/--bam and -f/--fai are mutually exclusive.', + message="Options -b/--bam and -f/--fai are mutually exclusive.", ) return options diff --git a/reditools/tools/find_repeats/__init__.py b/reditools/tools/find_repeats/__init__.py deleted file mode 100644 index 4606718..0000000 --- a/reditools/tools/find_repeats/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from reditools.tools.find_repeats.main import find_homo_seqs, main diff --git a/reditools/tools/find_repeats/main.py b/reditools/tools/find_repeats/main.py index 2f30876..433eecd 100644 --- a/reditools/tools/find_repeats/main.py +++ b/reditools/tools/find_repeats/main.py @@ -1,3 +1,6 @@ +"""Entry point for find-repeats tools.""" +from __future__ import annotations + import argparse import csv import sys @@ -23,7 +26,7 @@ def find_homo_seqs(seq: str, length: int = 5) -> Iterator[tuple[int, int, str]]: Iterator[tuple[int, int, str]] A tuple containing (start, end, base) for each homopolymeric sequence. """ - h_base = '' + h_base = "" start = 0 count = 0 @@ -50,26 +53,28 @@ def parse_options() -> argparse.Namespace: """ parser = argparse.ArgumentParser( prog="reditools find-repeats", - description='REDItools3', + description="REDItools3", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( - 'file', - help='The fasta file to be analyzed', + "file", + help="The fasta file to be analyzed", ) parser.add_argument( - '-l', - '--min-length', + "-l", + "--min-length", type=int, default=5, - help='Minimum length of repeat region', + help="Minimum length of repeat region", ) parser.add_argument( - '-o', - '--output', - default='/dev/stdout', - help='Destination to write results. Default is to use STDOUT. ' + - 'If the filename ends in .gz, the contents will be gzipped.', + "-o", + "--output", + default="/dev/stdout", + help=( + "Destination to write results. Default is to use STDOUT. " + "If the filename ends in .gz, the contents will be gzipped." + ), ) return parser.parse_args() @@ -113,11 +118,11 @@ def main() -> None: if options.output: stream = file_utils.open_stream( options.output, - 'wt', - encoding='utf-8', + "wt", + encoding="utf-8", ) else: stream = sys.stdout - writer = csv.writer(stream, delimiter='\t') + writer = csv.writer(stream, delimiter="\t") writer.writerows(iter_homo_output(fasta, options.min_length)) diff --git a/reditools/tools/index/__init__.py b/reditools/tools/index/__init__.py deleted file mode 100644 index b0ff873..0000000 --- a/reditools/tools/index/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from reditools.tools.index.main import main diff --git a/reditools/tools/index/main.py b/reditools/tools/index/main.py index a6cc314..eda5c9c 100644 --- a/reditools/tools/index/main.py +++ b/reditools/tools/index/main.py @@ -1,3 +1,4 @@ +"""Entry point for the index tool.""" import sys @@ -28,7 +29,7 @@ def main() -> None: indexer.add_target_from_bed(_) if options.output_file: - stream = open_stream(options.output_file, 'w') + stream = open_stream(options.output_file, "w") else: stream = sys.stdout @@ -36,8 +37,8 @@ def main() -> None: indexer.add_rt_output(_) for nuc, idx in sorted(indexer.calc_index().items()): - stream.write(f'{nuc}\t{idx}\n') + stream.write(f"{nuc}\t{idx}\n") -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/reditools/tools/index/parse_args.py b/reditools/tools/index/parse_args.py index e0f7a09..886873f 100644 --- a/reditools/tools/index/parse_args.py +++ b/reditools/tools/index/parse_args.py @@ -1,3 +1,5 @@ +"""Functions for parsing CLI options for the analyze tool.""" + import argparse @@ -11,36 +13,36 @@ def parse_args() -> argparse.Namespace: """ parser = argparse.ArgumentParser( prog="reditools index", - description='REDItools3', + description="REDItools3", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( - 'file', - nargs='+', - help='The REDItools output file to be analyzed', + "file", + nargs="+", + help="The REDItools output file to be analyzed", ) parser.add_argument( - '-o', - '--output-file', - default='/dev/stdout', - help='The output statistics file', + "-o", + "--output-file", + default="/dev/stdout", + help="The output statistics file", ) parser.add_argument( - '-g', - '--region', - help='The genomic region to be analyzed', + "-g", + "--region", + help="The genomic region to be analyzed", ) parser.add_argument( - '-B', - '--bed_file', - nargs='+', - help='Path of BED file containing target regions', + "-B", + "--bed_file", + nargs="+", + help="Path of BED file containing target regions", ) parser.add_argument( - '-k', - '--exclude_regions', - nargs='+', - help='Path of BED file containing regions to exclude from analysis', + "-k", + "--exclude_regions", + nargs="+", + help="Path of BED file containing regions to exclude from analysis", ) return parser.parse_args() diff --git a/ruff.toml b/ruff.toml deleted file mode 100644 index b40097d..0000000 --- a/ruff.toml +++ /dev/null @@ -1,5 +0,0 @@ -line-length = 80 - -[per-file-ignores] -"test/__main__.py" = ["F401"] -"*/__init__.py" = ["F401", "E501"] diff --git a/test/__init__.py b/test/__init__.py index e69de29..0c59cd5 100644 --- a/test/__init__.py +++ b/test/__init__.py @@ -0,0 +1 @@ +"""Perform unittest quality controls.""" diff --git a/test/__main__.py b/test/__main__.py index 95014a4..cffc4d7 100644 --- a/test/__main__.py +++ b/test/__main__.py @@ -1,3 +1,4 @@ +"""Perform unittest quality controls.""" import unittest from test.alignment_file import TestRTAlignmentFile from test.alignment_manager import TestAlignmentManager diff --git a/test/aligner.py b/test/aligner.py index 3ad036b..5a1c403 100644 --- a/test/aligner.py +++ b/test/aligner.py @@ -1,10 +1,45 @@ +"""Needleman-Wunsch sequence aligner.""" +from __future__ import annotations + + class Aligner: - def __init__(self, match=1, mismatch=1, gap=1): - self.match = 1 - self.mismatch = 1 - self.gap = 1 + """Needleman-Wunsch sequence aligner.""" + + _trace_match = frozenset((2, 5, 6, 9)) + _trace_del = frozenset((3, 7)) + _trace_ins = 4 + + def __init__(self, match: int=1, mismatch: int=1, gap: int=1) -> None: + """Initialize self. + + Parameters + ---------- + match : int + Match score. + mismatch : int + Mismatch penalty. + gap : int + Gap penalty. + """ + self.match = match + self.mismatch = mismatch + self.gap = gap + + def align(self, ref_seq: str, query_seq: str) -> tuple[str, str]: + """Perform alignment. + + Parameters + ---------- + ref_seq : str + Reference sequence. + query_seq : str + Query sequence. - def align(self, ref_seq, query_seq): + Returns + ------- + tuple[str, str] + Reference and query alignment strings, respectively. + """ matrix = NWMatrix( ref_seq, query_seq, @@ -19,36 +54,81 @@ def align(self, ref_seq, query_seq): matrix.trace_matrix, ) - def trace_matrix(self, ref_seq, qry_seq, trace_mat): + def trace_matrix( + self, + ref_seq: str, + qry_seq: str, + trace_mat: list[list[int]], + ) -> tuple[str, str]: + """Trace path from matrix to create alignment strings. + + Parameters + ---------- + ref_seq : str + Reference sequence. + qry_seq : str + Query sequence. + trace_mat : list[list[int]] + Trace matrix. + + Returns + ------- + tuple[str, str] + Reference and query alignment strings, respectively. + """ row_idx = len(ref_seq) col_idx = len(qry_seq) - ref_align = [] - qry_align = [] + ref_align: list[str] = [] + qry_align: list[str] = [] while row_idx > 0 or col_idx > 0: trace_val = trace_mat[row_idx][col_idx] - if trace_val in [2, 5, 6, 9]: + if trace_val in self._trace_match: row_idx -= 1 col_idx -= 1 ref_align.insert(0, ref_seq[row_idx]) qry_align.insert(0, qry_seq[col_idx]) - elif trace_val in [3, 7]: + elif trace_val in self._trace_del: row_idx -= 1 ref_align.insert(0, ref_seq[row_idx]) - qry_align.insert(0, '-') - elif trace_val == 4: + qry_align.insert(0, "-") + elif trace_val == self._trace_ins: col_idx -= 1 - ref_align.insert(0, '-') + ref_align.insert(0, "-") qry_align.insert(0, qry_seq[col_idx]) return ( - ''.join(ref_align), - ''.join(qry_align), + "".join(ref_align), + "".join(qry_align), ) class NWMatrix: - def __init__(self, ref_seq, query_seq, match, mismatch, gap): + """Needleman-Wunsch matrix and trace matrix object.""" + + def __init__( + self, + ref_seq: str, + query_seq: str, + match: int, + mismatch: int, + gap: int, + ) -> None: + """Initialize self. + + Parameters + ---------- + ref_seq : str + Reference sequence. + query_seq : str + Query sequence. + match : int + Match score. + mismatch : int + Mismatch penalty. + gap : int + Gap penalty. + """ self.ref_seq = ref_seq self.qry_seq = query_seq self.gap = gap @@ -58,25 +138,59 @@ def __init__(self, ref_seq, query_seq, match, mismatch, gap): self.nw_matrix = self.init_nw_matrix(ref_seq, query_seq) self.trace_matrix = self.init_trace_matrix(ref_seq, query_seq) - def assess_cell(self, col_idx, ref_base, row_idx, query_base): + def assess_cell( + self, + col_idx: int, + ref_base: str, + row_idx: int, + query_base: str, + ) -> None: + """Fill in a cell of the matrices. + + Parameters + ---------- + col_idx : int + Column index - 1. + ref_base : str + Reference nucleotide. + row_idx : int + Row index - 1. + query_base : str + Query nucleotide. + """ align_val = self.match if ref_base == query_base else -self.mismatch t_list = [ self.nw_matrix[row_idx][col_idx] + align_val, self.nw_matrix[row_idx][col_idx + 1] - self.gap, self.nw_matrix[row_idx + 1][col_idx] - self.gap, ] - t_max = max(t_list) + t_max = max(t_list) self.nw_matrix[row_idx + 1][col_idx + 1] = t_max self.trace_matrix[row_idx + 1][col_idx + 1] += sum(( idx + 2 for idx, tv in enumerate(t_list) if tv == t_max )) - def run_dp(self): + def run_dp(self) -> None: + """Run dynamic programming.""" for col_idx, ref_base in enumerate(self.qry_seq): for row_idx, query_base in enumerate(self.ref_seq): self.assess_cell(col_idx, ref_base, row_idx, query_base) - def init_nw_matrix(self, ref_seq, query_seq): + def init_nw_matrix(self, ref_seq: str, query_seq: str) -> list[list[int]]: + """Create a starting Needleman-Wunsch matrix. + + Parameters + ---------- + ref_seq : str + Reference sequence. + query_seq : str + Query sequence. + + Returns + ------- + list[list[int]] + Initialized matrix. + """ nw_mat = [ [-self.gap * (row_idx + 1)] + \ [0 for _ in range(len(query_seq))] @@ -90,7 +204,25 @@ def init_nw_matrix(self, ref_seq, query_seq): ) return nw_mat - def init_trace_matrix(self, ref_seq, query_seq): + def init_trace_matrix( + self, + ref_seq: str, + query_seq: str, + ) -> list[list[int]]: + """Create an initialized trace matrix. + + Parameters + ---------- + ref_seq : str + Reference sequence. + query_seq : str + Query sequence. + + Returns + ------- + list[list[int]] + Initialized matrix. + """ trace_mat = [ [3] + [0 for _ in range(len(query_seq))] for _ in range(len(ref_seq)) diff --git a/test/alignment_file.py b/test/alignment_file.py index 38a8772..15cc2ca 100644 --- a/test/alignment_file.py +++ b/test/alignment_file.py @@ -1,24 +1,32 @@ -import os +"""Test cases for RTAlignmentFile.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, Sequence, ntf from reditools.alignment_file import RTAlignmentFile class TestRTAlignmentFile(unittest.TestCase): - def setUp(self): + """Test cases for RTAlignmentFile.""" + + def setUp(self) -> None: + """Preflight setup.""" self.sam_obj = SAM() - self.sam_obj.add_contig('chr1', length=60) - self.refseq = self.sam_obj.genome['chr1'] + self.sam_obj.add_contig("chr1", length=60) + self.refseq = self.sam_obj.genome["chr1"] - self.genome_fname = ntf(suffix='.fa') - self.bam_fname = ntf(suffix='.bam') + self.genome_fname = ntf(suffix=".fa") + self.bam_fname = ntf(suffix=".bam") - def tearDown(self): - os.remove(self.genome_fname) - os.remove(self.bam_fname) + def tearDown(self) -> None: + """Posttest teardown.""" + Path(self.genome_fname).unlink() + Path(self.bam_fname).unlink() - def test_fetch_by_position(self): + def test_fetch_by_position(self) -> None: + """Check fetch_by_position standard functionality.""" for start, stop in ( (0, 20), (20, None), @@ -26,28 +34,29 @@ def test_fetch_by_position(self): (40, None), ): read_seq = self.refseq[start:stop] - self.sam_obj.add_read('chr1', Sequence(read_seq, start)) + self.sam_obj.add_read("chr1", Sequence(read_seq, start)) - self.sam_obj.add_contig('chr2') - self.sam_obj.add_read('chr2', Sequence(self.sam_obj.genome['chr2'], 0)) + self.sam_obj.add_contig("chr2") + self.sam_obj.add_read("chr2", Sequence(self.sam_obj.genome["chr2"], 0)) self.sam_obj.genome.save_to_fasta(self.genome_fname) self.sam_obj.save_to_sam(self.bam_fname, self.genome_fname) with RTAlignmentFile(self.bam_fname) as rtaf: - reads_iter = rtaf.fetch_by_position('chr1') + reads_iter = rtaf.fetch_by_position("chr1") self.assertEqual(len(next(reads_iter)), 1) self.assertEqual(len(next(reads_iter)), 2) self.assertEqual(len(next(reads_iter)), 1) - def test_exclude_reads(self): + def test_exclude_reads(self) -> None: + """Check ability to exclude reads by name.""" self.sam_obj.add_read( - 'chr1', - Sequence(self.refseq, 0, qname='exclude_me'), + "chr1", + Sequence(self.refseq, 0, read_name="exclude_me"), ) self.sam_obj.add_read( - 'chr1', - Sequence(self.refseq, 0, qname='include_me'), + "chr1", + Sequence(self.refseq, 0, read_name="include_me"), ) self.sam_obj.genome.save_to_fasta(self.genome_fname) @@ -55,94 +64,102 @@ def test_exclude_reads(self): with RTAlignmentFile( self.bam_fname, - excluded_read_names={'exclude_me'}, + excluded_read_names={"exclude_me"}, ) as rtaf: - reads = next(rtaf.fetch_by_position('chr1')) + reads = next(rtaf.fetch_by_position("chr1")) self.assertEqual(len(reads), 1) - self.assertEqual(reads[0].qname, 'include_me') + self.assertEqual(reads[0].query_name, "include_me") - def test_check_quality(self): + def test_check_quality(self) -> None: + """Check MAPQ quality filter.""" self.sam_obj.add_read( - 'chr1', + "chr1", Sequence(self.refseq, 0, mapq=10), ) self.sam_obj.add_read( - 'chr1', - Sequence(self.refseq, 0, mapq=30, qname='include_me'), + "chr1", + Sequence(self.refseq, 0, mapq=30, read_name="include_me"), ) self.sam_obj.genome.save_to_fasta(self.genome_fname) self.sam_obj.save_to_sam(self.bam_fname, self.genome_fname) with RTAlignmentFile(self.bam_fname, min_quality=20) as rtaf: - reads = next(rtaf.fetch_by_position('chr1')) + reads = next(rtaf.fetch_by_position("chr1")) self.assertEqual(len(reads), 1) - self.assertEqual(reads[0].qname, 'include_me') + self.assertEqual(reads[0].query_name, "include_me") - def test_check_length(self): + def test_check_length(self) -> None: + """Check minimum read length filter.""" self.sam_obj.add_read( - 'chr1', + "chr1", Sequence(self.refseq[:20], 0), ) self.sam_obj.add_read( - 'chr1', - Sequence(self.refseq, 0, qname='include_me'), + "chr1", + Sequence(self.refseq, 0, read_name="include_me"), ) self.sam_obj.genome.save_to_fasta(self.genome_fname) self.sam_obj.save_to_sam(self.bam_fname, self.genome_fname) with RTAlignmentFile(self.bam_fname, min_length=30) as rtaf: - reads = next(rtaf.fetch_by_position('chr1')) + reads = next(rtaf.fetch_by_position("chr1")) self.assertEqual(len(reads), 1) - self.assertEqual(reads[0].qname, 'include_me') + self.assertEqual(reads[0].query_name, "include_me") - def test_check_se_flags(self): + def test_check_se_flags(self) -> None: + """Check for filtering by SAM flags.""" for idx, flag in enumerate([0, 16]): read = Sequence( self.refseq, 0, flag=flag, - qname=f'se_good_{idx}', + read_name=f"se_good_{idx}", ) - self.sam_obj.add_read('chr1', read) + self.sam_obj.add_read("chr1", read) for idx, flag in enumerate([4, 256, 272, 512, 1024, 2048, 2064]): read = Sequence( self.refseq, 0, flag=flag, - qname=f'se_bad_{idx}', + read_name=f"se_bad_{idx}", ) - self.sam_obj.add_read('chr1', read) + self.sam_obj.add_read("chr1", read) self.sam_obj.genome.save_to_fasta(self.genome_fname) self.sam_obj.save_to_sam(self.bam_fname, self.genome_fname) with RTAlignmentFile(self.bam_fname) as rtaf: - reads = next(rtaf.fetch_by_position('chr1')) + reads = next(rtaf.fetch_by_position("chr1")) self.assertEqual(len(reads), 2) - self.assertTrue(all(_.qname.startswith('se_good') for _ in reads)) - - def test_check_pe_flags(self): + for algn_seg in reads: + self.assertTrue( + algn_seg.query_name is not None and \ + algn_seg.query_name.startswith("se_good"), + ) + + def test_check_pe_flags(self) -> None: + """Check for SAM paired end flags.""" for idx, flag in enumerate([83, 99]): self.sam_obj.add_read_pair( - 'chr1', + "chr1", Sequence( self.refseq, 0, flag=flag, - qname=f'pe_good_{idx}', + read_name=f"pe_good_{idx}", ), ) bad_flags = [73, 89, 137, 153, 329, 339, 345, 355, 393, 409] for idx, flag in enumerate(bad_flags): self.sam_obj.add_read_pair( - 'chr1', + "chr1", Sequence( self.refseq, 0, flag=flag, - qname=f'pe_bad_{idx}', + read_name=f"pe_bad_{idx}", ), ) @@ -150,6 +167,10 @@ def test_check_pe_flags(self): self.sam_obj.save_to_sam(self.bam_fname, self.genome_fname) with RTAlignmentFile(self.bam_fname) as rtaf: - reads = next(rtaf.fetch_by_position('chr1')) + reads = next(rtaf.fetch_by_position("chr1")) self.assertEqual(len(reads), 4) - self.assertTrue(all(_.qname.startswith('pe_good') for _ in reads)) + for read in reads: + self.assertTrue( + read.query_name is not None and \ + read.query_name.startswith("pe_good"), + ) diff --git a/test/alignment_manager.py b/test/alignment_manager.py index c88ebe9..fc177ac 100644 --- a/test/alignment_manager.py +++ b/test/alignment_manager.py @@ -1,72 +1,67 @@ -import os +"""Test cases for AlignmentManager class.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, Sequence, ntf from reditools.alignment_manager import AlignmentManager class TestAlignmentManager(unittest.TestCase): + """Test cases for AlignmentManager class.""" - def test_propagation(self): - genome_fname = ntf(suffix='.fa') - bam_fname = ntf(suffix='.bam') + def setUp(self) -> None: + """Pre-flight setup.""" + self.genome_fname = ntf(suffix=".fa") + self.bam_fname1 = ntf(suffix=".bam") + self.bam_fname2 = ntf(suffix=".bam") sam_obj = SAM() - sam_obj.add_contig('chr1') - sam_obj.genome.save_to_fasta(genome_fname) - sam_obj.save_to_sam(bam_fname, genome_fname) + sam_obj.add_contig("chr1", length=80) + refseq = sam_obj.genome["chr1"] + sam_obj.genome.save_to_fasta(self.genome_fname) + + sam_obj.add_read("chr1", Sequence(refseq, 0, read_name="1_1")) + sam_obj.add_read("chr1", Sequence(refseq[20:], 20, read_name="1_2")) + sam_obj.add_read("chr1", Sequence(refseq[40:], 40, read_name="1_3")) + sam_obj.save_to_sam(self.bam_fname1, self.genome_fname) + sam_obj = SAM() + sam_obj.add_contig("chr1", sequence=refseq) + sam_obj.add_read("chr1", Sequence(refseq[20:], 20, read_name="2_1")) + sam_obj.add_read("chr1", Sequence(refseq[50:], 50, read_name="2_2")) + sam_obj.save_to_sam(self.bam_fname2, self.genome_fname) + + def tearDown(self) -> None: + """Post checks cleanup.""" + Path(self.genome_fname).unlink() + Path(self.bam_fname1).unlink() + Path(self.bam_fname2).unlink() + + def test_propagation(self) -> None: + """Check that properties of AlignmentManager propagate to sub files.""" rtam = AlignmentManager(min_length=10, min_quality=30) - rtam.add_file(bam_fname) + rtam.add_file(self.bam_fname1) self.assertEqual(rtam._bams[0].readqc.min_length, 10) self.assertEqual(rtam._bams[0].readqc.min_quality, 30) - os.remove(genome_fname) - os.remove(bam_fname) - - def test_fetch_by_position(self): - genome_fname, bam_fnames = self.setup_dummy_data() - + def test_fetch_by_position(self) -> None: + """Check fetch_by_position works for all sub files.""" rtam = AlignmentManager(min_length=10, min_quality=30) - rtam.add_file(bam_fnames[0]) - rtam.add_file(bam_fnames[1]) + rtam.add_file(self.bam_fname1) + rtam.add_file(self.bam_fname2) - read_iter = rtam.fetch_by_position('chr1') + read_iter = rtam.fetch_by_position("chr1") read_group = next(read_iter) self.assertEqual(len(read_group), 1) - self.assertEqual(read_group[0].qname, '1_1') + self.assertEqual(read_group[0].query_name, "1_1") self.assertEqual(rtam.next_read_start, 20) read_group = next(read_iter) self.assertEqual(len(read_group), 2) - self.assertIn('2_1', (_.qname for _ in read_group)) - self.assertIn('1_2', (_.qname for _ in read_group)) + self.assertIn("2_1", (_.query_name for _ in read_group)) + self.assertIn("1_2", (_.query_name for _ in read_group)) self.assertEqual(rtam.next_read_start, 40) - - os.remove(genome_fname) - for fname in bam_fnames: - os.remove(fname) - - def setup_dummy_data(self): - genome_fname = ntf(suffix='.fa') - bam_fnames = [ntf(suffix='.bam') for _ in range(2)] - - sam_obj = SAM() - sam_obj.add_contig('chr1', length=80) - refseq = sam_obj.genome['chr1'] - sam_obj.genome.save_to_fasta(genome_fname) - - sam_obj.add_read('chr1', Sequence(refseq, 0, qname='1_1')) - sam_obj.add_read('chr1', Sequence(refseq[20:], 20, qname='1_2')) - sam_obj.add_read('chr1', Sequence(refseq[40:], 40, qname='1_3')) - sam_obj.save_to_sam(bam_fnames[0], genome_fname) - - sam_obj = SAM() - sam_obj.add_contig('chr1', sequence=refseq) - sam_obj.add_read('chr1', Sequence(refseq[20:], 20, qname='2_1')) - sam_obj.add_read('chr1', Sequence(refseq[50:], 50, qname='2_2')) - sam_obj.save_to_sam(bam_fnames[1], genome_fname) - - return genome_fname, bam_fnames diff --git a/test/analyze/parse_args.py b/test/analyze/parse_args.py index d74f4e8..fc66d13 100644 --- a/test/analyze/parse_args.py +++ b/test/analyze/parse_args.py @@ -1,84 +1,104 @@ +"""Test cases for analyze parse_args.""" +from __future__ import annotations + import sys import unittest from contextlib import contextmanager from io import StringIO +from typing import Iterator from reditools import reditools from reditools.tools.analyze.parse_args.parse_args import parse_args class TestParseArgs(unittest.TestCase): - def test_legacy_pruning(self): - args = parse_args(['test/test.bam']) - self.assertFalse(hasattr(args, 'dna')) - self.assertFalse(hasattr(args, 'exclude_multis')) - self.assertFalse(hasattr(args, 'strict')) - self.assertFalse(hasattr(args, 'load_omopolymeric_file')) - - def test_dna_mode(self): + """Test cases for analyze parse_args.""" + + def test_legacy_pruning(self) -> None: + """Check that legacy options are removed from the Namespace.""" + args = parse_args(["test/test.bam"]) + self.assertFalse(hasattr(args, "dna")) + self.assertFalse(hasattr(args, "exclude_multis")) + self.assertFalse(hasattr(args, "strict")) + self.assertFalse(hasattr(args, "load_omopolymeric_file")) + + def test_dna_mode(self) -> None: + """Check that --dna sets --strand to zero.""" args = parse_args([ - 'test/test.bam', - '--dna', + "test/test.bam", + "--dna", ]) self.assertEqual(args.strand, reditools.UNSTRANDED_MODE) - self.assertFalse(hasattr(args, 'dna')) + self.assertFalse(hasattr(args, "dna")) - def test_exclude_multis(self): + def test_exclude_multis(self) -> None: + """Check that --exclude-multis sets --max-editing-nucleotides to one.""" args = parse_args([ - 'test/test.bam', - '--exclude-multis', + "test/test.bam", + "--exclude-multis", ]) self.assertEqual(args.max_editing_nucleotides, 1) - self.assertFalse(hasattr(args, 'exclude_multis')) + self.assertFalse(hasattr(args, "exclude_multis")) - def test_strict(self): + def test_strict(self) -> None: + """Check that --strict sets --min-edits to one.""" args = parse_args([ - 'test/test.bam', - '--strict', + "test/test.bam", + "--strict", ]) self.assertEqual(args.min_edits, 1) - self.assertFalse(hasattr(args, 'strict')) + self.assertFalse(hasattr(args, "strict")) - def test_load_omopolymeric_file(self): + def test_load_omopolymeric_file(self) -> None: + """Check that --load-omopolymeric-file appends to exclude_regions.""" args = parse_args([ - 'test/test.bam', - '--load-omopolymeric-file', - 'test/test.bed', + "test/test.bam", + "--load-omopolymeric-file", + "test/test.bed", ]) - self.assertEqual(args.exclude_regions, ['test/test.bed']) - self.assertFalse(hasattr(args, 'load_omopolymeric_file')) + self.assertEqual(args.exclude_regions, ["test/test.bed"]) + self.assertFalse(hasattr(args, "load_omopolymeric_file")) args = parse_args([ - 'test/test.bam', - '--exclude-regions', 'example.bed', - '--load-omopolymeric-file', 'test/test.bed', + "test/test.bam", + "--exclude-regions", "example.bed", + "--load-omopolymeric-file", "test/test.bed", ]) self.assertEqual( args.exclude_regions, - ['example.bed', 'test/test.bed'], + ["example.bed", "test/test.bed"], ) - self.assertFalse(hasattr(args, 'load_omopolymeric_file')) - - def test_edit_frequency(self): - with self.assertRaises(SystemExit): - with self.capture_sys_output() as (stdout, stderr): - parse_args([ - 'test/test.bam', - '--max-editing-nucleotides', '1', - '--min-edits', '3', - ]) - - def test_unstranded(self): - with self.assertRaises(SystemExit): - with self.capture_sys_output() as (stdout, stderr): - parse_args([ - 'test/test.bam', - '--strand', '0', - '--strand-correction', - ]) + self.assertFalse(hasattr(args, "load_omopolymeric_file")) + + def test_edit_frequency(self) -> None: + """Check for error when min-edits is above max-editing-nucleotides.""" + with self.assertRaises(SystemExit), self.capture_sys_output(): + parse_args([ + "test/test.bam", + "--max-editing-nucleotides", "1", + "--min-edits", "3", + ]) + + def test_unstranded(self) -> None: + """Check for error when strand is zero whiel using strand-correction.""" + with self.assertRaises(SystemExit), self.capture_sys_output(): + parse_args([ + "test/test.bam", + "--strand", "0", + "--strand-correction", + ]) @contextmanager - def capture_sys_output(self): + def capture_sys_output(self) -> Iterator[tuple[StringIO, StringIO]]: + """Capture standard and error output. + + This method is used to silence the output for parse_args(). + + Returns + ------- + Iterator[tuple[StringIO, StringIO]] + Stdout and stderr, respectively. + """ capture_out, capture_err = StringIO(), StringIO() current_out, current_err = sys.stdout, sys.stderr try: # noqa: WPS229 diff --git a/test/analyze/parse_args_utils.py b/test/analyze/parse_args_utils.py index 235bfb7..e50573d 100644 --- a/test/analyze/parse_args_utils.py +++ b/test/analyze/parse_args_utils.py @@ -1,59 +1,80 @@ -import argparse +"""Test cases for analyze bounded_types module.""" +from __future__ import annotations + import unittest -from reditools.tools.analyze.parse_args.parse_args import (bounded_float, - bounded_int, - check_number_bounds) +from reditools.tools.analyze.parse_args.bounded_types import ( + CastFloatError, + CastIntError, + ValueAboveMaximumError, + ValueBelowMinimumError, + bounded_float, + bounded_int, + check_number_bounds, +) class TestParseArgsUtils(unittest.TestCase): - def test_check_number_bounds_valid(self): - # No error for value in bounds + """Test cases for analyze bounded_types module.""" + + def test_check_number_bounds_valid(self) -> None: + """Check check_number_bounds() method for acceptable values.""" check_number_bounds(5, min_value=1, max_value=10) check_number_bounds(1, min_value=1) check_number_bounds(10, max_value=10) - def test_check_number_bounds_too_low(self): - with self.assertRaises(argparse.ArgumentTypeError): + def test_check_number_bounds_too_low(self) -> None: + """Check check_number_bounds() with below minimum input.""" + with self.assertRaises(ValueBelowMinimumError): check_number_bounds(0, min_value=1) - def test_check_number_bounds_too_high(self): - with self.assertRaises(argparse.ArgumentTypeError): + def test_check_number_bounds_too_high(self) -> None: + """Check check_number_bounds() with above maximum input.""" + with self.assertRaises(ValueAboveMaximumError): check_number_bounds(11, max_value=10) - def test_bounded_int_valid(self): + def test_bounded_int_valid(self) -> None: + """Check bounded_int() method with valid input.""" conv = bounded_int(min_value=2, max_value=6) - self.assertEqual(conv('4'), 4) - self.assertEqual(conv('2'), 2) - self.assertEqual(conv('6'), 6) + self.assertEqual(conv("4"), 4) + self.assertEqual(conv("2"), 2) + self.assertEqual(conv("6"), 6) - def test_bounded_int_invalid_type(self): + def test_bounded_int_invalid_type(self) -> None: + """Check bounded_int() method with non-int input.""" conv = bounded_int() - with self.assertRaises(argparse.ArgumentTypeError): - conv('foo') + with self.assertRaises(CastIntError): + conv("2.1") + with self.assertRaises(CastIntError): + conv("foo") - def test_bounded_int_out_of_bounds(self): + def test_bounded_int_out_of_bounds(self) -> None: + """Check bounded_int method with out of bounds input.""" conv = bounded_int(min_value=3) - with self.assertRaises(argparse.ArgumentTypeError): - conv('1') + with self.assertRaises(ValueBelowMinimumError): + conv("1") conv = bounded_int(max_value=1) - with self.assertRaises(argparse.ArgumentTypeError): - conv('2') + with self.assertRaises(ValueAboveMaximumError): + conv("2") - def test_bounded_float_valid(self): + def test_bounded_float_valid(self) -> None: + """Check bounded_float() method with valid input.""" conv = bounded_float(min_value=0.5, max_value=2.6) - self.assertEqual(conv('1.2'), 1.2) - self.assertEqual(conv('0.5'), 0.5) - self.assertEqual(conv('2.6'), 2.6) + self.assertEqual(conv("1.2"), 1.2) + self.assertEqual(conv("0.5"), 0.5) + self.assertEqual(conv("2.6"), 2.6) + self.assertEqual(conv("1"), 1.0) - def test_bounded_float_invalid_type(self): + def test_bounded_float_invalid_type(self) -> None: + """Check bounded_float() with non-float input.""" conv = bounded_float() - with self.assertRaises(argparse.ArgumentTypeError): - conv('hello') + with self.assertRaises(CastFloatError): + conv("hello") - def test_bounded_float_out_of_bounds(self): + def test_bounded_float_out_of_bounds(self) -> None: + """Check bounded_float() with out of bounds input.""" conv = bounded_float(min_value=0.1, max_value=1.2) - with self.assertRaises(argparse.ArgumentTypeError): - conv('0.01') - with self.assertRaises(argparse.ArgumentTypeError): - conv('2.0') + with self.assertRaises(ValueBelowMinimumError): + conv("0.01") + with self.assertRaises(ValueAboveMaximumError): + conv("2.0") diff --git a/test/analyze/region_args.py b/test/analyze/region_args.py index 95209cd..400e270 100644 --- a/test/analyze/region_args.py +++ b/test/analyze/region_args.py @@ -1,5 +1,8 @@ -import os +"""Test cases for region_args module.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, ntf from reditools.region import Region @@ -8,44 +11,52 @@ class TestRegionArgs(unittest.TestCase): - def setUp(self): - self.fasta_fname = ntf(suffix='.fa') - self.bam_fname = ntf(suffix='.bam') + """Test cases for region_args module.""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.fasta_fname = ntf(suffix=".fa") + self.bam_fname = ntf(suffix=".bam") sam_obj = SAM() - sam_obj.add_contig('chr1', length=120) - sam_obj.add_contig('chr2', length=80) - sam_obj.add_contig('chr3', length=60) + sam_obj.add_contig("chr1", length=120) + sam_obj.add_contig("chr2", length=80) + sam_obj.add_contig("chr3", length=60) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) - def tearDown(self): - os.remove(self.fasta_fname) - os.remove(self.bam_fname) + def tearDown(self) -> None: + """Post-checks cleanup.""" + Path(self.fasta_fname).unlink() + Path(self.bam_fname).unlink() - def test_no_input(self): + def test_no_input(self) -> None: + """Check region_args() with no options.""" options = parse_args([self.bam_fname]) regions = region_args(options) self.assertEqual(len(regions), 3) - def test_region_input(self): - options = parse_args([self.bam_fname, '--region', 'chr1:1-100']) + def test_region_input(self) -> None: + """Check region_args() with specified region.""" + options = parse_args([self.bam_fname, "--region", "chr1:1-100"]) regions = region_args(options) - self.assertEqual(regions, [Region('chr1', 0, 100)]) + self.assertEqual(regions, [Region("chr1", 0, 100)]) - def test_region_window(self): + def test_region_window(self) -> None: + """Check region_args() with specified region and window size.""" options = parse_args([ self.bam_fname, - '--region', - 'chr1:1-100', - '--window', - '10', + "--region", + "chr1:1-100", + "--window", + "10", ]) regions = region_args(options) self.assertEqual(len(regions), 10) - def test_bam_window(self): - options = parse_args([self.bam_fname, '--window', '70']) + def test_bam_window(self) -> None: + """Check region_args() with window size.""" + options = parse_args([self.bam_fname, "--window", "70"]) regions = region_args(options) self.assertEqual(len(regions), 5) diff --git a/test/analyze/rtchecks.py b/test/analyze/rtchecks.py index 7f8256a..087c277 100644 --- a/test/analyze/rtchecks.py +++ b/test/analyze/rtchecks.py @@ -1,6 +1,9 @@ -import os +"""Test cases for RTChecks class.""" +from __future__ import annotations + import unittest from argparse import Namespace +from pathlib import Path from tempfile import NamedTemporaryFile from reditools.compiled_position import CompiledPosition, RTResult @@ -8,28 +11,44 @@ class TestRTChecks(unittest.TestCase): - def setUp(self): - self.bases = CompiledPosition(contig='chr1', position=1, ref='A') + """Test cases for RTChecks class.""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.bases = CompiledPosition(contig="chr1", position=1, ref="A") self.options = Namespace( max_editing_nucleotides=4, min_read_depth=0, min_edits=0, min_edits_per_nucleotide=0, - variants=['all'], + variants=["all"], exclude_regions=None, splicing_file=None, bed_file=None, ) - def run_check(self, rtc): - return rtc.check(RTResult(self.bases, '*')) + def run_check(self, rtc: RTChecks) -> tuple | None: + """Perform a quality control check. + + Parameters + ---------- + rtc : RTChecks + An RTChecks QC object. + + Returns + ------- + tuple | None + Output fo rtc.check() + """ + return rtc.check(RTResult(self.bases, "*")) - def test_check_column_edit_frequency(self): + def test_check_column_edit_frequency(self) -> None: + """Check --min-edits input.""" self.options.min_edits = 1 rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='A', strand='*') + self.bases.add_base(quality=30, base="A", strand="*") self.assertIsNotNone(self.run_check(rtc)) self.options.min_edits = 0 @@ -38,57 +57,60 @@ def test_check_column_edit_frequency(self): self.options.min_edits = 1 rtc = RTChecks(self.options) - self.bases.add_base(quality=30, base='T', strand='*') - self.bases.add_base(quality=30, base='T', strand='*') + self.bases.add_base(quality=30, base="T", strand="*") + self.bases.add_base(quality=30, base="T", strand="*") self.assertIsNone(self.run_check(rtc)) self.options.min_edits = 3 rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - def test_check_column_min_edits(self): + def test_check_column_min_edits(self) -> None: + """Check min-edits-per-nucleotide input.""" self.options.min_edits_per_nucleotide = 1 rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='A', strand='*') - self.bases.add_base(quality=30, base='A', strand='*') + self.bases.add_base(quality=30, base="A", strand="*") + self.bases.add_base(quality=30, base="A", strand="*") self.assertIsNone(self.run_check(rtc)) self.options.min_edits_per_nucleotide = 2 rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='T', strand='*') + self.bases.add_base(quality=30, base="T", strand="*") self.assertIsNotNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='T', strand='*') + self.bases.add_base(quality=30, base="T", strand="*") self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='C', strand='*') + self.bases.add_base(quality=30, base="C", strand="*") self.assertIsNotNone(self.run_check(rtc)) - def test_check_min_read_depth(self): + def test_check_min_read_depth(self) -> None: + """Check min-read-depth input.""" self.options.min_read_depth = 2 rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='C', strand='*') + self.bases.add_base(quality=30, base="C", strand="*") self.assertIsNotNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='A', strand='*') + self.bases.add_base(quality=30, base="A", strand="*") self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='A', strand='*') + self.bases.add_base(quality=30, base="A", strand="*") self.assertIsNone(self.run_check(rtc)) - def test_check_exclusions(self): + def test_check_exclusions(self) -> None: + """Check exclude-regions input.""" with NamedTemporaryFile( delete=False, - suffix='.bed', - mode='w+', + suffix=".bed", + mode="w+", ) as stream: - stream.write('chr1\t20\t30\n') + stream.write("chr1\t20\t30\n") bed_file = stream.name self.options.exclude_regions = [bed_file] @@ -96,25 +118,26 @@ def test_check_exclusions(self): rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - with open(bed_file, mode='a') as stream: - stream.write('chr1\t0\t10\n') + with Path(bed_file).open("a") as stream: + stream.write("chr1\t0\t10\n") rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - with open(bed_file, mode='a') as stream: - stream.write('chr2\t0\t10\n') + with Path(bed_file).open("a") as stream: + stream.write("chr2\t0\t10\n") rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - os.remove(bed_file) + Path(bed_file).unlink() - def test_check_splicing(self): + def test_check_splicing(self) -> None: + """Check splicing-file input.""" with NamedTemporaryFile( delete=False, - suffix='.txt', - mode='w+', + suffix=".txt", + mode="w+", ) as stream: - stream.write('chr1 1 4 A +\n') + stream.write("chr1 1 4 A +\n") splice_file = stream.name self.options.splicing_file = splice_file @@ -123,62 +146,64 @@ def test_check_splicing(self): rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - with open(splice_file, mode='w') as stream: - stream.write('chr1 1 4 A -\n') + with Path(splice_file).open("w") as stream: + stream.write("chr1 1 4 A -\n") rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - with open(splice_file, mode='w') as stream: - stream.write('chr1 1 4 D +\n') + with Path(splice_file).open("w") as stream: + stream.write("chr1 1 4 D +\n") rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - with open(splice_file, mode='w') as stream: - stream.write('chr1 1 4 D -\n') + with Path(splice_file).open("w") as stream: + stream.write("chr1 1 4 D -\n") rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - os.remove(splice_file) + Path(splice_file).unlink() - def test_check_max_editing_nucleotides(self): + def test_check_max_editing_nucleotides(self) -> None: + """Check max-editing-nucleotides input.""" self.options.max_editing_nucleotides = 1 rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='A', strand='*') + self.bases.add_base(quality=30, base="A", strand="*") self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='T', strand='*') - self.bases.add_base(quality=30, base='T', strand='*') + self.bases.add_base(quality=30, base="T", strand="*") + self.bases.add_base(quality=30, base="T", strand="*") self.assertIsNone(self.run_check(rtc)) - self.bases.add_base(quality=30, base='C', strand='*') + self.bases.add_base(quality=30, base="C", strand="*") self.assertIsNotNone(self.run_check(rtc)) self.options.max_editing_nucleotides = 2 rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - def test_check_target_positions(self): + def test_check_target_positions(self) -> None: + """Check bed-file input.""" with NamedTemporaryFile( delete=False, - suffix='.bed', - mode='w+', + suffix=".bed", + mode="w+", ) as stream: - stream.write('chr1\t10\t20\n') + stream.write("chr1\t10\t20\n") bed_file = stream.name self.options.bed_file = [bed_file] rtc = RTChecks(self.options) self.assertIsNotNone(self.run_check(rtc)) - with open(bed_file, mode='a') as stream: - stream.write('chr1\t0\t20\n') + with Path(bed_file).open("a") as stream: + stream.write("chr1\t0\t20\n") rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - with open(bed_file, mode='a') as stream: - stream.write('chr2\t0\t20\n') + with Path(bed_file).open("a") as stream: + stream.write("chr2\t0\t20\n") rtc = RTChecks(self.options) self.assertIsNone(self.run_check(rtc)) - os.remove(bed_file) + Path(bed_file).unlink() diff --git a/test/analyze/setup_alignment_manager.py b/test/analyze/setup_alignment_manager.py index e314179..5ba8bd4 100644 --- a/test/analyze/setup_alignment_manager.py +++ b/test/analyze/setup_alignment_manager.py @@ -1,28 +1,35 @@ -import os +"""Test cases for setup_alignment_manager() method.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, ntf -from reditools.tools.analyze.setup_alignment_manager import \ - setup_alignment_manager +from reditools.tools.analyze.setup_alignment_manager import ( + setup_alignment_manager, +) class TestSetupAlignmentManager(unittest.TestCase): - def test_setup(self): - fasta_fname = ntf(suffix='.fa') - bam_fname = ntf(suffix='.bam') + """Test cases for setup_alignment_manager() method.""" + + def test_setup(self) -> None: + """Check setup_alignment_manager() method.""" + fasta_fname = ntf(suffix=".fa") + bam_fname = ntf(suffix=".bam") sam_obj = SAM() - sam_obj.add_contig('chr1', length=120) - sam_obj.add_contig('chr2', length=80) - sam_obj.add_contig('chr3', length=60) + sam_obj.add_contig("chr1", length=120) + sam_obj.add_contig("chr2", length=80) + sam_obj.add_contig("chr3", length=60) sam_obj.genome.save_to_fasta(fasta_fname) sam_obj.save_to_sam(bam_fname, fasta_fname) - exclusions_fname = ntf(suffix='.bed') - with open(exclusions_fname, 'w') as stream: - stream.write('bad_read') + exclusions_fname = ntf(suffix=".bed") + with Path(exclusions_fname).open("w") as stream: + stream.write("bad_read") rtam = setup_alignment_manager( [bam_fname], @@ -33,8 +40,8 @@ def test_setup(self): self.assertEqual(rtam.min_quality, 50) self.assertEqual(rtam.min_length, 123) - self.assertIn('bad_read', rtam.excluded_read_names) + self.assertIn("bad_read", rtam.excluded_read_names) - os.remove(fasta_fname) - os.remove(bam_fname) - os.remove(exclusions_fname) + Path(fasta_fname).unlink() + Path(bam_fname).unlink() + Path(exclusions_fname).unlink() diff --git a/test/analyze/setup_rtools.py b/test/analyze/setup_rtools.py index 37a0cdc..9fd9350 100644 --- a/test/analyze/setup_rtools.py +++ b/test/analyze/setup_rtools.py @@ -1,3 +1,6 @@ +"""Test cases for setup_rtools() method.""" +from __future__ import annotations + import unittest from reditools.reditools import REDItools @@ -6,20 +9,23 @@ class TestSetupRTools(unittest.TestCase): - def test_options(self): + """Test cases for setup_rtools() method.""" + + def test_options(self) -> None: + """Check setup_rtools() method.""" options = parse_args([ - 'example.bam', - '-r', 'genome.fa', - '--min-base-position', '10', - '--max-base-position', '30', - '--min-base-quality', '23', - '--strand', '1', - '--strand-confidence-threshold', '0.567', - '--strand-correction', + "example.bam", + "-r", "genome.fa", + "--min-base-position", "10", + "--max-base-position", "30", + "--min-base-quality", "23", + "--strand", "1", + "--strand-confidence-threshold", "0.567", + "--strand-correction", ]) rtools = setup_rtools(options) self.assertIsInstance(rtools, REDItools) - self.assertEqual(rtools.reference, 'genome.fa') + self.assertEqual(rtools.reference, "genome.fa") self.assertEqual(rtools.min_base_position, 10) self.assertEqual(rtools.max_base_position, 30) self.assertEqual(rtools.min_base_quality, 23) diff --git a/test/compiled_position.py b/test/compiled_position.py index 8bc5ceb..5a56009 100644 --- a/test/compiled_position.py +++ b/test/compiled_position.py @@ -1,115 +1,127 @@ +"""Test cases for CompiledPosition and RTResult classes.""" +from __future__ import annotations + import unittest from reditools.compiled_position import CompiledPosition, RTResult class TestCompiledPosition(unittest.TestCase): - def setUp(self): - self.cp = CompiledPosition('A', 'chr1', 100) + """Test cases for CompiledPosition and RTResult classes.""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.cp = CompiledPosition(ref="A", contig="chr1", position=100) - def test_add_base_and_len(self): - self.cp.add_base(40, '+', 'A') - self.cp.add_base(35, '-', 'C') - self.cp.add_base(30, '+', 'G') + def test_len(self) -> None: + """Check len() functions for both classes.""" + self.cp.add_base(40, "+", "A") + self.cp.add_base(35, "-", "C") + self.cp.add_base(30, "+", "G") self.assertEqual(len(self.cp), 3) + rtresult = RTResult(self.cp, "*") + self.assertEqual(len(rtresult), 3) - def test_complement(self): - self.cp.add_base(11, '+', 'A') - self.cp.add_base(12, '-', 'C') + def test_complement(self) -> None: + """Check compelement() function of CompiledPosition.""" + self.cp.add_base(11, "+", "A") + self.cp.add_base(12, "-", "C") self.cp.complement() - self.assertEqual(self.cp.bases, ['T', 'G']) - self.assertEqual(self.cp.ref, 'T') - - def test_calculate_strand(self): - self.cp.add_base(20, '+', 'A') - self.cp.add_base(21, '-', 'A') - self.cp.add_base(22, '+', 'C') - self.assertEqual(self.cp.calculate_strand(), '+') - self.assertEqual(self.cp.calculate_strand(0.7), '*') - - def test_filter_by_strand(self): - self.cp.add_base(5, '+', 'A') - self.cp.add_base(5, '+', 'A') - self.cp.add_base(6, '-', 'C') - self.cp.filter_by_strand('+') + self.assertEqual(self.cp.bases, ["T", "G"]) + self.assertEqual(self.cp.ref, "T") + + def test_calculate_strand(self) -> None: + """Check calculate_strand() function of CompiledPosition.""" + self.cp.add_base(20, "+", "A") + self.cp.add_base(21, "-", "A") + self.cp.add_base(22, "+", "C") + self.assertEqual(self.cp.calculate_strand(), "+") + self.assertEqual(self.cp.calculate_strand(0.7), "*") + + def test_filter_by_strand(self) -> None: + """Check fitler_by_strand() function.""" + self.cp.add_base(5, "+", "A") + self.cp.add_base(5, "+", "A") + self.cp.add_base(6, "-", "C") + self.cp.filter_by_strand("+") self.assertEqual(len(self.cp), 2) - rtresult = RTResult(self.cp, '*') - self.assertEqual(rtresult['A'], 2) - self.assertEqual(rtresult['C'], 0) - - def test_filter_by_strand_star(self): - self.cp.add_base(5, '*', 'A') - self.cp.add_base(5, '*', 'A') - self.cp.add_base(6, '+', 'C') - self.cp.add_base(6, '-', 'T') + rtresult = RTResult(self.cp, "*") + self.assertEqual(rtresult["A"], 2) + self.assertEqual(rtresult["C"], 0) + + def test_filter_by_strand_star(self) -> None: + """Check filter_by_strand() with undetermined strand.""" + self.cp.add_base(5, "*", "A") + self.cp.add_base(5, "*", "A") + self.cp.add_base(6, "+", "C") + self.cp.add_base(6, "-", "T") self.assertEqual( self.cp.calculate_strand(threshold=1), - '*', + "*", ) - self.cp.filter_by_strand('*') - rtresult = RTResult(self.cp, '*') + self.cp.filter_by_strand("*") + rtresult = RTResult(self.cp, "*") self.assertEqual(len(rtresult), 4) - self.assertEqual(rtresult['A'], 2) - self.assertEqual(rtresult['C'], 1) - - def test_reference(self): - self.assertEqual(self.cp.ref, 'A') - rtresult = RTResult(self.cp, '*') - self.assertEqual(rtresult.reference, 'A') - - def test_len(self): - self.assertEqual(len(self.cp), 0) - self.cp.add_base(40, '+', 'A') - self.assertEqual(len(self.cp), 1) - rtresult = RTResult(self.cp, '*') - self.assertEqual(len(rtresult), 1) - - def test_get_base_counts(self): - self.cp.add_base(40, '+', 'A') - self.cp.add_base(35, '-', 'A') - self.cp.add_base(30, '+', 'C') - rtresult = RTResult(self.cp, '*') - self.assertEqual(rtresult['A'], 2) - self.assertEqual(rtresult['C'], 1) - self.assertEqual(rtresult['REF'], 2) - - def test_iter(self): - self.cp.add_base(41, '+', 'A') - self.cp.add_base(42, '+', 'C') - self.cp.add_base(43, '+', 'G') - counts = list(RTResult(self.cp, '*')) + self.assertEqual(rtresult["A"], 2) + self.assertEqual(rtresult["C"], 1) + + def test_reference(self) -> None: + """Check reference base propagates to RTResult class.""" + self.assertEqual(self.cp.ref, "A") + rtresult = RTResult(self.cp, "*") + self.assertEqual(rtresult.reference, "A") + + def test_get_base_counts(self) -> None: + """Check RTResult base count summary.""" + self.cp.add_base(40, "+", "A") + self.cp.add_base(35, "-", "A") + self.cp.add_base(30, "+", "C") + rtresult = RTResult(self.cp, "*") + self.assertEqual(rtresult["A"], 2) + self.assertEqual(rtresult["C"], 1) + self.assertEqual(rtresult["REF"], 2) + + def test_iter(self) -> None: + """Check list casting for RTResult.""" + self.cp.add_base(41, "+", "A") + self.cp.add_base(42, "+", "C") + self.cp.add_base(43, "+", "G") + counts = list(RTResult(self.cp, "*")) self.assertEqual(counts, [1, 1, 1, 0]) - def test_variants(self): - self.cp.add_base(10, '+', 'C') - self.cp.add_base(10, '+', 'A') - rtresult = RTResult(self.cp, '*') - self.assertEqual(rtresult.variants, ['AC']) + def test_variants(self) -> None: + """Check variant summary for RTResult.""" + self.cp.add_base(10, "+", "C") + self.cp.add_base(10, "+", "A") + rtresult = RTResult(self.cp, "*") + self.assertEqual(rtresult.variants, ["AC"]) - def test_edit_ratio(self): - rtresult = RTResult(self.cp, '*') + def test_edit_ratio(self) -> None: + """Check edit_ratio cacluation from RTResult.""" + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.edit_ratio, 0) - self.cp.add_base(10, '+', 'A') - self.cp.add_base(10, '+', 'C') - rtresult = RTResult(self.cp, '*') + self.cp.add_base(10, "+", "A") + self.cp.add_base(10, "+", "C") + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.edit_ratio, 0.5) - self.cp.add_base(10, '+', 'T') - rtresult = RTResult(self.cp, '*') + self.cp.add_base(10, "+", "T") + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.edit_ratio, 0.5) - self.cp.add_base(10, '+', 'C') - self.cp.add_base(10, '+', 'C') - rtresult = RTResult(self.cp, '*') + self.cp.add_base(10, "+", "C") + self.cp.add_base(10, "+", "C") + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.edit_ratio, 0.75) - def test_mean_quality(self): - rtresult = RTResult(self.cp, '*') + def test_mean_quality(self) -> None: + """Check mean_quality calculation from RTResult.""" + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.mean_quality, 0) - self.cp.add_base(10, '+', 'C') - self.cp.add_base(20, '+', 'C') - self.cp.add_base(30, '+', 'C') - rtresult = RTResult(self.cp, '*') + self.cp.add_base(10, "+", "C") + self.cp.add_base(20, "+", "C") + self.cp.add_base(30, "+", "C") + rtresult = RTResult(self.cp, "*") self.assertEqual(rtresult.mean_quality, 20) diff --git a/test/compiled_reads.py b/test/compiled_reads.py index a40c320..8d54b32 100644 --- a/test/compiled_reads.py +++ b/test/compiled_reads.py @@ -1,5 +1,8 @@ -import os +"""Test cases for CompiledReads class.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, Sequence, ntf from pysam import AlignmentFile @@ -8,22 +11,27 @@ class TestCompiledReads(unittest.TestCase): - def setUp(self): - self.fasta_fname = ntf(suffix='.fa') - self.bam_fname = ntf(suffix='.bam') + """Test cases for CompiledReads class.""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.fasta_fname = ntf(suffix=".fa") + self.bam_fname = ntf(suffix=".bam") - def tearDown(self): - os.remove(self.fasta_fname) - os.remove(self.bam_fname) + def tearDown(self) -> None: + """Post-check cleanup.""" + Path(self.fasta_fname).unlink() + Path(self.bam_fname).unlink() - def test_ref_seq_spliced(self): + def test_ref_seq_spliced(self) -> None: + """Check ability to get reference sequence from spliced reads.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=60) - spliceseq = sam_obj.genome['chr1'] + sam_obj.add_contig("chr1", length=60) + spliceseq = sam_obj.genome["chr1"] spliceseq = spliceseq[:20] + spliceseq[40:60] sam_obj.add_read( - 'chr1', - Sequence(spliceseq, 0, _cigar_str='20M20D20M'), + "chr1", + Sequence(spliceseq, 0, _cigar_str="20M20D20M"), ) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -33,14 +41,15 @@ def test_ref_seq_spliced(self): with AlignmentFile(self.bam_fname) as af: read = next(af.fetch()) - self.assertEqual(''.join(md_ref_fetch.get_refseq(read)), spliceseq) - self.assertEqual(''.join(fa_ref_fetch.get_refseq(read)), spliceseq) + self.assertEqual("".join(md_ref_fetch.get_refseq(read)), spliceseq) + self.assertEqual("".join(fa_ref_fetch.get_refseq(read)), spliceseq) - def test_ref_seq_unspliced(self): + def test_ref_seq_unspliced(self) -> None: + """Check ability to get reference sequence from contiguous reads.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=60) - refseq = sam_obj.genome['chr1'] - sam_obj.add_read('chr1', Sequence(refseq, 0)) + sam_obj.add_contig("chr1", length=60) + refseq = sam_obj.genome["chr1"] + sam_obj.add_read("chr1", Sequence(refseq, 0)) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -49,16 +58,20 @@ def test_ref_seq_unspliced(self): with AlignmentFile(self.bam_fname) as af: read = next(af.fetch()) - self.assertEqual(''.join(md_ref_fetch.get_refseq(read)), refseq) - self.assertEqual(''.join(fa_ref_fetch.get_refseq(read)), refseq) + self.assertEqual("".join(md_ref_fetch.get_refseq(read)), refseq) + self.assertEqual("".join(fa_ref_fetch.get_refseq(read)), refseq) - def test_ref_seq_snp(self): + def test_ref_seq_snp(self) -> None: + """Check ability to get reference sequence from reads with SNPs.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=60) - snpseq = list(sam_obj.genome['chr1']) - snpseq[30] = 'A' if snpseq[30] == 'T' else 'T' - snpseq = ''.join(snpseq) - sam_obj.add_read('chr1', Sequence(snpseq, 0, _cigar_str='30M1X29M')) + sam_obj.add_contig("chr1", length=60) + snpseq_list = list(sam_obj.genome["chr1"]) + snpseq_list[30] = "A" if snpseq_list[30] == "T" else "T" + sam_obj.add_read("chr1", Sequence( + "".join(snpseq_list), + 0, + _cigar_str="30M1X29M", + )) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -68,25 +81,29 @@ def test_ref_seq_snp(self): with AlignmentFile(self.bam_fname) as af: read = next(af.fetch()) self.assertEqual( - ''.join(fa_ref_fetch.get_refseq(read)), - sam_obj.genome['chr1'], + "".join(fa_ref_fetch.get_refseq(read)), + sam_obj.genome["chr1"], ) self.assertEqual( - ''.join(md_ref_fetch.get_refseq(read)), - sam_obj.genome['chr1'], + "".join(md_ref_fetch.get_refseq(read)), + sam_obj.genome["chr1"], ) - def test_se_strands(self): + def test_se_strands(self) -> None: + """Check ability to recognize single-stranded data. + + Checks for unstranded, forward, and reverse strand data. + """ sam_obj = SAM() - sam_obj.add_contig('chr1') - ref_seq = sam_obj.genome['chr1'] + sam_obj.add_contig("chr1") + ref_seq = sam_obj.genome["chr1"] sam_obj.add_read( - 'chr1', - Sequence(ref_seq, 0, flag=0, qname='read1'), + "chr1", + Sequence(ref_seq, 0, flag=0, read_name="read1"), ) sam_obj.add_read( - 'chr1', - Sequence(ref_seq[1:], 1, flag=16, qname='read2'), + "chr1", + Sequence(ref_seq[1:], 1, flag=16, read_name="read2"), ) sam_obj.genome.save_to_fasta(self.fasta_fname) @@ -113,12 +130,16 @@ def test_se_strands(self): [False, True], ) - def test_pe_strands(self): + def test_pe_strands(self) -> None: + """Check ability to recognize paired-end data. + + Checks for unstranded, forward, and reverse strand data. + """ sam_obj = SAM() - sam_obj.add_contig('chr1') - ref_seq = sam_obj.genome['chr1'] - sam_obj.add_read_pair('chr1', Sequence(ref_seq, 0, flag=99)) - sam_obj.add_read_pair('chr1', Sequence(ref_seq[1:], 1, flag=83)) + sam_obj.add_contig("chr1") + ref_seq = sam_obj.genome["chr1"] + sam_obj.add_read_pair("chr1", Sequence(ref_seq, 0, flag=99)) + sam_obj.add_read_pair("chr1", Sequence(ref_seq[1:], 1, flag=83)) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -135,7 +156,7 @@ def test_pe_strands(self): self.assertEqual( [cr.get_strand(_) for _ in reads], [True, True, False, False], - ) + ) cr = CompiledReads(strand=2) self.assertEqual( @@ -143,10 +164,11 @@ def test_pe_strands(self): [False, False, True, True], ) - def test_trim(self): + def test_trim(self) -> None: + """Check read trimming.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=20) - sam_obj.add_read('chr1', Sequence(sam_obj.genome['chr1'], 0)) + sam_obj.add_contig("chr1", length=20) + sam_obj.add_read("chr1", Sequence(sam_obj.genome["chr1"], 0)) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -157,27 +179,37 @@ def test_trim(self): self.assertEqual(min(cr._nucleotides.keys()), 5) self.assertEqual(max(cr._nucleotides.keys()), 15) - def test_base_quality(self): + def test_base_quality(self) -> None: + """Check base quality filters.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=20) - read = Sequence(sam_obj.genome['chr1'], 0, phred=range(20)) - sam_obj.add_read('chr1', read) + sam_obj.add_contig("chr1", length=20) + sam_obj.add_read("chr1", Sequence( + sam_obj.genome["chr1"], + 0, + phred_list=list(range(20)), + )) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) with AlignmentFile(self.bam_fname) as af: - read = next(af.fetch()) - cr = CompiledReads(min_base_quality=10) - for _, _, phred, _ in cr._prep_read(read): - self.assertTrue(phred >= 10) - cr.add_reads([read]) + algn_seg = next(af.fetch()) + mbq = 10 + cr = CompiledReads(min_base_quality=mbq) + for _, _, phred, _ in cr._prep_read(algn_seg): + self.assertTrue(phred >= mbq) + cr.add_reads([algn_seg]) self.assertEqual(len(cr._nucleotides), 10) - def test_pop_range(self): + def test_pop_range(self) -> None: + """Check pop_range function.""" sam_obj = SAM() - sam_obj.add_contig('chr1', length=20) - read = Sequence(sam_obj.genome['chr1'], 0, phred=range(20)) - sam_obj.add_read('chr1', read) + sam_obj.add_contig("chr1", length=20) + read = Sequence( + sam_obj.genome["chr1"], + 0, + phred_list=list(range(20)), + ) + sam_obj.add_read("chr1", read) sam_obj.genome.save_to_fasta(self.fasta_fname) sam_obj.save_to_sam(self.bam_fname, self.fasta_fname) @@ -187,4 +219,4 @@ def test_pop_range(self): cp_list = list(cr.pop_range(18, 30)) self.assertEqual(len(cp_list), 2) - self.assertEqual(cp_list[0].ref, sam_obj.genome['chr1'][18]) + self.assertEqual(cp_list[0].ref, sam_obj.genome["chr1"][18]) diff --git a/test/fasta_file.py b/test/fasta_file.py index 18e243c..44d7b07 100644 --- a/test/fasta_file.py +++ b/test/fasta_file.py @@ -1,79 +1,123 @@ -import os +"""Test cases for RTFastaFaile class.""" +from __future__ import annotations + import random import unittest from itertools import chain +from pathlib import Path from tempfile import NamedTemporaryFile +from test.sam_gen import Genome -from reditools.fasta_file import RTFastaFile +from reditools.constants import bases +from reditools.fasta_file import ( + MissingContigError, + PastContigEndError, + RTFastaFile, +) class TestRTFastaFile(unittest.TestCase): - def setUp(self): - self.contig1 = 'test1' - self.seq1 = self.random_seq(80) - self.contig2 = 'chrtest2' - self.seq2 = self.random_seq(80) + """Test cases for RTFastaFaile class.""" + + def setUp(self) -> None: + """Pre-flight setup.""" + self.genome = Genome() + self.naked_contig_name = "test1" + self.genome.add_contig(self.naked_contig_name, 80) + self.chr_contig_name = "chrtest2" + self.genome.add_contig(self.chr_contig_name, 80) with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8', + mode="w", + encoding="utf-8", ) as stream: self.fasta_fname = stream.name - stream.write(f'>{self.contig1}\n{self.seq1}\n') - stream.write(f'>{self.contig2}\n{self.seq2}\n') + self.genome.save_to_fasta(self.fasta_fname) - def tearDown(self): - os.remove(self.fasta_fname) + def tearDown(self) -> None: + """Post-check cleanup.""" + Path(self.fasta_fname).unlink() - def test_get_base(self): + def test_get_base(self) -> None: + """Check get_base() function.""" + refseq = self.genome[self.naked_contig_name] with RTFastaFile(self.fasta_fname) as rff: - positions = list(range(len(self.seq1))) - fasta_seq = rff.get_base(self.contig1, *positions) - self.assertEqual(self.seq1, ''.join(fasta_seq)) + fasta_seq = rff.get_base( + self.naked_contig_name, + *range(len(refseq)), + ) + self.assertEqual(refseq, "".join(fasta_seq)) - def test_get_base_splice(self): + def test_get_base_splice(self) -> None: + """Check get_base() function with spliced reads.""" + refseq = self.genome[self.naked_contig_name] with RTFastaFile(self.fasta_fname) as rff: - positions = list(chain( + positions = chain( range(20), - range(len(self.seq1) - 20, len(self.seq1)), - )) - fasta_seq = rff.get_base(self.contig1, *positions) + range(len(refseq) - 20, len(refseq)), + ) + fasta_seq = rff.get_base(self.naked_contig_name, *positions) self.assertEqual( - self.seq1[:20] + self.seq1[-20:], - ''.join(fasta_seq), + refseq[:20] + refseq[-20:], + "".join(fasta_seq), ) - def test_get_base_prefix(self): + def test_get_base_prefix(self) -> None: + """Check get_base() function with contig name variants. + + Specifically, get_base should work regardless of whether a chromosome + starts with "chr". + """ + refseq = self.genome[self.chr_contig_name] with RTFastaFile(self.fasta_fname) as rff: - positions = range(len(self.seq2)) - fasta_seq = rff.get_base('test2', *positions) - self.assertEqual(self.seq2, ''.join(fasta_seq)) + fasta_seq = rff.get_base( + self.chr_contig_name, + *range(len(refseq)), + ) + self.assertEqual(refseq, "".join(fasta_seq)) - fasta_seq = rff.get_base('chrtest2', *positions) - self.assertEqual(self.seq2, ''.join(fasta_seq)) + fasta_seq = rff.get_base( + f"chr{self.chr_contig_name}", + *range(len(refseq)), + ) + self.assertEqual(refseq, "".join(fasta_seq)) - def test_get_base_missing_contig(self): + def test_get_base_missing_contig(self) -> None: + """Check errors when accessing non-existent chromosomes.""" with RTFastaFile(self.fasta_fname) as rff: - with self.assertRaises(KeyError): - rff.get_base('test3', 0) - with self.assertRaises(KeyError): - rff.get_base('chrtest3', 0) + with self.assertRaises(MissingContigError): + list(rff.get_base("test3", 0)) + with self.assertRaises(MissingContigError): + list(rff.get_base("chrtest3", 0)) + + def test_get_base_out_of_bounds(self) -> None: + """Check errors when accessing bases outside of chromosome boundns.""" + refseq = self.genome[self.naked_contig_name] + with RTFastaFile(self.fasta_fname) as rff, \ + self.assertRaises(PastContigEndError): + positions = range( + len(refseq) - 20, + len(refseq) + 20, + ) + seq_iter = rff.get_base( + self.naked_contig_name, + *positions, + ) + list(seq_iter) - def test_get_base_out_of_bounds(self): - with RTFastaFile(self.fasta_fname) as rff: - with self.assertRaises(IndexError): - start = len(self.seq1) - 20 - stop = len(self.seq1) + 20 - positions = range(start, stop) - seq_iter = rff.get_base( - self.contig1, - *positions, - ) - list(seq_iter) @classmethod - def random_seq(cls, length): - sequence = [] - for _ in range(length): - sequence.append(random.choice('ACTG')) - return ''.join(sequence) + def random_seq(cls, length: int) -> str: + """Generate a random nucleotide sequence. + + Parameters + ---------- + length : int + Sequence length. + Returns + ------- + str + Random sequence. + """ + sequence = [random.choice(bases) for _ in range(length)] + return "".join(sequence) diff --git a/test/file_utils.py b/test/file_utils.py index 4ed7e2f..93c3d46 100644 --- a/test/file_utils.py +++ b/test/file_utils.py @@ -1,123 +1,161 @@ +"""Test cases for file_utils.""" +from __future__ import annotations + import gzip -import os import unittest +from pathlib import Path from tempfile import NamedTemporaryFile +from typing import Iterable from reditools import file_utils from reditools.region import Region class TestFileUtils(unittest.TestCase): - def write_file(self, data_list, sep=' '): + """Test cases for file_utils.""" + + def write_file(self, data_list: str | Iterable, sep: str=" ") -> str: + """Write data to file. + + Parameters + ---------- + data_list : str | Iterable + Data to write. If data_list is not a string, it will be joined + using sep. + sep : str + Field seperator. Only used if data_list is not a string. + + Returns + ------- + str + Filename data saved to. + """ with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8', + mode="w", + encoding="utf-8", ) as stream: for row in data_list: if isinstance(row, str): stream.write(row) else: stream.write(sep.join([str(_) for _ in row])) - stream.write('\n') + stream.write("\n") return stream.name - def check_test_data(self, test_data, real_data): + def check_test_data( + self, + test_data: Iterable[tuple], + real_data: list, + ) -> None: + """Check test data matches real data. + + Parameters + ---------- + test_data : Iterable[tuple] + Assumes second element of each tuple what the output should be. + real_data : list + What the output actually was. + """ self.assertEqual([_[1] for _ in test_data], real_data) - def test_open_stream_plain(self): - test_str = 'test123' + def test_open_stream_plain(self) -> None: + """Check read/write plain text files.""" + test_str = "test123" with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8', + mode="w", + encoding="utf-8", ) as stream: stream.write(test_str) fname = stream.name - with file_utils.open_stream(fname, 'rt') as stream: + with file_utils.open_stream(fname, "rt") as stream: file_content = stream.read() self.assertEqual(file_content, test_str) - os.remove(fname) + Path(fname).unlink() - def test_open_stream_gzip(self): - test_str = 'test_gzip' + def test_open_stream_gzip(self) -> None: + """Check read/write gzipped files.""" + test_str = "test_gzip" with NamedTemporaryFile( delete=False, - suffix='.gz', - mode='wb', + suffix=".gz", + mode="wb", ) as stream: - stream.write(gzip.compress(bytes(test_str, 'utf-8'))) + stream.write(gzip.compress(bytes(test_str, "utf-8"))) fname = stream.name - with file_utils.open_stream(fname, 'rt') as stream: + with file_utils.open_stream(fname, "rt") as stream: file_content = stream.read() self.assertEqual(file_content, test_str) - os.remove(fname) + Path(fname).unlink() - def test_read_bed_file(self): + def test_read_bed_file(self) -> None: + """Check read BED files.""" bed_data = ( ( - ('chr1', 10, 20), - Region('chr1', 10, 20), + ("chr1", 10, 20), + Region("chr1", 10, 20), ), ( - ('chr1', 30, 40), - Region('chr1', 30, 40), + ("chr1", 30, 40), + Region("chr1", 30, 40), ), ) - fname = self.write_file((_[0] for _ in bed_data), sep='\t') + fname = self.write_file((_[0] for _ in bed_data), sep="\t") region_list = list(file_utils.read_bed_file(fname)) self.check_test_data(bed_data, region_list) - os.remove(fname) + Path(fname).unlink() - def test_read_many_bed_files(self): + def test_read_many_bed_files(self) -> None: + """Check read multiple BED files.""" bed_data = ( ( - ('chr1', 10, 20), - Region('chr1', 10, 20), + ("chr1", 10, 20), + Region("chr1", 10, 20), ), ( - ('chr1', 30, 40), - Region('chr1', 30, 40), + ("chr1", 30, 40), + Region("chr1", 30, 40), ), ) - fnames = [] - for row in bed_data: - fnames.append(self.write_file([row[0]], sep='\t')) + fnames = [self.write_file([row[0]], sep="\t") for row in bed_data] region_list = list(file_utils.read_bed_file(*fnames)) self.check_test_data(bed_data, sorted(region_list)) for fname in fnames: - os.remove(fname) + Path(fname).unlink() - def test_concat(self): - file_contents = ('file1', 'file2', 'file3') + def test_concat(self) -> None: + """Check file concatenation.""" + file_contents = ("file1", "file2", "file3") file_names = [self.write_file([_]) for _ in file_contents] with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8') as stream: - file_utils.concat(stream, *file_names, encoding='utf-8') + mode="w", + encoding="utf-8") as stream: + file_utils.concat(stream, *file_names, encoding="utf-8") concat_filename = stream.name for fname in file_names: - self.assertFalse(os.path.exists(fname)) + self.assertFalse(Path(fname).exists()) - with open(concat_filename, 'r') as stream: + with Path(concat_filename).open("r") as stream: self.assertEqual( stream.read(), - ''.join([f'{_}\n' for _ in file_contents]), + "".join([f"{_}\n" for _ in file_contents]), ) - os.remove(concat_filename) + Path(concat_filename).unlink() - def test_load_text_file(self): + def test_load_text_file(self) -> None: + """Check read plaintext files.""" text_lines = ["rowA", "rowB", "rowC"] with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8', + mode="w", + encoding="utf-8", ) as stream: fname = stream.name - stream.write('\n'.join(text_lines)) + stream.write("\n".join(text_lines)) loaded_text = file_utils.load_text_file(fname) self.assertEqual(loaded_text, text_lines) - os.remove(fname) + Path(fname).unlink() diff --git a/test/reditools.py b/test/reditools.py index 5709304..51a5ee6 100644 --- a/test/reditools.py +++ b/test/reditools.py @@ -1,25 +1,30 @@ -import os +"""Test cases for REDItools class.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, Genome, Sequence, ntf from reditools import reditools from reditools.alignment_manager import AlignmentManager from reditools.compiled_position import CompiledPosition +from reditools.constants import comp_map from reditools.region import Region class TestREDItools(unittest.TestCase): - complement = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C'} + """Test cases for REDItools class.""" - def setUp(self): + def setUp(self) -> None: + """Pre-flight setup.""" self.rtools = reditools.REDItools() - self.bam_file = ntf(suffix='.bam') - self.fa_file = ntf(suffix='.fa') + self.bam_file = ntf(suffix=".bam") + self.fa_file = ntf(suffix=".fa") self.sam_obj = SAM() - self.sam_obj.add_contig('chr1', length=10) - self.sam_obj.add_read('chr1', Sequence(self.sam_obj.genome['chr1'], 0)) + self.sam_obj.add_contig("chr1", length=10) + self.sam_obj.add_read("chr1", Sequence(self.sam_obj.genome["chr1"], 0)) self.sam_obj.genome.save_to_fasta(self.fa_file) self.sam_obj.save_to_sam(self.bam_file, self.fa_file) @@ -27,55 +32,60 @@ def setUp(self): self.rtam = AlignmentManager() self.rtam.add_file(self.bam_file) - self.cp = CompiledPosition(ref='A', position=1, contig='chr1') - self.cp.add_base(30, '-', 'A') - self.cp.add_base(30, '-', 'A') - self.cp.add_base(30, '-', 'A') - self.cp.add_base(30, '-', 'T') - self.cp.add_base(30, '+', 'G') - self.cp.add_base(30, '+', 'G') - - def tearDown(self): - os.remove(self.bam_file) - os.remove(self.fa_file) - - def test_process_bases(self): + self.cp = CompiledPosition(ref="A", position=1, contig="chr1") + self.cp.add_base(30, "-", "A") + self.cp.add_base(30, "-", "A") + self.cp.add_base(30, "-", "A") + self.cp.add_base(30, "-", "T") + self.cp.add_base(30, "+", "G") + self.cp.add_base(30, "+", "G") + + def tearDown(self) -> None: + """Post-checks cleanup.""" + Path(self.bam_file).unlink() + Path(self.fa_file).unlink() + + def test_process_bases(self) -> None: + """Check _process_bases() function.""" rtresult = self.rtools._process_bases(self.cp) - self.assertEqual(rtresult.reference, 'A') - self.assertEqual(rtresult.strand, '*') - self.assertEqual(rtresult.variants, ['AG', 'AT']) + self.assertEqual(rtresult.reference, "A") + self.assertEqual(rtresult.strand, "*") + self.assertEqual(rtresult.variants, ["AG", "AT"]) - def test_strand_filter(self): + def test_strand_filter(self) -> None: + """Check strand modes.""" self.rtools.strand = reditools.FORWARD_STRAND_MODE self.rtools.strand_confidence_threshold = 0.5 rtresult = self.rtools._process_bases(self.cp) - self.assertEqual(rtresult.strand, '-') - self.assertEqual(rtresult.reference, 'A') - self.assertEqual(rtresult.variants, ['AT']) + self.assertEqual(rtresult.strand, "-") + self.assertEqual(rtresult.reference, "A") + self.assertEqual(rtresult.variants, ["AT"]) - def test_strand_correction(self): + def test_strand_correction(self) -> None: + """Check strand correction mode.""" self.rtools.strand = reditools.FORWARD_STRAND_MODE self.rtools.strand_confidence_threshold = 0.5 self.rtools.use_strand_correction() rtresult = self.rtools._process_bases(self.cp) - self.assertEqual(rtresult.strand, '-') - self.assertEqual(rtresult.reference, 'T') - self.assertEqual(rtresult.variants, ['TA']) + self.assertEqual(rtresult.strand, "-") + self.assertEqual(rtresult.reference, "T") + self.assertEqual(rtresult.variants, ["TA"]) - def test_add_reference(self): + def test_add_reference(self) -> None: + """Check using MD tags and FASTA files as references.""" rtresult = next( self.rtools.analyze( self.rtam, - Region.from_string('chr1', self.bam_file), + Region.from_string("chr1", self.bam_file), ), ) - self.assertEqual(rtresult.reference, self.sam_obj.genome['chr1'][0]) + self.assertEqual(rtresult.reference, self.sam_obj.genome["chr1"][0]) new_genome = Genome() new_genome.add_contig( - 'chr1', - sequence=''.join( - self.complement[_] for _ in self.sam_obj.genome['chr1'] + "chr1", + sequence="".join( + comp_map[_] for _ in self.sam_obj.genome["chr1"] ), ) new_genome.save_to_fasta(self.fa_file) @@ -83,20 +93,21 @@ def test_add_reference(self): rtresult = next( self.rtools.analyze( self.rtam, - Region.from_string('chr1', self.bam_file), + Region.from_string("chr1", self.bam_file), ), ) - self.assertEqual(rtresult.reference, self.sam_obj.genome['chr1'][0]) + self.assertEqual(rtresult.reference, self.sam_obj.genome["chr1"][0]) - def test_region(self): + def test_analyze(self) -> None: + """Check the analyze function.""" rtresults = list(self.rtools.analyze( self.rtam, - Region.from_string('chr1:3-7', self.bam_file), + Region.from_string("chr1:3-7", self.bam_file), )) self.assertEqual(len(rtresults), 5) rtresults = list(self.rtools.analyze( self.rtam, - Region.from_string('chr1:8-20', self.bam_file), + Region.from_string("chr1:8-20", self.bam_file), )) self.assertEqual(len(rtresults), 3) diff --git a/test/region.py b/test/region.py index ee8eb53..0ed65d9 100644 --- a/test/region.py +++ b/test/region.py @@ -1,91 +1,105 @@ -import os +"""Test cases for Region class.""" +from __future__ import annotations + import unittest +from pathlib import Path from test.sam_gen import SAM, ntf from reditools.region import Region class TestRegion(unittest.TestCase): + """Test cases for Region class.""" - def test_str(self): - self.assertEqual(str(Region('chr1', 100, 200)), 'chr1:101-200') - self.assertEqual(str(Region('chr1', 0, None)), 'chr1') - self.assertEqual(str(Region('chr1', 50, None)), 'chr1:51') + def test_str(self) -> None: + """Check cast to string.""" + self.assertEqual(str(Region("chr1", 100, 200)), "chr1:101-200") - def test_even_split(self): - region = Region('chr1', 0, 1000) + def test_even_split(self) -> None: + """Check split() when window sizes are a perfect fit.""" + region = Region("chr1", 0, 1000) windows = region.split(250) self.assertEqual(len(windows), 4) - self.assertEqual(windows[0], Region('chr1', 0, 250)) - self.assertEqual(windows[-1], Region('chr1', 750, 1000)) - - def test_uneven_split(self): - region = Region('chr1', 0, 950) + self.assertEqual(windows[0], Region("chr1", 0, 250)) + self.assertEqual(windows[1], Region("chr1", 250, 500)) + self.assertEqual(windows[2], Region("chr1", 500, 750)) + self.assertEqual(windows[3], Region("chr1", 750, 1000)) + + def test_uneven_split(self) -> None: + """Check split() when window sizes have a remainder.""" + region = Region("chr1", 0, 950) windows = region.split(300) self.assertEqual(len(windows), 4) - self.assertEqual(windows[0], Region('chr1', 0, 300)) - self.assertEqual(windows[1], Region('chr1', 300, 600)) - self.assertEqual(windows[2], Region('chr1', 600, 900)) - self.assertEqual(windows[3], Region('chr1', 900, 950)) - - def test_impossible_split(self): - region = Region('chr1', 0, 100) + self.assertEqual(windows[0], Region("chr1", 0, 300)) + self.assertEqual(windows[1], Region("chr1", 300, 600)) + self.assertEqual(windows[2], Region("chr1", 600, 900)) + self.assertEqual(windows[3], Region("chr1", 900, 950)) + + def test_impossible_split(self) -> None: + """Check split() when window is bigger than the Region.""" + region = Region("chr1", 0, 100) windows = region.split(200) self.assertEqual(len(windows), 1) - self.assertEqual(windows[0], Region('chr1', 0, 100)) + self.assertEqual(windows[0], Region("chr1", 0, 100)) - def test_nonzero_split(self): - region = Region('chr2', 5, 122) + def test_nonzero_split(self) -> None: + """Check split() when Region does not start at zero.""" + region = Region("chr2", 5, 122) windows = region.split(50) self.assertEqual(len(windows), 3) - self.assertEqual(windows[0], Region('chr2', 5, 55)) - self.assertEqual(windows[1], Region('chr2', 55, 105)) - self.assertEqual(windows[2], Region('chr2', 105, 122)) + self.assertEqual(windows[0], Region("chr2", 5, 55)) + self.assertEqual(windows[1], Region("chr2", 55, 105)) + self.assertEqual(windows[2], Region("chr2", 105, 122)) - def test_none_split(self): + def test_none_split(self) -> None: + """Check split() when Region bounds are undefined.""" with self.assertRaises(IndexError): - Region('chr1', None, 100).split(50) + Region("chr1", None, 100).split(50) # type: ignore[arg-type] with self.assertRaises(IndexError): - Region('chr1', 50, None).split(50) + Region("chr1", 50, None).split(50) # type: ignore[arg-type] - def test_from_string(self): - fasta_fname = ntf(suffix='.fa') - bam_fname = ntf(suffix='.bam') + def test_from_string(self) -> None: + """Check from_string() method.""" + fasta_fname = ntf(suffix=".fa") + bam_fname = ntf(suffix=".bam") sam_obj = SAM() chr1_len = 600 - sam_obj.add_contig('chr1', length=chr1_len) + sam_obj.add_contig("chr1", length=chr1_len) sam_obj.genome.save_to_fasta(fasta_fname) sam_obj.save_to_sam(bam_fname, fasta_fname) - region = Region.from_string('chr1:101-200', bam_fname) - self.assertEqual(region, Region('chr1', 100, 200)) - region = Region.from_string('chr1:104', bam_fname) - self.assertEqual(region, Region('chr1', 103, chr1_len)) - region = Region.from_string('chr1', bam_fname) - self.assertEqual(region, Region('chr1', 0, chr1_len)) + region = Region.from_string("chr1:101-200", bam_fname) + self.assertEqual(region, Region("chr1", 100, 200)) + region = Region.from_string("chr1:104", bam_fname) + self.assertEqual(region, Region("chr1", 103, chr1_len)) + region = Region.from_string("chr1", bam_fname) + self.assertEqual(region, Region("chr1", 0, chr1_len)) - os.remove(fasta_fname) - os.remove(bam_fname) + Path(fasta_fname).unlink() + Path(bam_fname).unlink() - def test_parse_string(self): - region = Region.parse_string('chr1:101-200') - self.assertEqual(region, ('chr1', 100, 200)) + def test_parse_string(self) -> None: + """Check parse_string() method.""" + region = Region.parse_string("chr1:101-200") + self.assertEqual(region, ("chr1", 100, 200)) with self.assertRaises(ValueError): - Region.parse_string('chr1:-2') + Region.parse_string("chr1:-2") - def test_to_int(self): + def test_to_int(self) -> None: + """Check _to_int() method.""" self.assertEqual(Region._to_int("10"), 10) self.assertEqual(Region._to_int("10,000"), 10000) with self.assertRaises(ValueError): Region._to_int("X") - def test_order(self): + def test_order(self) -> None: + """Check stortability.""" regions_list = [ - Region('chr1', 20, 30), - Region('chr1', 10, 30), - Region('chr1', 10, 20), + Region("chr1", 20, 30), + Region("chr1", 10, 30), + Region("chr1", 10, 20), ] self.assertEqual(sorted(regions_list), list(reversed(regions_list))) diff --git a/test/region_collection.py b/test/region_collection.py index 1ee85cf..2d5e4a3 100644 --- a/test/region_collection.py +++ b/test/region_collection.py @@ -1,3 +1,6 @@ +"""Test cases for RegonCollection class.""" +from __future__ import annotations + import unittest from reditools.region import Region @@ -5,32 +8,35 @@ class TestRegionCollection(unittest.TestCase): + """Test cases for RegonCollection class.""" - def setUp(self): + def setUp(self) -> None: + """Pre-flight setup.""" self.rc = RegionCollection() self.rc.add_regions([ - Region('chr1', 0, 99), - Region('chr1', 100, 199), - Region('chr2', 50, 150), + Region("chr1", 0, 99), + Region("chr1", 100, 199), + Region("chr2", 50, 150), ]) - def test_add_region_and_contains(self): - # RegionCollection contains method requires ordered queries. - self.assertTrue(self.rc.contains('chr1', 50)) - self.assertTrue(self.rc.contains('chr1', 150)) - self.assertFalse(self.rc.contains('chr1', 200)) - self.assertTrue(self.rc.contains('chr2', 100)) - self.assertFalse(self.rc.contains('chr2', 200)) - self.assertFalse(self.rc.contains('chrX', 1)) + def test_add_region_and_contains(self) -> None: + """Check contains() method.""" + self.assertTrue(self.rc.contains("chr1", 50)) + self.assertTrue(self.rc.contains("chr1", 150)) + self.assertFalse(self.rc.contains("chr1", 200)) + self.assertTrue(self.rc.contains("chr2", 100)) + self.assertFalse(self.rc.contains("chr2", 200)) + self.assertFalse(self.rc.contains("chrX", 1)) - def test_add_regions(self): + def test_add_regions(self) -> None: + """Check add_regions() method.""" regions = [ - Region('chr3', 0, 10), - Region('chr3', 11, 20), - Region('chr1', 200, 299), + Region("chr3", 0, 10), + Region("chr3", 11, 20), + Region("chr1", 200, 299), ] self.rc.add_regions(regions) - self.assertTrue(self.rc.contains('chr3', 5)) - self.assertTrue(self.rc.contains('chr3', 15)) - self.assertFalse(self.rc.contains('chr3', 21)) - self.assertTrue(self.rc.contains('chr1', 250)) + self.assertTrue(self.rc.contains("chr3", 5)) + self.assertTrue(self.rc.contains("chr3", 15)) + self.assertFalse(self.rc.contains("chr3", 21)) + self.assertTrue(self.rc.contains("chr1", 250)) diff --git a/test/rtannotater.py b/test/rtannotater.py index e9ac4fb..9b2802a 100644 --- a/test/rtannotater.py +++ b/test/rtannotater.py @@ -1,343 +1,356 @@ -import os +"""Test cases for RTAnnotater class.""" +from __future__ import annotations + import unittest +from pathlib import Path from tempfile import NamedTemporaryFile -from reditools.rtannotater import RTAnnotater +from reditools.rtannotater import AnalyzeMismatchError, RTAnnotater class TestRTAnnotater(unittest.TestCase): - def test_legacy_translate(self): + """Test cases for RTAnnotater class.""" + + def test_legacy_translate(self) -> None: + """Check legacy_translate() method.""" test_dict = { - 'Coverage-q30': '100', - 'gCoverage-q30': '200', - 'AnotherField': '123', + "Coverage-q30": "100", + "gCoverage-q30": "200", + "AnotherField": "123", } RTAnnotater.legacy_translate(test_dict) self.assertEqual(test_dict, { - 'Coverage': '100', - 'gCoverage': '200', - 'AnotherField': '123', + "Coverage": "100", + "gCoverage": "200", + "AnotherField": "123", }) - - def test_cmp_position(self): + + def test_cmp_position(self) -> None: + """Check cmp_position() method.""" contig_order = { - 'chrZ': 1, - 'chr1': 2, - 'chr2': 3, + "chrZ": 1, + "chr1": 2, + "chr2": 3, } rta = RTAnnotater(contig_order) self.assertEqual( rta.cmp_position( - {'Region': 'chr1', 'Position': '5'}, - {'Region': 'chr1', 'Position': '5'}, + {"Region": "chr1", "Position": "5"}, + {"Region": "chr1", "Position": "5"}, ), 0, ) self.assertTrue( rta.cmp_position( - {'Region': 'chr1', 'Position': '5'}, + {"Region": "chr1", "Position": "5"}, None, ) < 0, ) self.assertTrue( rta.cmp_position( - {'Region': 'chrZ', 'Position': '1'}, - {'Region': 'chr2', 'Position': '1'}, + {"Region": "chrZ", "Position": "1"}, + {"Region": "chr2", "Position": "1"}, ) < 0, ) self.assertTrue( rta.cmp_position( - {'Region': 'chr1', 'Position': '10'}, - {'Region': 'chr1', 'Position': '1'}, + {"Region": "chr1", "Position": "10"}, + {"Region": "chr1", "Position": "1"}, ) > 0, ) - def test_annotate_row(self): + def test_annotate_row(self) -> None: + """Check annotate_row() method.""" rta = RTAnnotater({}) self.assertEqual( rta.annotate_row( { - 'Reference': 'A', - 'Coverage': '-', - 'MeanQ': '-', - 'BaseCount[A,C,G,T]': '[0, 0, 0, 0]', - 'AllSubs': '-', - 'Frequency': '-', - 'AnotherField': 'ABCD', + "Reference": "A", + "Coverage": "-", + "MeanQ": "-", + "BaseCount[A,C,G,T]": "[0, 0, 0, 0]", + "AllSubs": "-", + "Frequency": "-", + "AnotherField": "ABCD", }, { - 'Reference': 'A', - 'Coverage': '100', - 'MeanQ': '40', - 'BaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'AllSubs': 'AC AG AT', - 'Frequency': '0.5', - 'AnotherField': 'EFGH', - 'YetAnotherField': 'IJKL', + "Reference": "A", + "Coverage": "100", + "MeanQ": "40", + "BaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "AllSubs": "AC AG AT", + "Frequency": "0.5", + "AnotherField": "EFGH", + "YetAnotherField": "IJKL", }, ), { - 'Reference': 'A', - 'Coverage': '-', - 'MeanQ': '-', - 'BaseCount[A,C,G,T]': '[0, 0, 0, 0]', - 'AllSubs': '-', - 'Frequency': '-', - 'gCoverage': '100', - 'gMeanQ': '40', - 'gBaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'gAllSubs': 'AC AG AT', - 'gFrequency': '0.5', - 'AnotherField': 'ABCD', + "Reference": "A", + "Coverage": "-", + "MeanQ": "-", + "BaseCount[A,C,G,T]": "[0, 0, 0, 0]", + "AllSubs": "-", + "Frequency": "-", + "gCoverage": "100", + "gMeanQ": "40", + "gBaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "gAllSubs": "AC AG AT", + "gFrequency": "0.5", + "AnotherField": "ABCD", }, ) - def test_annotate_complement_row_no_dna_edit(self): - rta = RTAnnotater({}, True) + def test_annotate_complement_row_no_dna_edit(self) -> None: + """Check annotate_row() on minus strand with unedited DNA.""" + rta = RTAnnotater({}, True) # noqa: FBT003 self.assertEqual( rta.annotate_row( { - 'Reference': 'T', - 'Coverage': '100', - 'MeanQ': '40', - 'BaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'AllSubs': 'AC AG AT', - 'Frequency': '0.5', - 'AnotherField': 'EFGH', - 'YetAnotherField': 'IJKL', + "Reference": "T", + "Coverage": "100", + "MeanQ": "40", + "BaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "AllSubs": "AC AG AT", + "Frequency": "0.5", + "AnotherField": "EFGH", + "YetAnotherField": "IJKL", }, { - 'Reference': 'A', - 'Coverage': '-', - 'MeanQ': '-', - 'BaseCount[A,C,G,T]': '[0, 0, 0, 0]', - 'AllSubs': '-', - 'Frequency': '-', - 'AnotherField': 'ABCD', + "Reference": "A", + "Coverage": "-", + "MeanQ": "-", + "BaseCount[A,C,G,T]": "[0, 0, 0, 0]", + "AllSubs": "-", + "Frequency": "-", + "AnotherField": "ABCD", }, ), { - 'Reference': 'T', - 'Coverage': '100', - 'MeanQ': '40', - 'BaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'AllSubs': 'AC AG AT', - 'Frequency': '0.5', - 'gCoverage': '-', - 'gMeanQ': '-', - 'gBaseCount[A,C,G,T]': '[0, 0, 0, 0]', - 'gAllSubs': '-', - 'gFrequency': '-', - 'AnotherField': 'EFGH', - 'YetAnotherField': 'IJKL', + "Reference": "T", + "Coverage": "100", + "MeanQ": "40", + "BaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "AllSubs": "AC AG AT", + "Frequency": "0.5", + "gCoverage": "-", + "gMeanQ": "-", + "gBaseCount[A,C,G,T]": "[0, 0, 0, 0]", + "gAllSubs": "-", + "gFrequency": "-", + "AnotherField": "EFGH", + "YetAnotherField": "IJKL", }, ) - def test_annotate_complement_row(self): - rta = RTAnnotater({}, True) + def test_annotate_complement_row(self) -> None: + """Check annotate_row() on minus strand.""" + rta = RTAnnotater({}, True) # noqa: FBT003 self.assertEqual( rta.annotate_row( { - 'Reference': 'A', - 'gCoverage': '-', - 'gMeanQ': '-', - 'gBaseCount[A,C,G,T]': '-', - 'gAllSubs': '-', - 'gFrequency': '-', - 'AnotherField': 'ABCD', + "Reference": "A", + "gCoverage": "-", + "gMeanQ": "-", + "gBaseCount[A,C,G,T]": "-", + "gAllSubs": "-", + "gFrequency": "-", + "AnotherField": "ABCD", }, { - 'Reference': 'T', - 'Coverage': '100', - 'MeanQ': '40', - 'BaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'AllSubs': 'AC AG AT', - 'Frequency': '0.5', - 'AnotherField': 'EFGH', - 'YetAnotherField': 'IJKL', + "Reference": "T", + "Coverage": "100", + "MeanQ": "40", + "BaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "AllSubs": "AC AG AT", + "Frequency": "0.5", + "AnotherField": "EFGH", + "YetAnotherField": "IJKL", }, ), { - 'Reference': 'A', - 'gCoverage': '100', - 'gMeanQ': '40', - 'gBaseCount[A,C,G,T]': '[4, 3, 2, 1]', - 'gAllSubs': 'TA TC TG', - 'gFrequency': '0.5', - 'AnotherField': 'ABCD', + "Reference": "A", + "gCoverage": "100", + "gMeanQ": "40", + "gBaseCount[A,C,G,T]": "[4, 3, 2, 1]", + "gAllSubs": "TA TC TG", + "gFrequency": "0.5", + "AnotherField": "ABCD", }, ) - def test_annotate_no_complement_row(self): + def test_annotate_no_complement_row(self) -> None: + """Check annotate_row() on minus strand without complementing.""" rta = RTAnnotater({}) self.assertEqual( rta.annotate_row( { - 'Reference': 'A', - 'gCoverage': '-', - 'gMeanQ': '-', - 'gBaseCount[A,C,G,T]': '-', - 'gAllSubs': '-', - 'gFrequency': '-', - 'AnotherField': 'ABCD', + "Reference": "A", + "gCoverage": "-", + "gMeanQ": "-", + "gBaseCount[A,C,G,T]": "-", + "gAllSubs": "-", + "gFrequency": "-", + "AnotherField": "ABCD", }, { - 'Reference': 'T', - 'Coverage': '100', - 'MeanQ': '40', - 'BaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'AllSubs': 'AC AG AT', - 'Frequency': '0.5', - 'AnotherField': 'EFGH', - 'YetAnotherField': 'IJKL', + "Reference": "T", + "Coverage": "100", + "MeanQ": "40", + "BaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "AllSubs": "AC AG AT", + "Frequency": "0.5", + "AnotherField": "EFGH", + "YetAnotherField": "IJKL", }, ), { - 'Reference': 'A', - 'gCoverage': '100', - 'gMeanQ': '40', - 'gBaseCount[A,C,G,T]': '[1, 2, 3, 4]', - 'gAllSubs': 'AC AG AT', - 'gFrequency': '0.5', - 'AnotherField': 'ABCD', + "Reference": "A", + "gCoverage": "100", + "gMeanQ": "40", + "gBaseCount[A,C,G,T]": "[1, 2, 3, 4]", + "gAllSubs": "AC AG AT", + "gFrequency": "0.5", + "AnotherField": "ABCD", }, ) - def test_mismatched_reference(self): + def test_mismatched_reference(self) -> None: + """Check error handling for reference mismatch.""" rta = RTAnnotater({}) - with self.assertRaises(ValueError): + with self.assertRaises(AnalyzeMismatchError): rta.annotate_row( - { 'Reference': 'A'}, - { 'Reference': 'G'}, + { "Reference": "A"}, + { "Reference": "G"}, ) - def test_merge_files(self): + def test_merge_files(self) -> None: + """Check merge_files() method.""" fieldnames = [ - 'Region', - 'Position', - 'Reference', - 'Strand', - 'Coverage', - 'MeanQ', - 'BaseCount[A,C,G,T]', - 'AllSubs', - 'Frequency', - 'gCoverage', - 'gMeanQ', - 'gBaseCount[A,C,G,T]', - 'gAllSubs', - 'gFrequency', + "Region", + "Position", + "Reference", + "Strand", + "Coverage", + "MeanQ", + "BaseCount[A,C,G,T]", + "AllSubs", + "Frequency", + "gCoverage", + "gMeanQ", + "gBaseCount[A,C,G,T]", + "gAllSubs", + "gFrequency", ] with NamedTemporaryFile( delete=False, - suffix='.out', - mode='w', + suffix=".out", + mode="w", ) as stream: rna_file = stream.name - stream.write('\t'.join(fieldnames)) - stream.write('\n') - stream.write('\t'.join([ - 'chr1', '3', 'A', '*', '5', '30', '[5, 0, 0, 0]', '', '0.0', - '-', '-', '-', '-', '-', + stream.write("\t".join(fieldnames)) + stream.write("\n") + stream.write("\t".join([ + "chr1", "3", "A", "*", "5", "30", "[5, 0, 0, 0]", "", "0.0", + "-", "-", "-", "-", "-", ])) - stream.write('\n') - stream.write('\t'.join([ - 'chr1', '5', 'A', '*', '5', '30', '[2, 0, 3, 0]', 'AG', '0.6', - '-', '-', '-', '-', '-', + stream.write("\n") + stream.write("\t".join([ + "chr1", "5", "A", "*", "5", "30", "[2, 0, 3, 0]", "AG", "0.6", + "-", "-", "-", "-", "-", ])) - stream.write('\n') - stream.write('\t'.join([ - 'chr2', '7', 'G', '*', '5', '30', '[2, 0, 3, 0]', 'GA', '0.4', - '-', '-', '-', '-', '-', + stream.write("\n") + stream.write("\t".join([ + "chr2", "7", "G", "*", "5", "30", "[2, 0, 3, 0]", "GA", "0.4", + "-", "-", "-", "-", "-", ])) - stream.write('\n') + stream.write("\n") with NamedTemporaryFile( delete=False, - suffix='.out', - mode='w', + suffix=".out", + mode="w", ) as stream: dna_file = stream.name - stream.write('\t'.join(fieldnames)) - stream.write('\n') - stream.write('\t'.join([ - 'chrZ', '3', 'A', '*', '5', '30', '[5, 0, 0, 0]', '', '0.0', - '-', '-', '-', '-', '-', + stream.write("\t".join(fieldnames)) + stream.write("\n") + stream.write("\t".join([ + "chrZ", "3", "A", "*", "5", "30", "[5, 0, 0, 0]", "", "0.0", + "-", "-", "-", "-", "-", ])) - stream.write('\n') - stream.write('\t'.join([ - 'chr1', '5', 'A', '*', '5', '35', '[2, 0, 0, 3]', 'AT', '0.6', - '-', '-', '-', '-', '-', + stream.write("\n") + stream.write("\t".join([ + "chr1", "5", "A", "*", "5", "35", "[2, 0, 0, 3]", "AT", "0.6", + "-", "-", "-", "-", "-", ])) - stream.write('\n') - stream.write('\t'.join([ - 'chr1', '9', 'G', '*', '5', '30', '[2, 0, 3, 0]', 'GA', '0.4', - '-', '-', '-', '-', '-', + stream.write("\n") + stream.write("\t".join([ + "chr1", "9", "G", "*", "5", "30", "[2, 0, 3, 0]", "GA", "0.4", + "-", "-", "-", "-", "-", ])) - stream.write('\n') + stream.write("\n") - rta = RTAnnotater({'chrZ': 1, 'chr1': 2, 'chr2': 3}) + rta = RTAnnotater({"chrZ": 1, "chr1": 2, "chr2": 3}) annotated_data = list(rta.merge_files(rna_file, dna_file)) self.assertEqual( annotated_data.pop(0), { - 'Region': 'chr1', - 'Position': '3', - 'Reference': 'A', - 'Strand': '*', - 'Coverage': '5', - 'MeanQ': '30', - 'BaseCount[A,C,G,T]': '[5, 0, 0, 0]', - 'AllSubs': '', - 'Frequency': '0.0', - 'gCoverage': '-', - 'gMeanQ': '-', - 'gBaseCount[A,C,G,T]': '-', - 'gAllSubs': '-', - 'gFrequency': '-', + "Region": "chr1", + "Position": "3", + "Reference": "A", + "Strand": "*", + "Coverage": "5", + "MeanQ": "30", + "BaseCount[A,C,G,T]": "[5, 0, 0, 0]", + "AllSubs": "", + "Frequency": "0.0", + "gCoverage": "-", + "gMeanQ": "-", + "gBaseCount[A,C,G,T]": "-", + "gAllSubs": "-", + "gFrequency": "-", }, ) self.assertEqual( annotated_data.pop(0), { - 'Region': 'chr1', - 'Position': '5', - 'Reference': 'A', - 'Strand': '*', - 'Coverage': '5', - 'MeanQ': '30', - 'BaseCount[A,C,G,T]': '[2, 0, 3, 0]', - 'AllSubs': 'AG', - 'Frequency': '0.6', - 'gCoverage': '5', - 'gMeanQ': '35', - 'gBaseCount[A,C,G,T]': '[2, 0, 0, 3]', - 'gAllSubs': 'AT', - 'gFrequency': '0.6', + "Region": "chr1", + "Position": "5", + "Reference": "A", + "Strand": "*", + "Coverage": "5", + "MeanQ": "30", + "BaseCount[A,C,G,T]": "[2, 0, 3, 0]", + "AllSubs": "AG", + "Frequency": "0.6", + "gCoverage": "5", + "gMeanQ": "35", + "gBaseCount[A,C,G,T]": "[2, 0, 0, 3]", + "gAllSubs": "AT", + "gFrequency": "0.6", }, ) self.assertEqual( annotated_data.pop(0), { - 'Region': 'chr2', - 'Position': '7', - 'Reference': 'G', - 'Strand': '*', - 'Coverage': '5', - 'MeanQ': '30', - 'BaseCount[A,C,G,T]': '[2, 0, 3, 0]', - 'AllSubs': 'GA', - 'Frequency': '0.4', - 'gCoverage': '-', - 'gMeanQ': '-', - 'gBaseCount[A,C,G,T]': '-', - 'gAllSubs': '-', - 'gFrequency': '-', + "Region": "chr2", + "Position": "7", + "Reference": "G", + "Strand": "*", + "Coverage": "5", + "MeanQ": "30", + "BaseCount[A,C,G,T]": "[2, 0, 3, 0]", + "AllSubs": "GA", + "Frequency": "0.4", + "gCoverage": "-", + "gMeanQ": "-", + "gBaseCount[A,C,G,T]": "-", + "gAllSubs": "-", + "gFrequency": "-", }, ) self.assertEqual(len(annotated_data), 0) - os.remove(rna_file) - os.remove(dna_file) + Path(rna_file).unlink() + Path(dna_file).unlink() diff --git a/test/rtindexer.py b/test/rtindexer.py index 9b9c01b..4c78787 100644 --- a/test/rtindexer.py +++ b/test/rtindexer.py @@ -1,47 +1,53 @@ +"""Test cases for RTIndexer class.""" +from __future__ import annotations + import csv -import os import unittest +from pathlib import Path from tempfile import NamedTemporaryFile from reditools.rtindexer import RTIndexer class TestRTIndexer(unittest.TestCase): - test_data = [ + """Test cases for RTIndexer class.""" + + test_data = ( { - 'Region': 'chr1', - 'Position': 1, - 'Reference': 'A', - 'BaseCount[A,C,G,T]': '[10, 0, 0, 0]', + "Region": "chr1", + "Position": 1, + "Reference": "A", + "BaseCount[A,C,G,T]": "[10, 0, 0, 0]", }, { - 'Region': 'chr1', - 'Position': 2, - 'Reference': 'A', - 'BaseCount[A,C,G,T]': '[0, 0, 10, 0]', + "Region": "chr1", + "Position": 2, + "Reference": "A", + "BaseCount[A,C,G,T]": "[0, 0, 10, 0]", }, { - 'Region': 'chr1', - 'Position': 3, - 'Reference': 'G', - 'BaseCount[A,C,G,T]': '[0, 10, 10, 0]', + "Region": "chr1", + "Position": 3, + "Reference": "G", + "BaseCount[A,C,G,T]": "[0, 10, 10, 0]", }, - ] + ) - def setUp(self): + def setUp(self) -> None: + """Pre-flight setup.""" with NamedTemporaryFile( delete=False, - suffix='.out', - mode='wt', + suffix=".out", + mode="wt", ) as stream: self.output_filename = stream.name writer = csv.DictWriter( stream, fieldnames=[ - 'Region', - 'Position', - 'Reference', - 'BaseCount[A,C,G,T]', + "Region", + "Position", + "Reference", + "BaseCount[A,C,G,T]", ], delimiter="\t", ) @@ -50,56 +56,61 @@ def setUp(self): with NamedTemporaryFile( delete=False, - suffix='.bed', - mode='wt', + suffix=".bed", + mode="wt", ) as stream: self.bed_filename = stream.name - stream.write('chr1\t0\t2\n') + stream.write("chr1\t0\t2\n") - def tearDown(self): - os.remove(self.output_filename) - os.remove(self.bed_filename) + def tearDown(self) -> None: + """Post-checks cleanup.""" + Path(self.output_filename).unlink() + Path(self.bed_filename).unlink() - def test_baseline(self): + def test_baseline(self) -> None: + """Check calc_index() method.""" rti = RTIndexer() rti.add_rt_output(self.output_filename) self.assertEqual(rti.calc_index(), { - 'A-C': 0, - 'A-T': 0, - 'A-G': 50, - 'C-A': 0, - 'C-T': 0, - 'C-G': 0, - 'G-A': 0, - 'G-C': 50, - 'G-T': 0, - 'T-A': 0, - 'T-C': 0, - 'T-G': 0, + "A-C": 0, + "A-T": 0, + "A-G": 50, + "C-A": 0, + "C-T": 0, + "C-G": 0, + "G-A": 0, + "G-C": 50, + "G-T": 0, + "T-A": 0, + "T-C": 0, + "T-G": 0, }) - def test_region(self): - rti = RTIndexer(region=('chr1', 100, 200)) - self.assertFalse(rti.do_ignore({'Region': 'chr1', 'Position': '150'})) - self.assertTrue(rti.do_ignore({'Region': 'chr1', 'Position': '50'})) - self.assertTrue(rti.do_ignore({'Region': 'chr1', 'Position': '250'})) - self.assertTrue(rti.do_ignore({'Region': 'chr2', 'Position': '150'})) + def test_region(self) -> None: + """Check do_ignore() method.""" + rti = RTIndexer(region=("chr1", 100, 200)) + self.assertFalse(rti.do_ignore({"Region": "chr1", "Position": "150"})) + self.assertTrue(rti.do_ignore({"Region": "chr1", "Position": "50"})) + self.assertTrue(rti.do_ignore({"Region": "chr1", "Position": "250"})) + self.assertTrue(rti.do_ignore({"Region": "chr2", "Position": "150"})) - rti =RTIndexer(region=('chr1', 100, None)) - self.assertFalse(rti.do_ignore({'Region': 'chr1', 'Position': '150'})) - self.assertTrue(rti.do_ignore({'Region': 'chr1', 'Position': '50'})) - self.assertTrue(rti.do_ignore({'Region': 'chr2', 'Position': '150'})) + rti =RTIndexer(region=("chr1", 100, None)) + self.assertFalse(rti.do_ignore({"Region": "chr1", "Position": "150"})) + self.assertTrue(rti.do_ignore({"Region": "chr1", "Position": "50"})) + self.assertTrue(rti.do_ignore({"Region": "chr2", "Position": "150"})) - def test_targets(self): + def test_targets(self) -> None: + """Check add_target_from_bed() method.""" rti = RTIndexer() rti.add_target_from_bed(self.bed_filename) - self.assertFalse(rti.do_ignore({'Region': 'chr1', 'Position': '1'})) - self.assertTrue(rti.do_ignore({'Region': 'chr1', 'Position': '2'})) - self.assertTrue(rti.do_ignore({'Region': 'chr2', 'Position': '1'})) + self.assertFalse(rti.do_ignore({"Region": "chr1", "Position": "1"})) + self.assertTrue(rti.do_ignore({"Region": "chr1", "Position": "2"})) + self.assertTrue(rti.do_ignore({"Region": "chr2", "Position": "1"})) - def test_exclusions(self): + def test_exclusions(self) -> None: + """Check add_exclusions_from_bed method.""" rti = RTIndexer() rti.add_exclusions_from_bed(self.bed_filename) - self.assertTrue(rti.do_ignore({'Region': 'chr1', 'Position': '1'})) - self.assertFalse(rti.do_ignore({'Region': 'chr1', 'Position': '2'})) - self.assertFalse(rti.do_ignore({'Region': 'chr2', 'Position': '1'})) + self.assertTrue(rti.do_ignore({"Region": "chr1", "Position": "1"})) + self.assertFalse(rti.do_ignore({"Region": "chr1", "Position": "2"})) + self.assertFalse(rti.do_ignore({"Region": "chr2", "Position": "1"})) diff --git a/test/sam_gen.py b/test/sam_gen.py index cce9de2..e65350c 100644 --- a/test/sam_gen.py +++ b/test/sam_gen.py @@ -1,88 +1,219 @@ -import os +"""Classes for SAM and FASTA file generation.""" +from __future__ import annotations + import random import re from dataclasses import InitVar, dataclass +from pathlib import Path from tempfile import NamedTemporaryFile from test.aligner import Aligner +from typing import Any, Iterator from pysam import samtools class Genome: - def __init__(self): - self.contigs = {} - - def __getitem__(self, contig_name): - return self.contigs.get(contig_name, None) - - def add_contig(self, name=None, length=120, sequence=None): + """Genomic sequences object.""" + + def __init__(self) -> None: + """Initialize self.""" + self.contigs: dict[str, str] = {} + + def __getitem__(self, contig_name: str) -> str: + """Retrive chromsomal sequence. + + Parameters + ---------- + contig_name : str + Chromosome name. + + Returns + ------- + str + Nucleotide sequence. + """ + return self.contigs[contig_name] + + def add_contig( + self, + name: str | None=None, + length: int=120, + sequence: str | None=None, + ) -> str: + """Create a new chromosome. + + Parameters + ---------- + name : str | None + Chromosome name. If None, a name will be generated. + length : int + Sequence length (ignored if sequence is not None). + sequence : str | None + Nucleotide sequence. If None, a random sequence will be generated. + + Returns + ------- + str + Chromosome name. + """ if name is None: n_contigs = len(self.contigs) - name = f'contig{len(self.contigs)}' + name = f"contig{len(self.contigs)}" while name in self.contigs: n_contigs += 1 - name = f'contig{len(self.contigs)}' + name = f"contig{len(self.contigs)}" if sequence is None: self.contigs[name] = self._random_seq(length) else: self.contigs[name] = sequence return name - def save_to_fasta(self, filename): - with open(filename, 'w') as stream: - for idx, (name, sequence) in enumerate(self.contigs.items(), 1): - stream.write(f'>{name} {idx}\n{sequence}\n') + def save_to_fasta(self, filename: str) -> None: + """Save the genome to a FASTA file. + + Parameters + ---------- + filename : str + Path to save file to. + """ + with Path(filename).open("w") as stream: + stream.writelines(( + f">{name} {idx}\n{sequence}\n" + for idx, (name, sequence) in enumerate(self.contigs.items(), 1) + )) samtools.faidx(filename) @classmethod - def _random_seq(cls, length): - return ''.join([random.choice('ACTG') for _ in range(length)]) + def _random_seq(cls, length: int) -> str: + """Generate a random nucleotide sequence. + + Parameters + ---------- + length : int + Sequence length. + + Returns + ------- + str + Random nucleotide sequence. + """ + return "".join([random.choice("ACTG") for _ in range(length)]) @dataclass -class Sequence: +class Sequence: + """SAM entry. + + Parameters + ---------- + seq : str + Nucleotide sequence. + start : int + Genomic start. + flag : int + SAM flags + phred : list[int] | None + PHRED scores (defaults to 30) + mapq : int + MAPQ score + qname : str + Read name. + pnext : int + Paired read start. + """ + seq: str start: int flag: int = 0 - phred: InitVar[list | None] = None + phred_list: InitVar[list | None] = None mapq: int = 255 _cigar_str: str | None = None - qname: InitVar[str | None] = None + read_name: InitVar[str | None] = None pnext: int = 0 - + read_n = 0 flag_reverse_strand = 16 phred_default = 30 - def __post_init__(self, phred, qname): - if phred is None: + def __post_init__( + self, + phred_list: list[int] | None, + read_name: str | None, + ) -> None: + """Post initialization. + + Parameters + ---------- + phred : list[int] | None + PHRED scores (defaults to 30) + qname : str | None + Read name (one will be generated if None) + """ + if phred_list is None: self.phred = [self.phred_default for _ in range(len(self.seq))] else: - self.phred = phred + self.phred = phred_list - if qname is None: + if read_name is None: self.qname = self.next_read_name() else: - self.qname = qname + self.qname = read_name + + def __len__(self) -> int: + """Return sequence length. - def __len__(self): + Returns + ------- + int + Sequence length. + """ return len(self.seq) - def __str__(self): + def __str__(self) -> str: + """Retrive nucleotide sequence. + + Returns + ------- + str + Nucleotide sequence. + """ return self.seq - def tlen(self, ref_seq): + def tlen(self, ref_seq: str) -> int: + """Calculate the transcript length. + + Parameters + ---------- + ref_seq : str + Reference sequence. + + Returns + ------- + int + Transcript length. + """ cigar = self.cigar_str(ref_seq) tlen = 0 - for count, op in re.findall(r'(?P\d+)(?P[A-Z])', cigar): - count = int(count) - if op not in ('S', 'I'): - tlen += count + for count, op in re.findall(r"(?P\d+)(?P[A-Z])", cigar): + if op not in ("S", "I"): + tlen += int(count) if self.flag & Sequence.flag_reverse_strand: return -tlen return tlen - def cigar_str(self, ref_seq): + def cigar_str(self, ref_seq: str) -> str: + """Generate alignment CIGAR string. + + Parameters + ---------- + ref_seq : str + Reference sequence + + Returns + ------- + str + CIGAR string. + """ if self._cigar_str is not None: return self._cigar_str alignment = Aligner().align( @@ -90,38 +221,89 @@ def cigar_str(self, ref_seq): str(self), ) cigar_iter = self.assemble_cigar_list(*alignment) - cigar_pieces = [f'{length}{op}' for length, op in cigar_iter] - self._cigar_str = ''.join(cigar_pieces) + cigar_pieces = [f"{length}{op}" for length, op in cigar_iter] + self._cigar_str = "".join(cigar_pieces) return self._cigar_str - def make_pair(self): + def make_pair(self) -> Sequence: + """Generate SAM paired entry. + + Returns + ------- + Sequence + SAM paired sequence. + """ return Sequence( seq=self.seq, start=self.start, flag=self.pair_flag(self.flag), - phred=self.phred, + phred_list=self.phred, mapq=self.mapq, _cigar_str=self._cigar_str, - qname=self.qname, + read_name=self.qname, pnext=self.start, ) @classmethod - def pair_flag(cls, flag_value): + def pair_flag(cls, flag_value: int) -> int: + """Create flags for paired read. + + Parameters + ---------- + flag_value : int + Flags of read to pair + + Returns + ------- + int + Flags for read mate. + """ return flag_value ^ 240 @classmethod - def cigar_op(cls, ref_base, query_base): - if ref_base == '-': - return 'I' - if query_base == '-': - return 'D' + def cigar_op(cls, ref_base: str, query_base: str) -> str: + """Determine CIGAR op from alignment base. + + Parameters + ---------- + ref_base : str + Reference sequence alignment base. + query_base : str + Query sequence alignment base. + + Returns + ------- + str + CIGAR operator. + """ + if ref_base == "-": + return "I" + if query_base == "-": + return "D" if query_base == ref_base: - return 'M' - return 'X' + return "M" + return "X" @classmethod - def assemble_cigar_list(cls, algn_ref, algn_query): + def assemble_cigar_list( + cls, + algn_ref: str, + algn_query: str, + ) -> Iterator[tuple[int, str]]: + """Iterate over CIGAR operators. + + Parameters + ---------- + algn_ref : str + Reference sequence alignment string. + algn_query : str + Query sequence alignment string. + + Yields + ------ + tuple[int, str] + CIGAR operator length and string. + """ last_op = None op_n = 0 for ref_base, query_base in zip(algn_ref, algn_query): @@ -136,97 +318,199 @@ def assemble_cigar_list(cls, algn_ref, algn_query): yield (op_n, cigar_op) @classmethod - def next_read_name(cls): + def next_read_name(cls) -> str: + """Get the next available read name. + + Returns + ------- + str + Next read name in the form "read#" + """ cls.read_n += 1 - return f'read{cls.read_n}' + return f"read{cls.read_n}" class SAM: - def __init__(self): + """SAM file object.""" + + def __init__(self) -> None: + """Initialize self.""" self.genome = Genome() - self.reads = {} + self.reads: dict[str, list[Sequence]] = {} - def __getitem__(self, contig_name): + def __getitem__(self, contig_name: str) -> list[Sequence]: + """Retrive reads for a specific contig. + + Parameters + ---------- + contig_name : str + Chromosome name. + + Returns + ------- + list[Sequence] + Reads aligned to the chromosome. + """ return self.reads[contig_name] - def header(self): - header = ['@HD\tVN:1.5'] + def header(self) -> str: + """Generate SAM header. + + Returns + ------- + str + Header. + """ + header = ["@HD\tVN:1.5"] for contig_name, seq in self.genome.contigs.items(): - header.append(f'@SQ\tSN:{contig_name}\tLN:{len(seq)}') + header.append(f"@SQ\tSN:{contig_name}\tLN:{len(seq)}") header.append( - '@RG\tID:1\tSM:1_AAAAA\tLB:default\tPU:xxx.1\tPL:ILLUMINA', + "@RG\tID:1\tSM:1_AAAAA\tLB:default\tPU:xxx.1\tPL:ILLUMINA", ) - header.append('@PG\tID:reditools\tPN:reditools\tCL:gen_sam.py') - return '\n'.join(header) - - def add_contig(self, contig_name=None, length=120, sequence=None): + header.append("@PG\tID:reditools\tPN:reditools\tCL:gen_sam.py") + return "\n".join(header) + + def add_contig( + self, + contig_name: str | None=None, + length: int=120, + sequence: str | None=None, + ) -> str: + """Add a chromosome to the genomic reference. + + Parameters + ---------- + contig_name : str | None + Chromosome name (generated if not provided). + length : int + Chromosome size (used if sequence is None). + sequence : str | None + Genomic sequence (random if not provided). + + Returns + ------- + str + Chromosome name. + """ contig_name = self.genome.add_contig(contig_name, length, sequence) self.reads[contig_name] = [] return contig_name - def add_read(self, contig_name, sequence_obj): + def add_read(self, contig_name: str, sequence_obj: Sequence) -> None: + """Add a read. + + Parameters + ---------- + contig_name : str + Chromosome the read aligns to. + sequence_obj : Sequence + Read. + """ self.reads[contig_name].append(sequence_obj) - def add_read_pair(self, contig_name, sequence_obj): + def add_read_pair(self, contig_name: str, sequence_obj: Sequence) -> None: + """Add a read and generate its mate. + + Parameters + ---------- + contig_name : str + Chromosome the read aligns to. + sequence_obj : Sequence + Read. + """ self.add_read(contig_name, sequence_obj) self.add_read(contig_name, sequence_obj.make_pair()) - def sam_entries(self): + def sam_entries(self) -> Iterator[str]: + """Iterate over SAM entries. + + Yields + ------ + SAM file lines. + """ for contig, reads in self.reads.items(): ref_seq = self.genome[contig] - for idx, sequence in enumerate(reads): - yield '\t'.join([str(_) for _ in ( + for sequence in reads: + yield "\t".join([str(_) for _ in ( sequence.qname, sequence.flag, contig, sequence.start + 1, sequence.mapq, sequence.cigar_str(ref_seq), - ('*', '=')[sequence.flag & 1], + ("*", "=")[sequence.flag & 1], sequence.pnext + 1, sequence.tlen(ref_seq), str(sequence), - ''.join([self._phred(_) for _ in sequence.phred]), + "".join([self._phred(_) for _ in sequence.phred]), )]) - def save_to_sam(self, bam_filename, genome_filename): + def save_to_sam(self, bam_filename: str, genome_filename: str) -> None: + """Save SAM data to file. + + Parameters + ---------- + bam_filename : str + Path to save to. + genome_filename : str + Path to reference FASTA file. + """ with NamedTemporaryFile( delete=False, - mode='w', - dir='.', - suffix='.sam', + mode="w", + dir=".", + suffix=".sam", ) as stream: sam_filename = stream.name stream.write(self.header()) - stream.write('\n') - stream.write('\n'.join(self.sam_entries())) + stream.write("\n") + stream.write("\n".join(self.sam_entries())) md_sam = samtools.calmd( sam_filename, genome_filename, catch_stdout=True, ) - with open(sam_filename, 'w') as stream: - stream.write(md_sam) - samtools.sort('-o', bam_filename, sam_filename) + with Path(sam_filename).open("w") as stream: + stream.writelines(md_sam) + samtools.sort("-o", bam_filename, sam_filename) samtools.index(bam_filename) - os.remove(sam_filename) - - def _covered_seqs(self, contig_name, position): - return [ - idx for idx, seq in enumerate(self[contig_name]) - if seq.start <= position < seq.stop - ] + Path(sam_filename).unlink() @classmethod - def _phred(cls, int_value): + def _phred(cls, int_value: int) -> str: + """Convert PHRED score to character. + + Parameters + ---------- + int_value : int + PHRED score. + + Returns + ------- + str + PHRED character. + """ return chr(33 + int_value) -def ntf(*args, **kwargs): - with NamedTemporaryFile( +def ntf(*args: Any, **kwargs: Any) -> str: # noqa: ANN401 + """Create a new temporary file. + + Parameters + ---------- + *args : Any + Positional arguments to NamedTemporaryFile. + **kwargs : Any + Named arguments to NamedTemporaryFile. + + Returns + ------- + str + Path to temporary file. + """ + with NamedTemporaryFile( # type: ignore[call-overload] *args, delete=False, - mode='w', + mode="w", **kwargs, ) as stream: - filename = stream.name - return filename + return stream.name diff --git a/test/splicing_file.py b/test/splicing_file.py index 63b38e1..eac6a39 100644 --- a/test/splicing_file.py +++ b/test/splicing_file.py @@ -1,65 +1,100 @@ -import os +"""Test cases for splice_file module.""" +from __future__ import annotations + import unittest +from pathlib import Path from tempfile import NamedTemporaryFile +from typing import Iterable from reditools.region import Region from reditools.splicing_file import load_splicing_file class TestSplicingFile(unittest.TestCase): - def write_file(self, data_list, sep=' '): + """Test cases for splice_file module.""" + + def write_file(self, data_list: str | list, sep: str=" ") -> str: + """Write data to a file. + + Parameters + ---------- + data_list : str | list + Data to write to file. If a list, will use sep to concatenate. + sep : str + Field seperator. Only applicable if data_list is a list. + + Returns + ------- + str + Path to output file. + """ with NamedTemporaryFile( delete=False, - mode='w', - encoding='utf-8', + mode="w", + encoding="utf-8", ) as stream: for row in data_list: if isinstance(row, str): stream.write(row) else: stream.write(sep.join([str(_) for _ in row])) - stream.write('\n') + stream.write("\n") return stream.name - def check_test_data(self, test_data, real_data): + def check_test_data( + self, + test_data: Iterable[list | tuple], + real_data: list, + ) -> None: + """Perform consistency check between real and expected output. + + Parameters + ---------- + test_data : Iterable[Iterable] + Expected output. Uses the second sub element of each element. + real_data : list + Actual output. + """ self.assertEqual([_[1] for _ in test_data], real_data) - def test_splicing_basic(self): + def test_splicing_basic(self) -> None: + """Check load_splicing_file() method.""" test_data = [ ( - ('chr1', '10', '25', 'A', '+'), - Region(contig='chr1', start=4, stop=9), + ("chr1", "10", "25", "A", "+"), + Region(contig="chr1", start=4, stop=9), ), ( - ('chr2', '20', '25', 'D', '-'), - Region(contig='chr2', start=14, stop=19), + ("chr2", "20", "25", "D", "-"), + Region(contig="chr2", start=14, stop=19), ), ( - ('chr3', '5', '15', 'A', '-'), - Region(contig='chr3', start=4, stop=9), + ("chr3", "5", "15", "A", "-"), + Region(contig="chr3", start=4, stop=9), ), ( - ('chr3', '5', '10', 'D', '+'), - Region(contig='chr3', start=4, stop=9), + ("chr3", "5", "10", "D", "+"), + Region(contig="chr3", start=4, stop=9), ), ] fname = self.write_file( - ['#Header'] + [_[0] for _ in test_data], + ["#Header"] + [_[0] for _ in test_data], ) splice_sites = list(load_splicing_file(fname, 5)) self.check_test_data(test_data, splice_sites) - os.remove(fname) + Path(fname).unlink() - def test_splicing_edge(self): + def test_splicing_edge(self) -> None: + """Check effects when at the very start of a contig.""" test_data = [ - ('chr1', '1', '25', 'A', '+'), - ('chr1', '1', '25', 'D', '-'), - ('chr1', '3', '25', 'D', '-'), + ("chr1", "1", "25", "A", "+"), + ("chr1", "1", "25", "D", "-"), + ("chr1", "3", "25", "D", "-"), ] - fname = self.write_file(['#Header'] + test_data) + fname = self.write_file(["#Header", *test_data]) splice_sites = list(load_splicing_file(fname, 5)) self.assertEqual( splice_sites, - [Region(contig='chr1', start=0, stop=2)], + [Region(contig="chr1", start=0, stop=2)], ) - os.remove(fname) + Path(fname).unlink()