class EstscanOrfAnnotator(object):
    'It annotates the given seqrecords'
    def __init__(self, usage_matrix):
        'Initiator'
        self._usage_matrix = usage_matrix

    def __call__(self, seqs):
        'It runs the actual annotations'
        if not seqs:
            return seqs
        pep_fhand = NamedTemporaryFile()
        dna_fhand = NamedTemporaryFile()
        _run_estscan(seqs, pep_fhand.name, dna_fhand.name,
                     self._usage_matrix)
        # now we read the result files
        estscan_result = _read_estcan_results(open(pep_fhand.name),
                                              open(dna_fhand.name))
        for seq in seqs:
            seq_name = get_name(seq)
            orfs = estscan_result.get(seq_name, {})
            feats = []
            for (start, end, strand), str_seqs in orfs.viewitems():
                start -= 1
                # end is fine  -- end[
                feat = SeqFeature(location=FeatureLocation(start, end, strand),
                                  type='ORF', qualifiers=str_seqs)
                feats.append(feat)
            if feats:
                seq.object.features.extend(feats)

        dna_fhand.close()
        pep_fhand.close()
        return seqs def _detect_polya_tail(seq, location, min_len, max_cont_mismatches):
    '''It detects 3' poylA or 5' polyT tails.

    This function is a re-implementation of the EMBOSS's trimest code.
    It will return the position of a poly-A in 3' or a poly-T in 5'.
    It returns the start and end of the tail. The nucleotide in the end
    position won't be included in the poly-A.
    '''
    if location == FIVE_PRIME:
        tail_nucl = 'T'
        inc = 1
        start = 0
        end = len(seq)
    elif location == THREE_PRIME:
        tail_nucl = 'A'
        inc = -1
        start = -1
        end = -len(seq) - 1
    else:
        msg = 'location should be five or three prime'
        raise ValueError(msg)

    mismatch_count = 0
    poly_count = 0
    tail_len = 0
    result = 0

    for index in range(start, end, inc):
        nucl = seq[index].upper()
        if nucl == tail_nucl:
            poly_count += 1
            mismatch_count = 0
        elif nucl == 'N':
            pass
        else:
            poly_count = 0
            mismatch_count += 1
        if poly_count >= min_len:
            result = tail_len + 1
        if mismatch_count > max_cont_mismatches:
            break
        tail_len += 1
    if result and location == FIVE_PRIME:
        start = 0
        end = result
        result = start, end
    elif result and location == THREE_PRIME:
        end = len(seq)
        start = end - result
        result = start, end
    else:
        result = None
    return result class PolyaAnnotator(object):
    'It annotates the given seqrecords with poly-A or poly-T regions'
    def __init__(self, min_len=get_setting('POLYA_ANNOTATOR_MIN_LEN'),
                max_cont_mismatches=get_setting('POLYA_ANNOTATOR_MISMATCHES')):
        '''It inits the class.

        min_len - minimum number of consecutive As (or Ts) to extend the tail
        max_cont_mismatches - maximum number of consecutive no A (or Ts) to
                              break a tail.
        '''
        self._min_len = min_len
        self._max_cont_mismatches = max_cont_mismatches

    def __call__(self, seqrecords):
        'It runs the actual annotations'
        max_cont_mismatches = self._max_cont_mismatches
        min_len = self._min_len

        for seq in seqrecords:
            _annotate_polya(seq, min_len, max_cont_mismatches)
        return seqrecords class BlastAnnotator(object):
    'It annotates using blast'
    def __init__(self, blastdb, program, dbtype=None, filters=None,
                 params=None, remote=False):
        'Initializes the class'
        self.blastdb = blastdb
        self._program = program
        self._filters = [] if filters is None else filters
        self._params = params
        self._dbtype = dbtype
        self._remote = remote

    def __call__(self, seqrecords):
        'It does the work'
        if not seqrecords:
            return seqrecords
        matcher = Blaster(seqrecords, self.blastdb, self._program,
                               self._dbtype, filters=self._filters,
                               params=self._params, remote=self._remote)
        blasts = matcher.blasts
        blastdb = os.path.basename(self.blastdb)
        for seqrecord in seqrecords:
            align_result = blasts.get(get_name(seqrecord), None)
            if not align_result:
                continue
            match_counter = 0
            for match in align_result['matches']:
                subject = match['subject']['name']
                match_counter += 1
                for match_part in  match['match_parts']:
                    if match_part['subject_end'] < match_part['subject_start']:
                        strand = -1
                        subject_start = match_part['subject_end']
                        subject_end = match_part['subject_start']
                    else:
                        strand = 1
                        subject_start = match_part['subject_start']
                        subject_end = match_part['subject_end']

                    query_start = match_part['query_start']
                    query_end = match_part['query_end']
                    qualifiers = {}
                    qualifiers['Target'] = {'start': subject_start,
                                            'end': subject_end,
                                            'name': subject}
                    qualifiers['score'] = match_part['scores']['expect']
                    qualifiers['identity'] = match_part['scores']['identity']
                    qualifiers['blastdb'] = blastdb
                    location = FeatureLocation(query_start, query_end, strand)
                    feature = SeqFeature(location=location, type='match_part',
                                         qualifiers=qualifiers,
                                       id='match{0:03d}'.format(match_counter))
                    seqrecord.object.features.append(feature)
        return seqrecords TEST_DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',
                                             '..', 'test', 'test_data')) def read_seqs(fhands, out_format=None, prefered_seq_classes=None):
    'It returns a stream of seqs in different codings: seqrecords, seqitems...'

    if not prefered_seq_classes:
        prefered_seq_classes = [SEQITEM, SEQRECORD]
    try:
        in_format = get_format(fhands[0])
    except FileIsEmptyError:
        return []
    # seqitems is incompatible with different input and output formats
    # or when in_format != a fasta or fastq
    if ((out_format not in (None, GUESS_FORMAT) and in_format != out_format
         and SEQITEM in prefered_seq_classes) or
        (in_format not in ('fasta',) + SANGER_FASTQ_FORMATS +
         ILLUMINA_FASTQ_FORMATS)):
        prefered_seq_classes.pop(prefered_seq_classes.index(SEQITEM))

    if not prefered_seq_classes:
        msg = 'No valid seq class left or prefered'
        raise ValueError(msg)

    for seq_class in prefered_seq_classes:
        if seq_class == SEQITEM:
            try:
                return _read_seqitems(fhands)
            except NotImplementedError:
                continue
        elif seq_class == SEQRECORD:
            try:
                seqs = _read_seqrecords(fhands)
                return assing_kind_to_seqs(SEQRECORD, seqs, None)
            except NotImplementedError:
                continue
        else:
            raise ValueError('Unknown class for seq: ' + seq_class)
    raise RuntimeError('We should not be here, fixme') 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)} FIVE_PRIME = "5'" THREE_PRIME = "3'" NUCL = 'nucl' SEQRECORD = 'seqrecord'
