biofiles 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

biofiles/bam.py

@@ -0,0 +1,199 @@
+import gzip
+import struct
+import sys
+from io import BytesIO
+from pathlib import Path
+from types import TracebackType
+from typing import Iterator, Any
+
+from biofiles.types.alignment import (
+    ReferenceSequence,
+    Alignment,
+    BAMTag,
+    CIGAR,
+    CIGAROpKind,
+    CIGAROperation,
+)
+
+
+class BAMReader:
+    def __init__(self, input_: BytesIO | Path | str) -> None:
+        if isinstance(input_, Path | str):
+            input_ = open(input_, "rb")
+        self._input = input_
+        self._ungzipped_input = gzip.open(input_)
+
+        self._header_text: str | None = None
+        self._ref_seqs: list[ReferenceSequence] = []
+
+        self._read_header()
+
+    def _read_header(self) -> None:
+        magic_bytes = self._ungzipped_input.read(8)
+        magic_data = struct.unpack("<ccccI", magic_bytes)
+        if b"".join(magic_data[:4]) != b"BAM\1":
+            raise ValueError("not a BAM file, invalid magic bytes")
+
+        header_text_length = magic_data[-1]
+        self._header_text = self._ungzipped_input.read(header_text_length)
+        (num_ref_seqs,) = struct.unpack("<I", self._ungzipped_input.read(4))
+
+        for _ in range(num_ref_seqs):
+            (ref_seq_name_length,) = struct.unpack("<I", self._ungzipped_input.read(4))
+            ref_seq_name = self._ungzipped_input.read(ref_seq_name_length)
+            (ref_seq_length,) = struct.unpack("<I", self._ungzipped_input.read(4))
+            ref_seq = ReferenceSequence(
+                id=ref_seq_name.rstrip(b"\0").decode("ascii"), length=ref_seq_length
+            )
+            self._ref_seqs.append(ref_seq)
+
+    def __iter__(self) -> Iterator[Alignment]:
+        return self
+
+    def __next__(self) -> Alignment:
+        block_size_bytes = self._ungzipped_input.read(4)
+        if not block_size_bytes:
+            raise StopIteration
+
+        (block_length,) = struct.unpack("<I", block_size_bytes)
+
+        body_format = "<iiBBHHHIiii"
+        body_bytes = self._ungzipped_input.read(struct.calcsize(body_format))
+        (
+            ref_seq_idx,
+            pos,
+            read_name_length,
+            mapping_quality,
+            bai_index_bin,
+            num_cigar_ops,
+            flags,
+            seq_length,
+            next_ref_seq_idx,
+            next_pos,
+            template_length,
+        ) = struct.unpack(body_format, body_bytes)
+        read_name_bytes = self._ungzipped_input.read(read_name_length)
+
+        cigar_format = "<" + "I" * num_cigar_ops
+        cigar_bytes = self._ungzipped_input.read(struct.calcsize(cigar_format))
+        encoded_cigar = struct.unpack(cigar_format, cigar_bytes)
+
+        seq_bytes = self._ungzipped_input.read((seq_length + 1) // 2)
+        encoded_seq = struct.unpack("<" + "B" * len(seq_bytes), seq_bytes)
+
+        quality = self._ungzipped_input.read(seq_length).decode("ascii")
+
+        remaining_length = (
+            block_length
+            - len(body_bytes)
+            - len(read_name_bytes)
+            - len(cigar_bytes)
+            - len(seq_bytes)
+            - len(quality)
+        )
+
+        tags: list[BAMTag] = []
+        while remaining_length > 0:
+            tag, used_length = self._read_tag()
+            tags.append(tag)
+            remaining_length -= used_length
+        if remaining_length < 0:
+            raise ValueError("invalid BAM file, wrong tag length")
+
+        ref_seq = self._ref_seqs[ref_seq_idx] if ref_seq_idx >= 0 else None
+        next_ref_seq = (
+            self._ref_seqs[next_ref_seq_idx] if next_ref_seq_idx >= 0 else None
+        )
+        return Alignment(
+            reference_sequence=ref_seq,
+            start_c=pos,
+            read_name=read_name_bytes.rstrip(b"\0").decode("utf-8"),
+            mapping_quality=mapping_quality,
+            bai_index_bin=bai_index_bin,
+            next_reference_sequence=next_ref_seq,
+            next_start_c=next_pos,
+            template_length=template_length,
+            cigar=self._decode_cigar(encoded_cigar),
+            read_sequence=self._decode_seq(encoded_seq),
+            quality=quality,
+            bam_flags=flags,
+            bam_tags=tuple(tags),
+        )
+
+    def _decode_cigar(self, encoded_cigar: tuple[int, ...]) -> CIGAR:
+        return CIGAR(
+            operations=tuple(
+                CIGAROperation(kind=_BAM_CIGAR_OP_KINDS[item & 0b1111], count=item >> 4)
+                for item in encoded_cigar
+            )
+        )
+
+    def _decode_seq(self, encoded_seq: tuple[int, ...]) -> str:
+        return "".join(
+            f"{_BAM_SEQUENCE_LETTERS[b >> 4]}{_BAM_SEQUENCE_LETTERS[b & 15]}"
+            for b in encoded_seq
+        )
+
+    def _read_tag(self) -> tuple[BAMTag, int]:
+        tag = self._ungzipped_input.read(2).decode("ascii")
+        value_type = self._ungzipped_input.read(1)
+        value, value_length = self._read_tag_value(value_type)
+        return BAMTag(tag=tag, value=value), 3 + value_length
+
+    def _read_tag_value(self, value_type: bytes) -> tuple[Any, int]:
+        if value_type in (b"Z", b"H"):
+            characters: list[bytes] = []
+            last_character = b""
+            while last_character != b"\0":
+                characters.append(last_character)
+                last_character = self._ungzipped_input.read(1)
+            value = b"".join(characters).decode("utf-8")
+            return value, len(characters)
+
+        elif value_type == b"B":
+            subtype, count = struct.unpack("<cI", self._ungzipped_input.read(5))
+            format_ = "<" + _BAM_FORMAT_TO_STRUCT_FORMAT[subtype] * count
+            length = struct.calcsize(format_)
+            value = struct.unpack(format_, self._ungzipped_input.read(length))
+            return value, 5 + length
+
+        else:
+            format_ = "<" + _BAM_FORMAT_TO_STRUCT_FORMAT[value_type]
+            length = struct.calcsize(format_)
+            (value,) = struct.unpack(format_, self._ungzipped_input.read(length))
+            return value, length
+
+    def __enter__(self):
+        self._input.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self._input.__exit__(exc_type, exc_val, exc_tb)
+
+
+_BAM_FORMAT_TO_STRUCT_FORMAT = {
+    b"A": "c",
+    b"c": "b",
+    b"C": "B",
+    b"s": "h",
+    b"S": "H",
+    b"i": "i",
+    b"I": "I",
+    b"f": "f",
+}
+
+_BAM_CIGAR_OP_KINDS: list[CIGAROpKind] = ["M", "I", "D", "N", "S", "H", "P", "=", "X"]
+_BAM_SEQUENCE_LETTERS = "=ACMGRSVTWYHKDBN"
+
+if __name__ == "__main__":
+    for path in sys.argv[1:]:
+        num_alignments = 0
+        with BAMReader(path) as reader:
+            for record in reader:
+                num_alignments += 1
+        print(f"Parsed {num_alignments} alignments from {path}")

biofiles/gff.py

@@ -3,7 +3,8 @@ from pathlib import Path
 from typing import Iterator, cast, TextIO
 
 from biofiles.common import Strand, Writer
-from biofiles.feature import FeatureReader, FeatureDraft, FeatureDrafts
+from biofiles.utility.cli import parse_pipeline_args
+from biofiles.utility.feature import FeatureReader, FeatureDraft, FeatureDrafts
 from biofiles.types.feature import Feature, Gene, Exon, UTR
 
 __all__ = ["GFFReader", "GFF3Writer"]
@@ -137,7 +138,15 @@ _VERSION_PREFIX = "##gff-version "
 
 
 if __name__ == "__main__":
-    for path in sys.argv[1:]:
+    pipeline = parse_pipeline_args(sys.argv[1:])
+    if pipeline.mapper is None:
+        writer = GFF3Writer(sys.stdout)
+        pipeline.mapper = writer.write
+    else:
+        old_mapper = pipeline.mapper
+        pipeline.mapper = lambda f: print(old_mapper(f))
+
+    for path in pipeline.inputs:
         with GFFReader(path) as r:
             total_features = 0
             annotated_genes = 0
@@ -154,9 +163,14 @@ if __name__ == "__main__":
                 parsed_genes += isinstance(feature, Gene)
                 parsed_exons += isinstance(feature, Exon)
                 parsed_utrs += isinstance(feature, UTR)
+
+                if pipeline.filter(feature):
+                    pipeline.map(feature)
+
         print(
             f"{path}: {total_features} features, "
             f"{parsed_genes} genes parsed out of {annotated_genes}, "
             f"{parsed_exons} exons parsed out of {annotated_exons}, "
-            f"{parsed_utrs} UTRs parsed out of {annotated_utrs}"
+            f"{parsed_utrs} UTRs parsed out of {annotated_utrs}",
+            file=sys.stderr,
         )

biofiles/gtf.py

@@ -1,8 +1,9 @@
-__all__ = ["GTFReader"]
+__all__ = ["GTFReader", "GTFWriter"]
 
 import sys
 from typing import Iterator
 
+from biofiles.common import Writer
 from biofiles.gff import GFFReader
 from biofiles.types.feature import Gene, Exon, Feature, UTR
 
@@ -13,12 +14,32 @@ class GTFReader(GFFReader):
 
     def _parse_attributes(self, line: str, attributes_str: str) -> dict[str, str]:
         return {
-            k: v.strip('"')
+            k: v.removeprefix('"').removesuffix('"').replace(r"\"", '"')
             for part in attributes_str.strip(";").split(";")
             for k, v in (part.strip().split(None, 1),)
         }
 
 
+class GTFWriter(Writer):
+    def write(self, feature: Feature) -> None:
+        fields = (
+            feature.sequence_id,
+            feature.source,
+            feature.type_,
+            str(feature.start_c + 1),
+            str(feature.end_c),
+            str(feature.score) if feature.score is not None else ".",
+            str(feature.strand) if feature.strand is not None else ".",
+            str(feature.phase) if feature.phase is not None else ".",
+            "; ".join(
+                f'{k} "' + v.replace('"', r"\"") + '"'
+                for k, v in feature.attributes.items()
+            ),
+        )
+        self._output.write("\t".join(fields))
+        self._output.write("\n")
+
+
 if __name__ == "__main__":
     for path in sys.argv[1:]:
         with GTFReader(path) as r:
@@ -41,5 +62,6 @@ if __name__ == "__main__":
             f"{path}: {total_features} features, "
             f"{parsed_genes} genes parsed out of {annotated_genes}, "
             f"{parsed_exons} exons parsed out of {annotated_exons}, "
-            f"{parsed_utrs} UTRs parsed out of {annotated_utrs}"
+            f"{parsed_utrs} UTRs parsed out of {annotated_utrs}",
+            file=sys.stderr,
         )

