class FilterByLength(_BaseFilter):
    'It removes the sequences according to their length.'
    def __init__(self, minimum=None, maximum=None, ignore_masked=False,
                 failed_drags_pair=True):
        '''The initiator.

        threshold - minimum length to pass the filter (integer)
        reverse - if True keep the short sequences and discard the long ones
        ignore_masked - If True only uppercase letters will be counted.
        '''
        if min is None and max is None:
            raise ValueError('min or max threshold must be given')
        self.min = minimum
        self.max = maximum
        self.ignore_masked = ignore_masked
        super(FilterByLength, self).__init__(reverse=False,
                                           failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        min_ = self.min
        max_ = self.max
        length = uppercase_length(get_str_seq(seq)) if self.ignore_masked else get_length(seq)

        passed = True
        if min_ is not None and length < min_:
            passed = False
        if max_ is not None and length > max_:
            passed = False
        return passed class FilterById(_BaseFilter):
    'It removes the sequences not found in the given set'
    def __init__(self, seq_ids, failed_drags_pair=True, reverse=False):
        '''The initiator.

        seq_ids - An iterator with the sequence ids to keep
        reverse - if True keep the sequences not found on the list
        '''
        if not isinstance(seq_ids, set):
            seq_ids = set(seq_ids)
        self.seq_ids = seq_ids
        super(FilterById, self).__init__(failed_drags_pair=failed_drags_pair,
                                              reverse=reverse)

    def _do_check(self, seq):
        return True if get_name(seq) in self.seq_ids else False class FilterByQuality(_BaseFilter):
    'It removes the sequences according to its quality'
    def __init__(self, threshold, ignore_masked=False, failed_drags_pair=True,
                 reverse=False):
        '''The initiator.

        threshold - minimum quality to pass the filter (float)
        reverse - if True keep the sequences not found on the list
        '''
        self.threshold = float(threshold)
        self.ignore_masked = ignore_masked
        super(FilterByQuality, self).__init__(reverse=reverse,
                                           failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        seq_object = seq.object
        try:
            quals = seq_object.letter_annotations['phred_quality']
        except KeyError:
            msg = 'Some of the input sequences do not have qualities: {}'
            msg = msg.format(get_name(seq))
            raise WrongFormatError(msg)
        if self.ignore_masked:
            str_seq = str(seq_object.seq)
            seg_quals = [quals[segment[0]: segment[1] + 1]
                            for segment in get_uppercase_segments(str_seq)]
            qual = sum(sum(q) * len(q) for q in seg_quals) / len(quals)
        else:
            qual = sum(quals) / len(quals)
        return True if qual >= self.threshold else False class FilterBlastMatch(_BaseFilter):
    'It filters a seq if there is a match against a blastdb'
    def __init__(self, database, program, filters, dbtype=None,
                 failed_drags_pair=True, reverse=False):
        '''The initiator
            database: path to a file with seqs or a blast database
            filter_params:
                expect_threshold
                similarty treshlod
                min_length_percentaje
        '''
        self._blast_db = database
        self._blast_program = program
        self._filters = filters
        self._dbtype = dbtype
        super(FilterBlastMatch, self).__init__(reverse=reverse,
                                          failed_drags_pair=failed_drags_pair)

    def _setup_checks(self, filterpacket):
        seqs = [s for seqs in filterpacket[SEQS_PASSED]for s in seqs]
        self._matcher = Blaster(seqs, self._blast_db, dbtype=self._dbtype,
                                program=self._blast_program,
                                filters=self._filters)

    def _do_check(self, seq):
        segments = self._matcher.get_matched_segments(get_name(seq))
        return True if segments is None else False class FilterBlastShort(_BaseFilter):
    'It filters a seq if there is a match against the given oligos'
    def __init__(self, oligos, failed_drags_pair=True, reverse=False):
        self.oligos = oligos
        super(FilterBlastShort, self).__init__(reverse=reverse,
                                          failed_drags_pair=failed_drags_pair)

    def _setup_checks(self, filterpacket):
        seqs = [s for seqs in filterpacket[SEQS_PASSED]for s in seqs]

        # we create a blastdb for these reads and then we use the oligos
        # as the blast query
        db_fhand = write_seqs(seqs, file_format='fasta')
        db_fhand.flush()
        params = {'task': 'blastn-short', 'expect': '0.0001'}
        filters = [{'kind': 'score_threshold', 'score_key': 'identity',
                    'min_score': 87},
                   {'kind': 'min_length', 'min_num_residues': 13,
                    'length_in_query': False}]
        self._matcher = BlasterForFewSubjects(db_fhand.name, self.oligos,
                                             program='blastn', filters=filters,
                                             params=params,
                                             elongate_for_global=False)

    def _do_check(self, seq):
        segments = self._matcher.get_matched_segments_for_read(get_name(seq))
        return True if segments is None else False class FilterDustComplexity(_BaseFilter):
    'It filters a sequence according to its dust score'
    def __init__(self, threshold=get_setting('DEFATULT_DUST_THRESHOLD'),
                 reverse=False, failed_drags_pair=True):
        '''The initiator
        '''
        self._threshold = threshold
        super(FilterDustComplexity, self).__init__(reverse=reverse,
                                          failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        threshold = self._threshold
        dustscore = calculate_dust_score(seq)
        return True if dustscore < threshold else False def seq_to_filterpackets(seq_packets, group_paired_reads=False):
    'It yields packets suitable for the filters'

    for packet in seq_packets:
        if group_paired_reads:
            packet = list(group_pairs_by_name(packet))
        else:
            packet = list(group_pairs(packet, n_seqs_in_pair=1))
        yield {SEQS_PASSED: packet, SEQS_FILTERED_OUT: []} class FilterByRpkm(_BaseFilter):
    def __init__(self, read_counts, min_rpkm, reverse=False,
                 failed_drags_pair=True):
        '''The init

        read_counts its a dict:
            - keys are the sequence names
            - values should be dicts with length, mapped_reads and
            unmmapped_reads, counts
        '''
        self._read_counts = read_counts
        self._min_rpkm = min_rpkm
        total_reads = sum([v['mapped_reads'] + v['unmapped_reads'] for v in read_counts.values()])
        self._million_reads = total_reads / 1e6
        super(FilterByRpkm, self).__init__(reverse=reverse,
                                           failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        count = self._read_counts[get_name(seq)]

        kb_len = count['length'] / 1000
        rpk = count['mapped_reads'] / kb_len  # rpks
        rpkm = rpk / self._million_reads  # rpkms

        return True if rpkm >= self._min_rpkm else False class FilterByBam(FilterById):
    'It filters the reads not mapped in the given BAM files'
    def __init__(self, bam_fpaths, min_mapq=0, reverse=False):
        seq_ids = self._get_mapped_reads(bam_fpaths, min_mapq)
        super(FilterByBam, self).__init__(seq_ids, reverse=reverse)

    def _get_mapped_reads(self, bam_fpaths, min_mapq):
        mapped_reads = []
        for fpath in bam_fpaths:
            mapped_reads.extend(_get_mapped_reads(fpath, min_mapq=min_mapq))
        return mapped_reads class FilterAllNs(_BaseFilter):
    'It filters a sequence completely composed by Ns'
    def __init__(self, reverse=False, failed_drags_pair=True):
        super(FilterAllNs, self).__init__(reverse=reverse,
                                          failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        seq = get_str_seq(seq)
        if not seq:
            return True
        chars = set(seq)
        good_chars = chars.difference(set(('N', 'n', '-', '*')))
        if good_chars:
            return True
        else:
            return False class FilterBowtie2Match(_BaseFilter):
    'It filters a seq if it maps against a bowtie2 index'
    def __init__(self, index_fpath, reverse=False, min_mapq=None,
                 failed_drags_pair=True, threads=None):
        self._index_fpath = index_fpath
        self._reverse = reverse
        self.min_mapq = min_mapq
        self.threads = threads
        super(FilterBowtie2Match, self).__init__(reverse=reverse,
                                          failed_drags_pair=failed_drags_pair)

    def _setup_checks(self, filterpacket):
        index_fpath = self._index_fpath
        seqs = [s for seqs in filterpacket[SEQS_PASSED]for s in seqs]
        seq_class = seqs[0].kind
        extra_params = []
        # Which format do we need for the bowtie2 input read file fasta or
        # fastq?
        if seq_class == SEQRECORD:
            if 'phred_quality' in seqs[0].object.letter_annotations.viewkeys():
                file_format = 'fastq'
            else:
                extra_params.append('-f')
                file_format = 'fasta'
        elif seq_class == SEQITEM:
            file_format = get_file_format(seqs[0])
            if 'illumina' in file_format:
                extra_params.append('--phred64')
            elif 'fasta' in file_format:
                extra_params.append('-f')
            elif 'fastq' in file_format:
                pass
            else:
                msg = 'For FilterBowtie2Match and SeqItems fastq or fasta '
                msg += 'files are required'
                raise RuntimeError(msg)
        else:
            raise NotImplementedError()

        reads_fhand = NamedTemporaryFile(suffix=file_format)
        write_seqs(seqs, reads_fhand, file_format=file_format)
        reads_fhand.flush()

        bam_fhand = NamedTemporaryFile(suffix='.bam')
        map_process = map_with_bowtie2(index_fpath,
                                       unpaired_fpath=reads_fhand.name,
                                       extra_params=extra_params,
                                       threads=self.threads)
        map_process_to_bam(map_process, bam_fhand.name)

        self.mapped_reads = _get_mapped_reads(bam_fhand.name, self.min_mapq)

    def _do_check(self, seq):
        return False if get_name(seq) in self.mapped_reads else True class FilterByFeatureTypes(_BaseFilter):
    'It filters out sequences not annotated with the given feature types'
    def __init__(self, feature_types, reverse=False, failed_drags_pair=True):
        '''The initiator

        feat_types is a list of types of features to use to filter'''

        self._feat_types = feature_types
        super(FilterByFeatureTypes, self).__init__(reverse=reverse,
                                           failed_drags_pair=failed_drags_pair)

    def _do_check(self, seq):
        seq = seq.object
        f_in_seq = [f.type for f in seq.features if f.type in self._feat_types]
        return True if f_in_seq else False BIN_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..',
                                       'bin')) TEST_DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                             '..', 'test', 'test_data')) NUCL = 'nucl' SEQS_FILTERED_OUT = 'seqs_filtered_out' SEQS_PASSED = 'seqs_passed' SEQITEM = 'seqitem' SEQRECORD = 'seqrecord' class SeqItem(_SeqItem):
    def __new__(cls, name, lines, annotations=None):
def get_title(seq):
def get_description(seq):
def get_name(seq):
def get_file_format(seq):
def _break():
def _is_fastq_plus_line(line, seq_name):
def _get_seqitem_quals(seq):
def get_str_seq(seq):
def get_length(seq):
def _get_seqitem_qualities(seqwrap):
def get_int_qualities(seq):
def _int_quals_to_str_quals(int_quals, out_format):
def get_str_qualities(seq, out_format=None):
def get_annotations(seq):
def _copy_seqrecord(seqrec, seq=None, name=None, id_=None):
def _copy_seqitem(seqwrapper, seq=None, name=None):
def copy_seq(seqwrapper, seq=None, name=None):
def _slice_seqitem(seqwrap, start, stop):
def slice_seq(seq, start=None, stop=None):
def assing_kind_to_seqs(kind, seqs, file_format):
SANGER_QUALS = {chr(i): i - 33 for i in range(33, 127)}
ILLUMINA_QUALS = {chr(i): i - 64 for i in range(64, 127)}
SANGER_STRS = {i - 33: chr(i)  for i in range(33, 127)}
ILLUMINA_STRS = {i - 64: chr(i) for i in range(64, 127)} def read_seq_packets(fhands, size=get_setting('PACKET_SIZE'), out_format=None,
                     prefered_seq_classes=None):
    '''It yields SeqItems in packets of the given size.'''
    seqs = read_seqs(fhands, out_format=out_format,
                     prefered_seq_classes=prefered_seq_classes)
    return group_in_packets(seqs, size)
import unittest
import os.path
from string import ascii_lowercase
from random import choice
from subprocess import check_output, call, CalledProcessError
from tempfile import NamedTemporaryFile
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
from Bio.SeqFeature import SeqFeature, FeatureLocation
from crumbs.seq.filters import (FilterByLength, FilterById, FilterByQuality,
                                FilterBlastMatch, FilterBlastShort,
                                FilterDustComplexity, seq_to_filterpackets,
                                FilterByRpkm, FilterByBam, FilterAllNs,
                                FilterBowtie2Match, FilterByFeatureTypes)
from crumbs.utils.bin_utils import BIN_DIR
from crumbs.utils.test_utils import TEST_DATA_DIR
from crumbs.utils.tags import (NUCL, SEQS_FILTERED_OUT, SEQS_PASSED, SEQITEM,
                               SEQRECORD)
from crumbs.seq.seq import get_name, get_str_seq, SeqWrapper
from crumbs.seq.seqio import read_seq_packets
        result = check_output([filter_bin, '-b', blastdb, '-s', '99',
                               seq_fhand.name])
        assert 'CATGAACACACACAT' in result

        # fail if -a an -l given i in the command line
        stderr = NamedTemporaryFile()
        try:
            check_output([filter_bin, '-b', blastdb, seq_fhand.name, '-a', '3',
                          '-l', '4'], stderr=stderr)
            self.fail()
        except CalledProcessError:
            assert 'not allowed with argument' in open(stderr.name).read()

        # minlen percentaje_filter
        filtered_fhand = NamedTemporaryFile()
        result = check_output([filter_bin, '-b', blastdb, '-l', '40',
                               seq_fhand.name, '-e', filtered_fhand.name])
        assert ">seq\nATCATGTAGTTACACATG" in open(filtered_fhand.name).read()
        assert result == ''

        result = check_output([filter_bin, '-b', blastdb, '-l', '80',
                               seq_fhand.name])
        assert 'CATGAACACACACAT' in result

        # min nucleotides in query filter
        result = check_output([filter_bin, '-b', blastdb, '-a', '40',
                               seq_fhand.name])
        assert result == ''

        result = check_output([filter_bin, '-b', blastdb, '-a', '80',
                               seq_fhand.name])
        assert 'CATGAACACACACAT' in result

        # reverse
        filtered_fhand = NamedTemporaryFile()
        result = check_output([filter_bin, '-b', blastdb, '-a', '80', '-r',
                               seq_fhand.name, '-e', filtered_fhand.name])
        assert result == ''
        assert ">seq\nATCATGTAGTTACACATG" in open(filtered_fhand.name).read()

    def test_filter_blast_bin_fastq(self):
        filter_bin = os.path.join(BIN_DIR, 'filter_by_blast')
        blastdb = os.path.join(TEST_DATA_DIR, 'blastdbs', 'arabidopsis_genes')
        # With fastq
        match = 'CCAAAGTACGGTCTCACAAGCGGTCTCTTACCGGACACCGTCACCGATTTCACCCTCT'
        seq = 'ATCATGTAGTTACACATGAACACACACATG'
        seq += match

        fastq = '@seq\n' + seq + '\n+\n' + 'a' * len(seq) + '\n'
        seq_fhand = _make_fhand(fastq + fastq)

        result = check_output([filter_bin, '-b', blastdb, '-x', '1e-27',
                               seq_fhand.name])

        assert 'CATGAACACACACAT' in result


class BlastShortFilterTest(unittest.TestCase):
    'It tests the filtering by blast-short match'
    @staticmethod
    def test_blastmatch_filter():
        'it test filter by blast'
        seq = 'CCAAAGTACGGTCTCCCAAGCGGTCTCTTACCGGACACCGTCACCGATTTCACCCTCT'
        oligo = 'GTGACTGGAGTTCAGACGTGTGCTCTTCCGATCT'
        seq_oligo = seq + oligo
        oligo = SeqRecord(Seq(oligo))
        oligo = SeqWrapper(SEQRECORD, oligo, None)

        seq = SeqRecord(Seq(seq), id='seq')
        seq = SeqWrapper(object=seq, kind=SEQRECORD, file_format=None)

        seq_oligo = SeqRecord(Seq(seq_oligo), id='seq_oligo')
        seq_oligo = SeqWrapper(object=seq_oligo, kind=SEQRECORD,
                               file_format=None)

        seqs = {SEQS_PASSED: [[seq], [seq_oligo]], SEQS_FILTERED_OUT: []}

        filter_ = FilterBlastShort([oligo])
        filt_packet = filter_(seqs)
        passed = [get_name(pair[0]) for pair in filt_packet[SEQS_PASSED]]
        fail = [get_name(pair[0]) for pair in filt_packet[SEQS_FILTERED_OUT]]
        assert passed == ['seq']
        assert fail == ['seq_oligo']

    def test_filter_blast_bin(self):
        filter_bin = os.path.join(BIN_DIR, 'filter_by_blast_short')
        # With fastq
        seq = 'CCAAAGTACGGTCTCCCAAGCGGTCTCTTACCGGACACCGTCACCGATTTCACCCTCT'
        oligo = 'GTGACTGGAGTTCAGACGTGTGCTCTTCCGATCT'
        seq_oligo = seq + oligo

        fastq = '@seq_oligo\n' + seq_oligo + '\n+\n' + 'a' * len(seq_oligo)
        fastq += '\n@seq\n' + seq + '\n+\n' + 'a' * len(seq) + '\n'
        seq_fhand = _make_fhand(fastq)

        result = check_output([filter_bin, '-l', oligo, seq_fhand.name])

        assert '@seq' in result
        assert 'oligo' not in result

        # with oligo file
        oligo_fhand = NamedTemporaryFile()
        oligo_fhand.write('>oligo\nGTGACTGGAGTTCAGACGTGTGCTCTTCCGATCT\n')
        oligo_fhand.flush()
        result = check_output([filter_bin, '-f', oligo_fhand.name,
                               seq_fhand.name])
        assert '@seq' in result
        assert 'oligo' not in result


class ComplexityFilterTest(unittest.TestCase):
    'It tests the filtering by complexity'
    @staticmethod
    def test_dust_filter():
        'It tests the complexity filter'
        seq1 = 'TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTAAAAAAAAAAAAAAAAAAAAAAAAA'
        seq2 = 'CATCGATTGCGATCGATCTTGTTGCACGACTAGCTATCGATTGCTAGCTTAGCTAGCTAGTT'
        seq1 = SeqRecord(Seq(seq1), id='seq1')
        seq2 = SeqRecord(Seq(seq2), id='seq2')
        seq1 = SeqWrapper(SEQRECORD, seq1, None)

        seq2 = SeqWrapper(SEQRECORD, seq2, None)