Strand symbols

reditools/compiled_position.py

@@ -5,7 +5,12 @@
 from typing import Iterator
 
 from reditools.constants import bases as base_order
-from reditools.constants import comp_map
+from reditools.constants import (
+    comp_map,
+    forward_strand_symbol,
+    reverse_strand_symbol,
+    undetermined_strand_symbol,
+)
 
 
 @dataclass
@@ -78,17 +83,17 @@ def calculate_strand(self, threshold: float = 0) -> str:
         pos_count = 0
         neg_count = 0
         for strand in self.strands:
-            if strand == "+":
+            if strand == forward_strand_symbol:
                 pos_count += 1
-            elif strand == "-":
+            elif strand == reverse_strand_symbol:
                 neg_count += 1
         if pos_count == neg_count:
-            return "*"
+            return undetermined_strand_symbol
         if pos_count / (pos_count + neg_count) >= threshold:
-            return "+"
+            return forward_strand_symbol
         if neg_count / (pos_count + neg_count) >= threshold:
-            return "-"
-        return "*"
+            return reverse_strand_symbol
+        return undetermined_strand_symbol
 
     def filter_by_strand(self, strand: str) -> None:
         """Filter observations to keep only those from a specific strand.
@@ -98,7 +103,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 == undetermined_strand_symbol:
             return
         keep = [
             idx for idx in range(len(self.bases))

reditools/compiled_reads.py

@@ -4,6 +4,11 @@
 from typing import TYPE_CHECKING
 
 from reditools.compiled_position import CompiledPosition
+from reditools.constants import (
+    forward_strand_symbol,
+    reverse_strand_symbol,
+    undetermined_strand_symbol,
+)
 from reditools.fasta_file import RTFastaFile
 
 if TYPE_CHECKING:
@@ -99,7 +104,11 @@ class CompiledReads:
     and strand.
     """
 
-    _strands = ("-", "+", "*")
+    _strands = (
+        reverse_strand_symbol,
+        forward_strand_symbol,
+        undetermined_strand_symbol,
+    )
 
     def __init__(
         self,

reditools/constants.py

@@ -1,4 +1,6 @@
-"""Dictionary for base complements."""
+"""Constants."""
+
+"""Complementary base dictionary."""
 comp_map = {
     "A": "T",
     "T": "A",
@@ -7,4 +9,16 @@
     "N": "N",
     "-": "-",
 }
+
+"""Nucleotides in output order."""
 bases = ("A", "C", "G", "T")
+
+"""Output strand symbols."""
+forward_strand_symbol = "+"
+reverse_strand_symbol = "-"
+undetermined_strand_symbol = "*"
+
+"""Output strand numbers."""
+forward_strand_number = "1"
+reverse_strand_number = "0"
+undetermined_strand_number = "2"

reditools/tools/analyze/parse_args/parse_args.py

@@ -89,6 +89,14 @@ def build_argument_parser() -> argparse.ArgumentParser:  # noqa: WPS213, WPS210
         action="store_true",
         help="Appends results to file (and creates if not existing).",
     )
+    output_group.add_argument(
+        "--number-strand-output",
+        action="store_true",
+        help=(
+            "Replaces the '-', '+', and '*' values with 0, 1, and 2 in the "
+            "Strand column of the output."
+        ),
+    )
     bqf_group = parser.add_argument_group(
         title="Base/Read Quality Controls",
     )

reditools/tools/analyze/redi_thread.py

@@ -44,6 +44,7 @@ def __init__(self, options: argparse.Namespace) -> None:
         )
         self.rtqc = RTChecks(options)
         self.temp_dir = options.temp_dir
