Annotated bases flag on api.query

Author: tombch

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "maptide"
-version = "0.2.1"
+version = "0.3.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

Cargo.toml

@@ -1,15 +1,19 @@
 import os
-from typing import Dict, Tuple, List, Optional
+from typing import Dict, Tuple, Optional, Any
 from . import maptide  #  type: ignore
 
 
+BASES = ["A", "C", "G", "T", "DS", "N"]
+
+
 def query(
     bam: str,
     region: Optional[str] = None,
     bai: Optional[str] = None,
     mapping_quality: int = 0,
     base_quality: int = 0,
-) -> Dict[str, Dict[Tuple[int, int], List[int]]]:
+    annotated: bool = False,
+) -> Dict[str, Dict[Tuple[int, int], Any]]:
     """Performs a pileup over a region, obtaining per-position base frequencies for the provided BAM file.
 
     Parameters
@@ -24,6 +28,8 @@ def query(
         Minimum mapping quality for a read to be included in the pileup (default: 0)
     base_quality : int, optional
         Minimum base quality for a base within a read to be included in the pileup (default: 0)
+    annotated : bool, optional
+        Return frequencies annotated with their bases, as a `dict[str, int]`. Default is to return frequencies only, as a `list[int]` (default: False)
 
     Returns
     -------
@@ -34,9 +40,16 @@ def query(
     if region:
         if not bai and os.path.isfile(bam + ".bai"):
             bai = bam + ".bai"
-        return maptide.query(bam, bai, region, mapping_quality, base_quality)
+        data = maptide.query(bam, bai, region, mapping_quality, base_quality)
     else:
-        return maptide.all(bam, mapping_quality, base_quality)
+        data = maptide.all(bam, mapping_quality, base_quality)
+
+    if annotated:
+        for _, positions in data.items():
+            for position, frequencies in positions.items():
+                positions[position] = dict(zip(BASES, frequencies))
+
+    return data
 
 
 def parse_region(region: str) -> Tuple[str, int, int]:

python/maptide/api.py

@@ -108,23 +108,12 @@ def run():
         "pos",
         "ins",
         "cov",
-        "a",
-        "c",
-        "g",
-        "t",
-        "ds",
-        "n",
-    ]
+    ] + [base.lower() for base in api.BASES]
 
     if args.stats:
         columns.extend(
-            [
-                "pc_a",
-                "pc_c",
-                "pc_g",
-                "pc_t",
-                "pc_ds",
-                "pc_n",
+            [f"pc_{base.lower()}" for base in api.BASES]
+            + [
                 "entropy",
                 "secondary_entropy",
             ]