<s> """<STR_LIT>""" <EOL> import os <EOL> import uuid <EOL> import shutil <EOL> import subprocess <EOL> import toolz as tz <EOL> import yaml <EOL> from bcbio import utils <EOL> from bcbio . distributed . transaction import file_transaction <EOL> from bcbio . log import logger <EOL> from bcbio . pipeline import config_utils <EOL> from bcbiovm . docker import remap <EOL> from bcbiovm . ship import pack as ship_n_pack <EOL> def prep_workdir ( pack , parallel , args ) : <EOL> """<STR_LIT>""" <EOL> if pack [ "<STR_LIT:type>" ] == "<STR_LIT>" : <EOL> workdir , remap_dict , new_args = _create_workdir_shared ( pack [ "<STR_LIT>" ] , args , parallel , pack [ "<STR_LIT>" ] ) <EOL> return workdir , new_args , _shared_finalizer ( new_args , workdir , remap_dict , parallel ) <EOL> elif pack [ "<STR_LIT:type>" ] == "<STR_LIT>" : <EOL> workdir , new_args = _unpack_s3 ( pack [ "<STR_LIT>" ] [ "<STR_LIT>" ] , args ) <EOL> datai , data = config_utils . get_dataarg ( new_args ) <EOL> if "<STR_LIT>" not in data : <EOL> data [ "<STR_LIT>" ] = { } <EOL> data [ "<STR_LIT>" ] [ "<STR_LIT>" ] = workdir <EOL> new_args [ datai ] = data <EOL> return workdir , new_args , ship_n_pack . send_run_integrated ( pack ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % pack ) <EOL> def prep_datadir ( pack , args ) : <EOL> if "<STR_LIT>" in pack : <EOL> return pack [ "<STR_LIT>" ] , args <EOL> elif pack [ "<STR_LIT:type>" ] == "<STR_LIT>" : <EOL> return _unpack_s3 ( pack [ "<STR_LIT>" ] [ "<STR_LIT>" ] , args ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % pack ) <EOL> def prep_systemconfig ( datadir , args ) : <EOL> """<STR_LIT>""" <EOL> default_system = os . path . join ( datadir , "<STR_LIT>" , "<STR_LIT>" ) <EOL> if not utils . file_exists ( default_system ) : <EOL> with open ( default_system , "<STR_LIT:w>" ) as out_handle : <EOL> _ , data = config_utils . get_dataarg ( args ) <EOL> out = { "<STR_LIT>" : tz . get_in ( [ "<STR_LIT>" , "<STR_LIT>" ] , data , { } ) } <EOL> yaml . safe_dump ( out , out_handle , default_flow_style = False , allow_unicode = False ) <EOL> def get_output ( target_file , pconfig ) : <EOL> """<STR_LIT>""" <EOL> if pconfig [ "<STR_LIT:type>" ] == "<STR_LIT>" : <EOL> keyname = "<STR_LIT>" % ( tz . get_in ( [ "<STR_LIT>" , "<STR_LIT>" ] , pconfig ) , os . path . basename ( target_file ) ) <EOL> _transfer_s3 ( target_file , keyname , tz . get_in ( [ "<STR_LIT>" , "<STR_LIT>" ] , pconfig ) ) <EOL> return target_file <EOL> else : <EOL> raise NotImplementedError ( "<STR_LIT>" % pconfig ) <EOL> def _transfer_s3 ( out_fname , keyname , bucket ) : <EOL> if not os . path . exists ( out_fname ) : <EOL> utils . safe_makedir ( os . path . dirname ( out_fname ) ) <EOL> with file_transaction ( out_fname ) as tx_out_fname : <EOL> subprocess . check_call ( [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , tx_out_fname , <EOL> "<STR_LIT>" , keyname , "<STR_LIT>" , bucket ] ) <EOL> def _unpack_s3 ( bucket , args ) : <EOL> """<STR_LIT>""" <EOL> local_dir = utils . safe_makedir ( os . path . join ( os . getcwd ( ) , bucket ) ) <EOL> remote_key = "<STR_LIT>" % bucket <EOL> def _get_s3 ( orig_fname , context , remap_dict ) : <EOL> """<STR_LIT>""" <EOL> if orig_fname . startswith ( remote_key ) : <EOL> if context [ <NUM_LIT:0> ] in [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> cur_dir = os . path . join ( local_dir , "<STR_LIT>" ) <EOL> else : <EOL> cur_dir = local_dir <EOL> for fname in utils . file_plus_index ( orig_fname ) : <EOL> out_fname = fname . replace ( remote_key , cur_dir ) <EOL> keyname = fname . replace ( remote_key + "<STR_LIT:/>" , "<STR_LIT>" ) <EOL> _transfer_s3 ( out_fname , keyname , bucket ) <EOL> return orig_fname . replace ( remote_key , cur_dir ) <EOL> else : <EOL> return orig_fname <EOL> new_args = remap . walk_files ( args , _get_s3 , { remote_key : local_dir } ) <EOL> return local_dir , new_args <EOL> def _remap_dict_shared ( workdir , new_workdir , args ) : <EOL> """<STR_LIT>""" <EOL> ignore_keys = set ( [ "<STR_LIT>" ] ) <EOL> out = { workdir : new_workdir } <EOL> def _update_remap ( fname , context , remap_dict ) : <EOL> """<STR_LIT>""" <EOL> if not fname . startswith ( tuple ( out . keys ( ) ) ) and context and context [ <NUM_LIT:0> ] not in ignore_keys : <EOL> dirname = os . path . normpath ( os . path . dirname ( fname ) ) <EOL> local_dir = utils . safe_makedir ( os . path . join ( new_workdir , "<STR_LIT>" , str ( len ( out ) ) ) ) <EOL> out [ dirname ] = local_dir <EOL> remap . walk_files ( args , _update_remap , { } ) <EOL> return out <EOL> def _create_workdir_shared ( workdir , args , parallel , tmpdir = None ) : <EOL> """<STR_LIT>""" <EOL> if not tmpdir : <EOL> return workdir , { } , args <EOL> else : <EOL> new_workdir = utils . safe_makedir ( os . path . join ( tmpdir , "<STR_LIT>" % uuid . uuid1 ( ) ) ) <EOL> remap_dict = _remap_dict_shared ( workdir , new_workdir , args ) <EOL> new_args = remap . walk_files ( args , _remap_copy_file ( parallel ) , remap_dict ) <EOL> return new_workdir , remap_dict , new_args <EOL> def is_required_resource ( context , parallel ) : <EOL> fresources = parallel . get ( "<STR_LIT>" ) <EOL> if not fresources : <EOL> return True <EOL> for fresource in fresources : <EOL> if context [ : len ( fresource ) ] == fresource : <EOL> return True <EOL> return False <EOL> def _remap_copy_file ( parallel ) : <EOL> """<STR_LIT>""" <EOL> def _do ( fname , context , orig_to_temp ) : <EOL> new_fname = remap . remap_fname ( fname , context , orig_to_temp ) <EOL> if os . path . isfile ( fname ) : <EOL> if is_required_resource ( context , parallel ) : <EOL> logger . info ( "<STR_LIT>" % ( context , fname ) ) <EOL> utils . safe_makedir ( os . path . dirname ( new_fname ) ) <EOL> for ext in [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> if os . path . exists ( fname + ext ) : <EOL> if not os . path . exists ( new_fname + ext ) : <EOL> shutil . copyfile ( fname + ext , new_fname + ext ) <EOL> else : <EOL> logger . info ( "<STR_LIT>" % ( context , fname ) ) <EOL> elif os . path . isdir ( fname ) : <EOL> utils . safe_makedir ( new_fname ) <EOL> return new_fname <EOL> return _do <EOL> def _shared_finalizer ( args , workdir , remap_dict , parallel ) : <EOL> """<STR_LIT>""" <EOL> def _do ( out ) : <EOL> if remap_dict : <EOL> new_remap_dict = { v : k for k , v in remap_dict . items ( ) } <EOL> new_out = ( remap . walk_files ( out , _remap_copy_file ( parallel ) , new_remap_dict ) <EOL> if out else None ) <EOL> if os . path . exists ( workdir ) : <EOL> shutil . rmtree ( workdir ) <EOL> return new_out <EOL> else : <EOL> return out <EOL> return _do </s>