+        self.number_strand_output = options.number_strand_output
 
     def analyze(
             self,
@@ -65,6 +66,7 @@ def analyze(
             filename,
             self.rtqc,
             self.rtools.log,
+            self.number_strand_output,
         )
 
 class REDIThreadManager:

reditools/tools/analyze/write_results.py

@@ -2,33 +2,57 @@
 
 import csv
 from pathlib import Path
-from typing import Callable, Iterator
+from typing import Callable, Iterable
 
 from reditools.compiled_position import RTResult
+from reditools.constants import (
+    forward_strand_number,
+    forward_strand_symbol,
+    reverse_strand_number,
+    reverse_strand_symbol,
+    undetermined_strand_number,
+    undetermined_strand_symbol,
+)
 from reditools.logger import Logger
 from reditools.tools.analyze.rtchecks import RTChecks
 
 _empty = "-"
 
 def write_results(
-        rtresults: Iterator[RTResult],
+        rtresults: Iterable[RTResult],
         filename: str,
         filters: RTChecks,
         logger: Callable,
+        strand_numbers: bool = False,
 ) -> None:
     """Write analysis results to a file.
 
     Parameters
     ----------
-    rtresults : Iterator[RTResult]
+    rtresults : Iterable[RTResult]
         The analysis results for each position.
     filename : str
         Where to save the results.
     filters : RTChecks
         The quality control checks to apply.
     logger : Callable
         The logger function for debug messages.
+    strand_numbers : bool
+        If False, uses -, +, and * for the Strand column.
+        If True, uses 0, 1, and 2 for the Strand column.
     """
+    if strand_numbers:
+        strand_lookup = {
+            forward_strand_symbol: forward_strand_number,
+            reverse_strand_symbol: reverse_strand_number,
+            undetermined_strand_symbol: undetermined_strand_number,
+        }
+    else:
+        strand_lookup = {
+            forward_strand_symbol: forward_strand_symbol,
+            reverse_strand_symbol: reverse_strand_symbol,
+            undetermined_strand_symbol: undetermined_strand_symbol,
+        }
     with Path(filename).open("w") as stream:
         writer = csv.writer(stream, delimiter="\t", lineterminator="\n")
         for rt_result in rtresults:
@@ -41,7 +65,7 @@ def write_results(
                 rt_result.contig,
                 rt_result.position + 1,
                 rt_result.reference,
-                rt_result.strand,
+                strand_lookup[rt_result.strand],
                 len(rt_result),
                 f"{rt_result.mean_quality:.2f}",
                 list(rt_result),

test/__main__.py

@@ -8,6 +8,7 @@
 from test.analyze.rtchecks import TestRTChecks
 from test.analyze.setup_alignment_manager import TestSetupAlignmentManager
 from test.analyze.setup_rtools import TestSetupRTools
+from test.analyze.write_results import TestWriteResults
 from test.compiled_position import TestCompiledPosition
 from test.compiled_reads import TestCompiledReads
 from test.fasta_file import TestRTFastaFile

test/analyze/write_results.py

@@ -0,0 +1,124 @@
+"""Test methods for analyze tool's write_results()."""
+from __future__ import annotations
+
+import unittest
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+from reditools.compiled_position import CompiledPosition, RTResult
+from reditools.logger import Logger
+from reditools.tools.analyze.parse_args.parse_args import parse_args
+from reditools.tools.analyze.rtchecks import RTChecks
+from reditools.tools.analyze.write_results import write_results
+
+
+class TestWriteResults(unittest.TestCase):
+    """Test methods for analyze tool's write_results()."""
+
+    def setUp(self) -> None:
+        """Pre-flight setup."""
+        self.rtresults = []
+
+        cp = CompiledPosition(ref="A", contig="chr1", position=100)
+        cp.add_base(40, "+", "A")
+        cp.add_base(30, "+", "G")
+        cp.add_base(30, "+", "T")
+        self.rtresults.append(RTResult(cp, "+"))
+        self.rtresults.append(RTResult(cp, "-"))
+        self.rtresults.append(RTResult(cp, "*"))
+
+        with NamedTemporaryFile(mode="w", dir=".", delete=False) as temp_file:
+            self.fname = temp_file.name
+        self.log = Logger(Logger.silent_level).log
+
+    def tearDown(self) -> None:
+        """Post-checks cleanup."""
+        Path(self.fname).unlink()
+
+    def rtchecks(self, cli_args: list[str]) -> RTChecks:
+        """Create RTChecks object from CLI arguments.
+
+        Parameters
+        ----------
+        cli_args : list[str]
+            Command line arguments.
+
+        Returns
+        -------
+        RTChecks
+            Checks for output filtering.
+        """
+        return RTChecks(parse_args(cli_args))
+
+    def test_write_results(self) -> None:
+        """Check basic functionality."""
+        write_results(
+            self.rtresults,
+            self.fname,
+            self.rtchecks(["x.bam"]),
+            self.log,
+        )
+        with Path(self.fname).open("r") as stream:
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "+", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "-", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "*", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+
+    def test_write_results_numbers(self) -> None:
+        """Check strand-number option."""
+        write_results(
+            self.rtresults,
+            self.fname,
+            self.rtchecks(["x.bam"]),
+            self.log,
+            True,  # noqa: FBT003
+        )
+        with Path(self.fname).open("r") as stream:
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "1", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "0", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+            self.assertEqual(
+                next(stream).strip().split("\t") ,
+                [
+                    "chr1", "101", "A", "2", "3", "33.33", "[1, 0, 1, 1]",
+                    "AG AT", "0.50", "-", "-", "-", "-", "-",
+                ],
+            )
+
+    def test_filter(self) -> None:
+        """Check that RTChecks filters output."""
+        write_results(
+            self.rtresults,
+            self.fname,
+            self.rtchecks(["x.bam", "--min-read-depth", "5"]),
+            self.log,
+        )
+        self.assertEqual(Path(self.fname).stat().st_size, 0)