biofiles/types/alignment.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+
+__all__ = ["ReferenceSequence", "Alignment", "BAMTag"]
+
+from enum import IntFlag
+
+from typing import Any, Literal
+
+
+@dataclass(frozen=True)
+class ReferenceSequence:
+    id: str
+    length: int
+
+
+@dataclass(frozen=True, slots=True)
+class BAMTag:
+    tag: str
+    value: Any
+
+
+CIGAROpKind = Literal["M", "I", "D", "N", "S", "H", "P", "=", "X"]
+
+
+@dataclass(frozen=True, slots=True)
+class CIGAROperation:
+    kind: CIGAROpKind
+    count: int
+
+
+@dataclass(frozen=True)
+class CIGAR:
+    operations: tuple[CIGAROperation, ...]
+
+    def __repr__(self) -> str:
+        return f'CIGAR("{self}")'
+
+    def __str__(self) -> str:
+        return "".join(f"{op.count}{op.kind}" for op in self.operations)
+
+
+class BAMFlag(IntFlag):
+    MULTIPLE_SEGMENTS = 1 << 0
+    EACH_SEGMENT_PROPERLY_ALIGNED = 1 << 1
+    SEGMENT_UNMAPPED = 1 << 2
+    NEXT_SEGMENT_UNMAPPED = 1 << 3
+    READ_SEQUENCE_REVERSE_COMPLEMENTED = 1 << 4
+    NEXT_SEGMENT_READ_SEQUENCE_REVERSE_COMPLEMENTED = 1 << 5
+    FIRST_SEGMENT = 1 << 6
+    LAST_SEGMENT = 1 << 7
+    SECONDARY_SEGMENT = 1 << 8
+    NOT_PASSING_QUALITY_CONTROL = 1 << 9
+    DUPLICATE = 1 << 10
+    SUPPLEMENTARY_ALIGNMENT = 1 << 11
+
+
+@dataclass(frozen=True)
+class Alignment:
+    reference_sequence: ReferenceSequence | None
+
+    start_c: int
+    # 0-based leftmost coordinate.
+    read_name: str
+    mapping_quality: int
+    bai_index_bin: int
+
+    next_reference_sequence: ReferenceSequence | None
+    next_start_c: int
+    template_length: int
+    cigar: CIGAR
+    read_sequence: str
+    quality: str
+
+    bam_flags: int
+    bam_tags: tuple[BAMTag, ...]