import unittest
import os
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from crumbs.seq.annotation import (EstscanOrfAnnotator, _detect_polya_tail,
                                   PolyaAnnotator, BlastAnnotator)
from crumbs.utils.test_utils import TEST_DATA_DIR
from crumbs.seq.seqio import read_seqs
from crumbs.seq.seq import SeqWrapper
from crumbs.utils.tags import FIVE_PRIME, THREE_PRIME, NUCL, SEQRECORD
# Copyright 2012 Jose Blanca, Peio Ziarsolo, COMAV-Univ. Politecnica Valencia
# This file is part of ngs_crumbs.
# ngs_crumbs is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# ngs_crumbs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR  PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with ngs_crumbs. If not, see <http://www.gnu.org/licenses/>.

# pylint: disable=R0201
# pylint: disable=R0904



_wrap_seq = lambda seq: SeqWrapper(SEQRECORD, seq, None)


class AnnotationTest(unittest.TestCase):
    'Test annotator classes'

    def test_orf_annotator(self):
        'It tests orf annotator'
        fpath = os.path.join(TEST_DATA_DIR, 'orf_test.fasta')
        estscan_matrix = os.path.join(TEST_DATA_DIR,
                                      'Arabidopsis_thaliana.smat')
        seq_records = list(read_seqs([open(fpath)],
                                     prefered_seq_classes=[SEQRECORD]))
        orf_annotator = EstscanOrfAnnotator(estscan_matrix)
        seq_records = orf_annotator(seq_records)
        orf1 = seq_records[0].object.features[0]
        orf2 = seq_records[1].object.features[0]
        assert orf1.strand == 1
        assert orf1.location.start.position == 0
        assert orf1.location.end.position == 541
        assert orf2.strand == -1
        assert orf2.location.start.position == 0
        assert orf2.location.end.position == 541
        assert not seq_records[2].object.features

    def test_polya_annotator(self):
        'It annotates poly-A or poly-T regions'
        seq1 = SeqRecord(seq=Seq('atccgtcagcatcCAATAAAAA'), id='seq1')
        seq2 = SeqRecord(seq=Seq('TTTTcTTcatccgtcag'), id='seq2')
        seq3 = SeqRecord(seq=Seq('TTTTcTTatccgtcagcatcCAATAAAAA'), id='seq3')
        seqs = [_wrap_seq(seq1), _wrap_seq(seq2), _wrap_seq(seq3)]

        annotator = PolyaAnnotator(min_len=5, max_cont_mismatches=0)