sequence_collection.py

from dataclasses import dataclass, field
from typing import List

class NamedSequence():
    def __init__(self, name, sequence, repr=None):
        self.name = name
        self.sequence = sequence
        if repr is None:
            self.repr_sample = name
        else:
            self.repr_sample = repr

@dataclass
class SequenceCollection():
    reads: List[NamedSequence] = field(default_factory=lambda: [])
    references: List[NamedSequence] = field(default_factory=lambda: [])
    haplotypes: List[NamedSequence] = field(default_factory=lambda: [])

    def get_named_sequences(self, descriptor="all"):
        result = []
        if "all" in descriptor:
            descriptor = "reads,refs,haps"
        if "reads" in descriptor:
            result.extend(self.reads)
        if "refs" in descriptor:
            result.extend(self.references)
        if "haps" in descriptor:
            result.extend(self.haplotypes)
        return result

    def get_sequences(self, descriptor="all"):
        return [r.sequence for r in self.get_named_sequences(descriptor)]

    def get_names(self, descriptor="all"):
        return [r.name for r in self.get_named_sequences(descriptor)]

    def get_labels(self):
        return self.get_names("refs,haps")

    def __len__(self):
        return len(self.get_named_sequences())

    def n_reads(self):
        return len(self.reads)
    def n_refs(self):
        return len(self.references)
    def n_haps(self):
        return len(self.haplotypes)

    def __add__(self, other):
        res = SequenceCollection()
        res.reads = self.reads + other.reads
        res.references = self.references + other.references
        res.haplotypes = self.haplotypes + other.haplotypes
        return res

    def __iadd__(self, other):
        self = self.__add__(other)
        return self