<s> """<STR_LIT>""" <EOL> import os <EOL> import os . path as op <EOL> from pbcore . io import FastaReader , FastaWriter <EOL> from pbtools . pbtranscript . Utils import mkdir <EOL> class FastaSplitter ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , input_fasta , reads_per_split , out_dir , out_prefix ) : <EOL> self . input_fasta = input_fasta <EOL> self . out_dir = out_dir <EOL> self . reads_per_split = reads_per_split <EOL> self . out_prefix = out_prefix <EOL> self . out_fns = None <EOL> mkdir ( self . out_dir ) <EOL> def __str__ ( self ) : <EOL> if self . out_fns is None or len ( self . out_fns ) == <NUM_LIT:0> : <EOL> return "<STR_LIT>" . format ( input_fasta = self . input_fasta ) + "<STR_LIT>" + "<STR_LIT>" . format ( n = self . reads_per_split ) <EOL> else : <EOL> return "<STR_LIT>" . format ( input_fasta = self . input_fasta ) + "<STR_LIT>" . format ( m = len ( self . out_fns ) , <EOL> n = self . reads_per_split ) + "<STR_LIT:;>" . join ( self . out_fns ) <EOL> def _out_fn ( self , split_index ) : <EOL> """<STR_LIT>""" <EOL> if split_index > <NUM_LIT> : <EOL> raise ValueError ( "<STR_LIT>" + <EOL> "<STR_LIT>" ) <EOL> name = "<STR_LIT>" . format ( prefix = self . out_prefix , <EOL> idx = split_index ) <EOL> return op . join ( self . out_dir , name ) <EOL> def split ( self , first_split = None ) : <EOL> """<STR_LIT>""" <EOL> split_index = <NUM_LIT:0> <EOL> self . out_fns = [ ] <EOL> writer = FastaWriter ( self . _out_fn ( split_index ) ) <EOL> self . out_fns . append ( self . _out_fn ( split_index ) ) <EOL> if first_split is None : <EOL> first_split = self . reads_per_split <EOL> with FastaReader ( self . input_fasta ) as reader : <EOL> for ridx , r in enumerate ( reader ) : <EOL> if ( ( split_index == <NUM_LIT:0> and ridx == first_split ) or ( split_index > <NUM_LIT:0> and ridx % self . reads_per_split == <NUM_LIT:0> ) ) and ridx != <NUM_LIT:0> : <EOL> split_index += <NUM_LIT:1> <EOL> writer . close ( ) <EOL> writer = FastaWriter ( self . _out_fn ( split_index ) ) <EOL> self . out_fns . append ( self . _out_fn ( split_index ) ) <EOL> writer . writeRecord ( r . name , r . sequence ) <EOL> writer . close ( ) <EOL> return list ( self . out_fns ) <EOL> def rmOutFNs ( self ) : <EOL> """<STR_LIT>""" <EOL> for f in self . out_fns : <EOL> os . remove ( f ) <EOL> self . out_fns = [ ] <EOL> def splitFasta ( input_fasta , reads_per_split , out_dir , out_prefix , first_split = None ) : <EOL> """<STR_LIT>""" <EOL> obj = FastaSplitter ( input_fasta = input_fasta , <EOL> reads_per_split = reads_per_split , <EOL> out_dir = out_dir , out_prefix = out_prefix ) <EOL> return obj . split ( first_split ) <EOL> def get_args ( ) : <EOL> """<STR_LIT>""" <EOL> import argparse <EOL> parser = argparse . ArgumentParser ( <EOL> description = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , <EOL> type = str , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , <EOL> type = int , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , <EOL> type = str , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , <EOL> type = str , <EOL> help = "<STR_LIT>" ) <EOL> this_args = parser . parse_args ( ) <EOL> return this_args <EOL> def main ( ) : <EOL> """<STR_LIT>""" <EOL> import logging <EOL> from pbtools . pbtranscript . __init__ import get_version <EOL> log = logging . getLogger ( __name__ ) <EOL> args = get_args ( ) <EOL> from pbtools . pbtranscript . Utils import setup_log <EOL> setup_log ( alog = log , level = logging . DEBUG ) <EOL> log . info ( "<STR_LIT>" . format ( f = op . basename ( __file__ ) , <EOL> v = get_version ( ) ) ) <EOL> splitFasta ( input_fasta = args . input_fasta , <EOL> reads_per_split = args . reads_per_split , <EOL> out_dir = args . out_dir , <EOL> out_prefix = args . out_prefix ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import sys <EOL> sys . exit ( main ( ) ) </s>