biofiles/utility/cli.py

@@ -0,0 +1,126 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TypeAlias, Callable, Any, Literal, Type
+
+from biofiles.types.feature import Feature, Gene, Transcript, UTR, Exon
+
+FeatureFilter: TypeAlias = Callable[[Feature], bool]
+FeatureMapper: TypeAlias = Callable[[Feature], Any]
+
+
+@dataclass
+class Pipeline:
+    inputs: list[Path]
+    filters: list[FeatureFilter]
+    mapper: FeatureMapper | None
+
+    def filter(self, feature: Feature) -> bool:
+        for f in self.filters:
+            if not f(feature):
+                return False
+        return True
+
+    def map(self, feature: Feature) -> Any:
+        if not self.mapper:
+            return feature
+        return self.mapper(feature)
+
+
+Mode: TypeAlias = Literal["inputs", "filters", "done"]
+
+
+def parse_pipeline_args(argv: list[str]) -> Pipeline:
+    pipeline = Pipeline(inputs=[], filters=[], mapper=None)
+
+    mode: Mode = "inputs"
+    i = 0
+    while i < len(argv):
+        match mode, argv[i:]:
+            case "inputs", [str_path, *_] if (path := Path(str_path)).is_file():
+                pipeline.inputs.append(path)
+                i += 1
+            case "inputs", ["--filter", *_]:
+                mode = "filters"
+                i += 1
+            case "inputs" | "filters", ["--attr", key]:
+                path = key.split(".")
+                pipeline.mapper = _produce_attr_mapper(path)
+                mode = "done"
+                i += 2
+            case "filters", [filter_str, *_]:
+                filter_ = _parse_filter(filter_str)
+                pipeline.filters.append(filter_)
+                i += 1
+            case other:
+                raise ValueError(f"can't parse command line arguments {argv[i:]}")
+
+    return pipeline
+
+
+def _parse_filter(filter_str: str) -> FeatureFilter:
+    if "=" not in filter_str:
+        # --filter gene,transcript
+        type_strs = filter_str.split(",")
+        types = tuple(_parse_feature_type(t) for t in type_strs)
+        return lambda f: isinstance(f, types)
+
+    # --filter attr=value1,value2
+    key, value = filter_str.split("=", maxsplit=1)
+    values = value.split(",")
+    match key:
+        case "chromosome":
+            return lambda f: f.sequence_id in values
+        case "type":
+            return lambda f: f.type_ in values
+        case "strand":
+            return lambda f: f.strand in values
+        case _:
+            path = key.split(".")
+            return _produce_attr_filter(path, values)
+
+    raise ValueError(f"can't parse filter {filter_str!r}")
+
+
+def _parse_feature_type(t: str) -> Type[Feature]:
+    if t not in _FEATURE_TYPES:
+        raise ValueError(f"unknown feature type {t!r}")
+    return _FEATURE_TYPES[t]
+
+
+def _produce_attr_filter(path: list[str], values: list[str]) -> FeatureFilter:
+    assert path
+    if len(path) == 1:
+        (key,) = path
+        match key:
+            case "chromosome" | "type" | "strand" | "id":
+                return lambda f: getattr(f, key) in values
+            # TODO other attributes
+            case _:
+                return lambda f: f.attributes.get(key) in values
+
+    if path[0] not in ("gene", "transcript", "parent"):
+        raise ValueError(f"unknown attribute {path[-2]!r}")
+
+    nested = _produce_attr_filter(path[1:], values)
+    return lambda f: (nested(nf) if (nf := getattr(f, path[0], None)) else False)
+
+
+def _produce_attr_mapper(path: list[str]) -> FeatureMapper:
+    assert path
+    if len(path) == 1:
+        (key,) = path
+        match key:
+            case "chromosome" | "type" | "strand" | "id":
+                return lambda f: getattr(f, key)
+            # TODO other attributes
+            case _:
+                return lambda f: f.attributes.get(key, "")
+
+    if path[0] not in ("gene", "transcript", "parent"):
+        raise ValueError(f"unknown attribute {path[-2]!r}")
+
+    nested = _produce_attr_mapper(path[1:])
+    return lambda f: (nested(nf) if (nf := getattr(f, path[0], None)) else None)
+
+
+_FEATURE_TYPES = {"gene": Gene, "transcript": Transcript, "exon": Exon, "utr": UTR}

