<s> '''<STR_LIT>''' <EOL> import sys <EOL> import os <EOL> import gzip <EOL> from ngsutils . support import revcomp , FASTA <EOL> from ngsutils . fastq import FASTQ <EOL> import swalign <EOL> sw = swalign . LocalAlignment ( swalign . NucleotideScoringMatrix ( <NUM_LIT:2> , - <NUM_LIT:1> ) ) <EOL> def fastx_barcode_split ( reader , outtempl , barcodes , edits = <NUM_LIT:0> , pos = <NUM_LIT:0> , allow_revcomp = False , gzip_output = False , stats_fname = None ) : <EOL> '''<STR_LIT>''' <EOL> outs = { } <EOL> if gzip_output : <EOL> outtempl += '<STR_LIT>' <EOL> outs [ '<STR_LIT>' ] = gzip . open ( outtempl % '<STR_LIT>' , '<STR_LIT:w>' ) <EOL> else : <EOL> outs [ '<STR_LIT>' ] = open ( outtempl % '<STR_LIT>' , '<STR_LIT:w>' ) <EOL> tag_count = { } <EOL> for tag in barcodes : <EOL> if gzip_output : <EOL> outs [ tag ] = gzip . open ( outtempl % tag , '<STR_LIT:w>' ) <EOL> else : <EOL> outs [ tag ] = open ( outtempl % tag , '<STR_LIT:w>' ) <EOL> tag_count [ tag ] = <NUM_LIT:0> <EOL> matched = <NUM_LIT:0> <EOL> perfect = <NUM_LIT:0> <EOL> mismatched = <NUM_LIT:0> <EOL> mispositioned = <NUM_LIT:0> <EOL> missing = <NUM_LIT:0> <EOL> for record in reader . fetch ( ) : <EOL> ismatch , ( tag , aln , is_forward , reason ) = check_tags ( barcodes , record . seq , edits , pos , allow_revcomp ) <EOL> if not ismatch : <EOL> missing += <NUM_LIT:1> <EOL> comment = '<STR_LIT>' % ( tag , barcodes [ tag ] [ <NUM_LIT:0> ] , '<STR_LIT>' if is_forward else '<STR_LIT>' , aln . q_pos , aln . extended_cigar_str , aln . mismatches , reason ) <EOL> record . subseq ( None , None , comment ) . write ( outs [ '<STR_LIT>' ] ) <EOL> else : <EOL> matched += <NUM_LIT:1> <EOL> if aln . mismatches == <NUM_LIT:0> : <EOL> if aln . q_pos == <NUM_LIT:0> or aln . q_end == len ( record . seq ) : <EOL> perfect += <NUM_LIT:1> <EOL> else : <EOL> mispositioned += <NUM_LIT:1> <EOL> else : <EOL> mismatched += <NUM_LIT:1> <EOL> comment = '<STR_LIT>' % ( tag , '<STR_LIT>' if is_forward else '<STR_LIT>' , aln . q_pos , aln . mismatches , ) <EOL> if not barcodes [ tag ] [ <NUM_LIT:2> ] : <EOL> record . write ( outs [ tag ] ) <EOL> else : <EOL> if is_forward : <EOL> if barcodes [ tag ] [ <NUM_LIT:1> ] == '<STR_LIT:5>' : <EOL> newrec = record . subseq ( aln . q_end , None , comment ) <EOL> else : <EOL> subseqlen = len ( barcodes [ tag ] [ <NUM_LIT:0> ] ) + edits + pos <EOL> newrec = record . subseq ( None , - <NUM_LIT:1> * ( subseqlen - aln . q_pos ) , comment ) <EOL> else : <EOL> if barcodes [ tag ] [ <NUM_LIT:1> ] == '<STR_LIT:5>' : <EOL> subseqlen = len ( barcodes [ tag ] [ <NUM_LIT:0> ] ) + edits + pos <EOL> newrec = record . subseq ( None , - <NUM_LIT:1> * ( subseqlen - aln . q_pos ) , comment ) <EOL> else : <EOL> newrec = record . subseq ( aln . q_end , None , comment ) <EOL> newrec . write ( outs [ tag ] ) <EOL> tag_count [ tag ] += <NUM_LIT:1> <EOL> for tag in outs : <EOL> outs [ tag ] . close ( ) <EOL> if stats_fname : <EOL> with open ( stats_fname , '<STR_LIT:w>' ) as f : <EOL> f . write ( "<STR_LIT>" % missing ) <EOL> f . write ( "<STR_LIT>" % matched ) <EOL> f . write ( "<STR_LIT>" % perfect ) <EOL> f . write ( "<STR_LIT>" % mismatched ) <EOL> f . write ( "<STR_LIT>" % mispositioned ) <EOL> f . write ( '<STR_LIT:\n>' ) <EOL> for tag in tag_count : <EOL> f . write ( "<STR_LIT>" % ( tag , tag_count [ tag ] ) ) <EOL> def _tag_aln_check ( aln , seqlen , barcodelen , orientation , edit = <NUM_LIT:0> , pos = <NUM_LIT:0> ) : <EOL> r_len = aln . r_end - aln . r_pos <EOL> mismatches = barcodelen - r_len + aln . mismatches <EOL> if mismatches <= edit : <EOL> if orientation == '<STR_LIT:5>' and aln . q_pos <= pos : <EOL> return True , '<STR_LIT>' <EOL> elif orientation == '<STR_LIT:3>' and aln . q_end >= seqlen - pos : <EOL> return True , '<STR_LIT>' <EOL> else : <EOL> return False , '<STR_LIT>' <EOL> else : <EOL> return False , '<STR_LIT>' <EOL> return False , '<STR_LIT:?>' <EOL> def check_tags ( barcodes , seq , edit , pos , allow_revcomp = False , verbose = False ) : <EOL> '''<STR_LIT>''' <EOL> best = None <EOL> for tag in barcodes : <EOL> barcodeseq , orientation , strip = barcodes [ tag ] <EOL> if orientation == '<STR_LIT:5>' : <EOL> testseq = seq [ : len ( barcodeseq ) + edit + pos ] <EOL> else : <EOL> testseq = seq [ - <NUM_LIT:1> * ( len ( barcodeseq ) + edit + pos ) : ] <EOL> aln = sw . align ( barcodeseq , testseq ) <EOL> valid , reason = _tag_aln_check ( aln , len ( testseq ) , len ( barcodeseq ) , orientation , edit , pos ) <EOL> if verbose : <EOL> print '<STR_LIT>' % ( str ( barcodes [ tag ] ) , testseq ) <EOL> aln . dump ( ) <EOL> print valid , reason <EOL> if valid : <EOL> return True , ( tag , aln , True , '<STR_LIT>' ) <EOL> if not best or aln . score > best [ <NUM_LIT:1> ] . score : <EOL> best = ( tag , aln , True , reason ) <EOL> if allow_revcomp : <EOL> if orientation == '<STR_LIT:5>' : <EOL> testseq = seq [ - <NUM_LIT:1> * ( len ( barcodeseq ) + edit + pos ) : ] <EOL> else : <EOL> testseq = seq [ : len ( barcodeseq ) + edit + pos ] <EOL> aln = sw . align ( revcomp ( barcodeseq ) , testseq ) <EOL> valid , reason = _tag_aln_check ( aln , len ( testseq ) , len ( barcodeseq ) , '<STR_LIT:5>' if orientation == '<STR_LIT:3>' else '<STR_LIT:3>' , edit , pos ) <EOL> if verbose : <EOL> print '<STR_LIT>' % ( str ( barcodes [ tag ] ) , testseq ) <EOL> aln . dump ( ) <EOL> print valid , reason <EOL> if valid : <EOL> return True , ( tag , aln , False , '<STR_LIT>' ) <EOL> if not best or aln . score > best [ <NUM_LIT:1> ] . score : <EOL> best = ( tag , aln , False , reason ) <EOL> if verbose : <EOL> print '<STR_LIT>' , best <EOL> best [ <NUM_LIT:1> ] . dump ( ) <EOL> return False , best <EOL> def usage ( ) : <EOL> sys . stdout . write ( '<STR_LIT>' % __doc__ ) <EOL> sys . stdout . write ( '''<STR_LIT>''' % os . path . basename ( sys . argv [ <NUM_LIT:0> ] ) ) <EOL> sys . exit ( <NUM_LIT:1> ) <EOL> def is_fasta_file ( fname ) : <EOL> '''<STR_LIT>''' <EOL> if fname [ - <NUM_LIT:3> : ] == '<STR_LIT>' : <EOL> f = gzip . open ( fname ) <EOL> else : <EOL> f = open ( fname ) <EOL> try : <EOL> while True : <EOL> line = f . next ( ) <EOL> if line [ <NUM_LIT:0> ] == '<STR_LIT:>>' : <EOL> f . close ( ) <EOL> return True <EOL> elif line [ <NUM_LIT:0> ] == '<STR_LIT:@>' : <EOL> f . close ( ) <EOL> return False <EOL> except : <EOL> f . close ( ) <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> barcodes = None <EOL> fname = None <EOL> templ = None <EOL> edit = <NUM_LIT:0> <EOL> pos = <NUM_LIT:0> <EOL> allow_revcomp = False <EOL> gz = False <EOL> stats = False <EOL> last = None <EOL> for arg in sys . argv [ <NUM_LIT:1> : ] : <EOL> if last == '<STR_LIT>' : <EOL> edit = int ( arg ) <EOL> last = None <EOL> elif last == '<STR_LIT>' : <EOL> pos = int ( arg ) <EOL> last = None <EOL> elif arg in [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> last = arg <EOL> elif arg == '<STR_LIT>' : <EOL> allow_revcomp = True <EOL> elif arg == '<STR_LIT>' : <EOL> gz = True <EOL> elif arg == '<STR_LIT>' : <EOL> stats = True <EOL> elif not barcodes and os . path . exists ( arg ) : <EOL> barcodes = { } <EOL> with open ( arg ) as f : <EOL> first = True <EOL> for line in f : <EOL> if first : <EOL> first = False <EOL> continue <EOL> cols = line . strip ( ) . split ( '<STR_LIT:\t>' ) <EOL> if len ( cols ) < <NUM_LIT:4> : <EOL> cols . append ( '<STR_LIT:y>' ) <EOL> barcodes [ cols [ <NUM_LIT:0> ] ] = ( cols [ <NUM_LIT:1> ] , cols [ <NUM_LIT:2> ] , True if cols [ <NUM_LIT:3> ] in '<STR_LIT>' else False ) <EOL> elif not fname and os . path . exists ( arg ) : <EOL> fname = arg <EOL> elif not templ : <EOL> templ = arg <EOL> if not barcodes or not fname or not templ : <EOL> usage ( ) <EOL> if is_fasta_file ( fname ) : <EOL> outtempl = '<STR_LIT>' % templ <EOL> reader = FASTA ( fname ) <EOL> else : <EOL> outtempl = '<STR_LIT>' % templ <EOL> reader = FASTQ ( fname ) <EOL> if stats : <EOL> stats_fname = '<STR_LIT>' % templ <EOL> else : <EOL> stats_fname = None <EOL> fastx_barcode_split ( reader , outtempl , barcodes , edit , pos , allow_revcomp , gz , stats_fname ) </s>