{biofiles-0.0.9.dist-info → biofiles-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: biofiles
-Version: 0.0.9
+Version: 0.0.10
 Summary: Pure-Python, zero-dependency collection of bioinformatics-related file readers and writers
 Author-email: Tigran Saluev <tigran@saluev.com>
 Maintainer-email: Tigran Saluev <tigran@saluev.com>

biofiles-0.0.10.dist-info/RECORD

@@ -0,0 +1,20 @@
+biofiles/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+biofiles/bam.py,sha256=w32LLOAuKWdGF7joTSrB4HYXCdfvvijZW44jizG36R8,6771
+biofiles/common.py,sha256=Yi0i85FpD2wR3vqL645LTUAE6TybGDxxZQsUmEGHqu4,1126
+biofiles/fasta.py,sha256=ctIt5I_fcZx-xQN921zpmlZS7e9_ICf-3_i6mTs5qbs,2135
+biofiles/gff.py,sha256=b3apOmJNoiy_qQHtyUSnNh0s999B6gyAODyjI7fN15g,6246
+biofiles/gtf.py,sha256=h_eFKnYWb8GQp-CX9EPZRodUba-bzQLGidGHOPUo4iM,2366
+biofiles/repeatmasker.py,sha256=txOYdw15ru88pUczsk0pDFzgGpplLu23CB8Ppz-MczY,3119
+biofiles/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+biofiles/types/alignment.py,sha256=5UvwKJ2psIpkkU5efGRHe8gYhMIoW35-RZ_Zoe5YDrY,1612
+biofiles/types/feature.py,sha256=3Ar45WRgiaDSh5iQt24Emtk6_57G01q5nHJ1GNIJ19Y,1190
+biofiles/types/repeat.py,sha256=63SqzAwEGIDIGP9pxC85RUdwXbbSm0S5WNL3lSiWlmc,641
+biofiles/types/sequence.py,sha256=EOw_oKuMR0THpCYJqVE__27z7qrRqcdIPrRWTL4OFMw,152
+biofiles/utility/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+biofiles/utility/cli.py,sha256=bkUzmT5R4qdJ0YtA4LNU5JYpimD1HmZlHtoSaKzDsUc,4032
+biofiles/utility/feature.py,sha256=tUTn16xV1e0qpgkZ1ZwQ4LJJGil5mgQJBJ9s1yFDgiI,8068
+biofiles-0.0.10.dist-info/LICENSE,sha256=CbR8ssdFyViKj25JAlMjIt1_FbiZ1tAC5t-uwUbxqak,1070
+biofiles-0.0.10.dist-info/METADATA,sha256=jthXyKbpyvig9dgiQmUaIyDH8hWKX2zMyojyIxr5mjM,3034
+biofiles-0.0.10.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+biofiles-0.0.10.dist-info/top_level.txt,sha256=laFaFv8hpkI4U-Pgs0yBaAJXN2_CJKl7jb-m3-tGfSc,9
+biofiles-0.0.10.dist-info/RECORD,,

{biofiles-0.0.9.dist-info → biofiles-0.0.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (76.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

biofiles-0.0.9.dist-info/RECORD

@@ -1,16 +0,0 @@
-biofiles/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-biofiles/common.py,sha256=Yi0i85FpD2wR3vqL645LTUAE6TybGDxxZQsUmEGHqu4,1126
-biofiles/fasta.py,sha256=ctIt5I_fcZx-xQN921zpmlZS7e9_ICf-3_i6mTs5qbs,2135
-biofiles/feature.py,sha256=tUTn16xV1e0qpgkZ1ZwQ4LJJGil5mgQJBJ9s1yFDgiI,8068
-biofiles/gff.py,sha256=6xmwnuU1CsFibIHzbggYJajzQC4KGsGAfWMxyYFFChw,5798
-biofiles/gtf.py,sha256=kAt_5ifb0f8cCR-kycnQhkyo78xOynqTUUGqgOP8tjA,1543
-biofiles/repeatmasker.py,sha256=txOYdw15ru88pUczsk0pDFzgGpplLu23CB8Ppz-MczY,3119
-biofiles/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-biofiles/types/feature.py,sha256=3Ar45WRgiaDSh5iQt24Emtk6_57G01q5nHJ1GNIJ19Y,1190
-biofiles/types/repeat.py,sha256=63SqzAwEGIDIGP9pxC85RUdwXbbSm0S5WNL3lSiWlmc,641
-biofiles/types/sequence.py,sha256=EOw_oKuMR0THpCYJqVE__27z7qrRqcdIPrRWTL4OFMw,152
-biofiles-0.0.9.dist-info/LICENSE,sha256=CbR8ssdFyViKj25JAlMjIt1_FbiZ1tAC5t-uwUbxqak,1070
-biofiles-0.0.9.dist-info/METADATA,sha256=gLu3ufoag4tZllgq9xCDZe_kA24RXuI4TqQdAI_QIKw,3033
-biofiles-0.0.9.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-biofiles-0.0.9.dist-info/top_level.txt,sha256=laFaFv8hpkI4U-Pgs0yBaAJXN2_CJKl7jb-m3-tGfSc,9
-biofiles-0.0.9.dist-info/RECORD,,