<s> """<STR_LIT>""" <EOL> from collections import namedtuple <EOL> import json <EOL> import math <EOL> import unittest <EOL> Pair = namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class GvcfExpander ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , bin_size = <NUM_LIT:100> , <EOL> filter_ref_matches = False , <EOL> emit_ref_blocks = True ) : <EOL> self . bin_size = bin_size <EOL> self . filter_ref_matches = filter_ref_matches <EOL> self . emit_ref_blocks = emit_ref_blocks <EOL> self . current_key = None <EOL> self . binned_calls = [ ] <EOL> self . sample_refs = { } <EOL> def is_variant ( self , fields ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in fields and fields [ '<STR_LIT>' ] : <EOL> return True <EOL> return False <EOL> def is_snp ( self , fields ) : <EOL> """<STR_LIT>""" <EOL> if ( self . is_variant ( fields ) <EOL> and fields [ '<STR_LIT>' ] in [ '<STR_LIT:A>' , '<STR_LIT:C>' , '<STR_LIT>' , '<STR_LIT:T>' ] <EOL> and all ( alt in [ '<STR_LIT:A>' , '<STR_LIT:C>' , '<STR_LIT>' , '<STR_LIT:T>' ] <EOL> for alt in fields [ '<STR_LIT>' ] ) ) : <EOL> return True <EOL> return False <EOL> def get_start ( self , fields ) : <EOL> return int ( fields [ '<STR_LIT:start>' ] ) <EOL> def get_end ( self , fields ) : <EOL> if '<STR_LIT>' in fields : <EOL> return int ( fields [ '<STR_LIT>' ] ) <EOL> else : <EOL> return int ( fields [ '<STR_LIT:end>' ] ) <EOL> def compute_start_bin ( self , fields ) : <EOL> return int ( math . floor ( self . get_start ( fields ) / self . bin_size ) ) <EOL> def compute_end_bin ( self , fields ) : <EOL> return int ( math . floor ( self . get_end ( fields ) / self . bin_size ) ) <EOL> def map ( self , fields ) : <EOL> results = [ ] <EOL> start_bin = self . compute_start_bin ( fields ) <EOL> if self . is_variant ( fields ) : <EOL> end_bin = start_bin <EOL> else : <EOL> end_bin = self . compute_end_bin ( fields ) <EOL> for current_bin in range ( start_bin , end_bin + <NUM_LIT:1> ) : <EOL> key = '<STR_LIT>' % ( fields [ '<STR_LIT>' ] , str ( current_bin ) ) <EOL> results . append ( Pair ( key , fields ) ) <EOL> return results <EOL> def reduce ( self , pair ) : <EOL> expanded_calls = [ ] <EOL> if None is self . current_key : <EOL> self . current_key = pair . k <EOL> if pair . k != self . current_key : <EOL> expanded_calls = self . expand_binned_calls ( ) <EOL> self . current_key = pair . k <EOL> self . binned_calls = [ ] <EOL> self . sample_refs = { } <EOL> self . binned_calls . append ( pair . v ) <EOL> return expanded_calls <EOL> def finalize ( self ) : <EOL> return self . expand_binned_calls ( ) <EOL> def expand_binned_calls ( self ) : <EOL> expanded_calls = [ ] <EOL> current_bin = int ( self . current_key . split ( '<STR_LIT::>' ) [ <NUM_LIT:1> ] ) <EOL> calls = sorted ( self . binned_calls , <EOL> key = lambda k : ( int ( k [ '<STR_LIT:start>' ] ) , <EOL> len ( k . get ( '<STR_LIT>' , [ ] ) ) ) ) <EOL> for call in calls : <EOL> if self . is_variant ( call ) : <EOL> expanded_calls . append ( self . expand_variant ( call ) ) <EOL> else : <EOL> self . accumulate_block ( call ) <EOL> if self . emit_ref_blocks : <EOL> if self . compute_start_bin ( call ) == current_bin : <EOL> expanded_calls . append ( call ) <EOL> return expanded_calls <EOL> def expand_variant ( self , variant_call ) : <EOL> if not self . is_snp ( variant_call ) : <EOL> return variant_call <EOL> expansion_calls = [ ] <EOL> for sample_id in self . sample_refs . keys ( ) : <EOL> ref_call = self . sample_refs [ sample_id ] <EOL> if ( self . get_start ( ref_call ) <= self . get_start ( variant_call ) and <EOL> self . get_end ( ref_call ) >= self . get_start ( variant_call ) + <NUM_LIT:1> ) : <EOL> expansion_calls . extend ( ref_call [ '<STR_LIT>' ] ) <EOL> else : <EOL> del self . sample_refs [ sample_id ] <EOL> if self . filter_ref_matches : <EOL> variant_sample_names = [ call [ '<STR_LIT>' ] for call in <EOL> variant_call [ '<STR_LIT>' ] ] <EOL> variant_call [ '<STR_LIT>' ] . extend ( <EOL> [ call for call in expansion_calls <EOL> if call [ '<STR_LIT>' ] not in variant_sample_names ] <EOL> ) <EOL> else : <EOL> variant_call [ '<STR_LIT>' ] . extend ( expansion_calls ) <EOL> return variant_call <EOL> def accumulate_block ( self , ref_call ) : <EOL> self . sample_refs [ ref_call [ '<STR_LIT>' ] [ <NUM_LIT:0> ] [ '<STR_LIT>' ] ] = ref_call <EOL> class GvcfExpanderTest ( unittest . TestCase ) : <EOL> """<STR_LIT>""" <EOL> def test_is_variant ( self ) : <EOL> expander = GvcfExpander ( ) <EOL> self . assertTrue ( expander . is_variant ( json . loads ( self . snp_1 ) ) ) <EOL> self . assertTrue ( expander . is_variant ( json . loads ( self . snp_2 ) ) ) <EOL> self . assertTrue ( expander . is_variant ( json . loads ( self . insertion_1 ) ) ) <EOL> self . assertTrue ( expander . is_variant ( json . loads ( self . deletion_1 ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . ref_a ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . ref_b ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . ref_c ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . ref_d ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . ref_ambiguous ) ) ) <EOL> self . assertFalse ( expander . is_variant ( json . loads ( self . no_call_1 ) ) ) <EOL> def test_is_snp ( self ) : <EOL> expander = GvcfExpander ( ) <EOL> self . assertTrue ( expander . is_snp ( json . loads ( self . snp_1 ) ) ) <EOL> self . assertTrue ( expander . is_snp ( json . loads ( self . snp_2 ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . insertion_1 ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . deletion_1 ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . ref_a ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . ref_b ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . ref_c ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . ref_d ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . ref_ambiguous ) ) ) <EOL> self . assertFalse ( expander . is_snp ( json . loads ( self . no_call_1 ) ) ) <EOL> def test_mapper_variant ( self ) : <EOL> test_input = json . loads ( self . snp_1 ) <EOL> expected_output = json . loads ( self . snp_1 ) <EOL> expander = GvcfExpander ( ) <EOL> result = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> ( key , value ) = result [ <NUM_LIT:0> ] <EOL> self . assertEqual ( '<STR_LIT>' , key ) <EOL> self . assertDictEqual ( expected_output , value ) <EOL> def test_mapper_ref_block ( self ) : <EOL> test_input = json . loads ( self . ref_a ) <EOL> expected_output = json . loads ( self . ref_a ) <EOL> expander = GvcfExpander ( ) <EOL> result = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:3> , len ( result ) ) <EOL> self . assertEqual ( '<STR_LIT>' , result [ <NUM_LIT:0> ] . k ) <EOL> self . assertEqual ( '<STR_LIT>' , result [ <NUM_LIT:1> ] . k ) <EOL> self . assertEqual ( '<STR_LIT>' , result [ <NUM_LIT:2> ] . k ) <EOL> for i in range ( <NUM_LIT:0> , <NUM_LIT:3> ) : <EOL> self . assertDictEqual ( expected_output , result [ i ] . v ) <EOL> def test_mapper_no_call ( self ) : <EOL> test_input = json . loads ( self . no_call_1 ) <EOL> expected_output = json . loads ( self . no_call_1 ) <EOL> expander = GvcfExpander ( ) <EOL> result = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> ( key , value ) = result [ <NUM_LIT:0> ] <EOL> self . assertEqual ( '<STR_LIT>' , key ) <EOL> self . assertDictEqual ( expected_output , value ) <EOL> def test_mr_variant ( self ) : <EOL> test_input = json . loads ( self . snp_1 ) <EOL> expected_output = json . loads ( self . snp_1 ) <EOL> expander = GvcfExpander ( ) <EOL> result = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> ( key , value ) = result [ <NUM_LIT:0> ] <EOL> self . assertEqual ( '<STR_LIT>' , key ) <EOL> self . assertDictEqual ( expected_output , value ) <EOL> result = expander . reduce ( pair = result [ <NUM_LIT:0> ] ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> result = expander . finalize ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> value = result [ <NUM_LIT:0> ] <EOL> self . assertDictEqual ( expected_output , value ) <EOL> def test_mr_ref ( self ) : <EOL> test_input = json . loads ( self . ref_a ) <EOL> expected_output = json . loads ( self . ref_a ) <EOL> expander = GvcfExpander ( ) <EOL> pairs = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:3> , len ( pairs ) ) <EOL> result = expander . reduce ( pairs [ <NUM_LIT:0> ] ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> result = expander . reduce ( pairs [ <NUM_LIT:1> ] ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> value = result [ <NUM_LIT:0> ] <EOL> self . assertDictEqual ( expected_output , value ) <EOL> for i in range ( <NUM_LIT:2> , <NUM_LIT:3> ) : <EOL> result = expander . reduce ( pairs [ i ] ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> result = expander . finalize ( ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> def test_mr_no_call ( self ) : <EOL> test_input = json . loads ( self . no_call_1 ) <EOL> expected_output = json . loads ( self . no_call_1 ) <EOL> expander = GvcfExpander ( ) <EOL> result = expander . map ( fields = test_input ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> ( key , value ) = result [ <NUM_LIT:0> ] <EOL> self . assertEqual ( '<STR_LIT>' , key ) <EOL> self . assertDictEqual ( expected_output , value ) <EOL> result = expander . reduce ( pair = result [ <NUM_LIT:0> ] ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> result = expander . finalize ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> value = result [ <NUM_LIT:0> ] <EOL> self . assertDictEqual ( expected_output , value ) <EOL> def test_mr ( self ) : <EOL> for filter_ref_matches in [ True , False ] : <EOL> expander = GvcfExpander ( filter_ref_matches = filter_ref_matches ) <EOL> pairs = [ ] <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . ref_a ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . ref_b ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . ref_c ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . ref_ambiguous ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . snp_1 ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . snp_2 ) ) ) <EOL> pairs . extend ( expander . map ( fields = json . loads ( self . no_call_1 ) ) ) <EOL> self . assertEqual ( <NUM_LIT:11> , len ( pairs ) ) <EOL> pairs = sorted ( pairs ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:0> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:1> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:2> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:3> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:4> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:5> ] ) ) ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( expander . reduce ( pairs [ <NUM_LIT:6> ] ) ) ) <EOL> result = expander . reduce ( pairs [ <NUM_LIT:7> ] ) <EOL> self . assertEqual ( <NUM_LIT:7> , len ( result ) ) <EOL> self . assertIn ( json . loads ( self . ref_b ) , result ) <EOL> self . assertIn ( json . loads ( self . ref_c ) , result ) <EOL> self . assertIn ( json . loads ( self . ref_a ) , result ) <EOL> self . assertIn ( json . loads ( self . ref_ambiguous ) , result ) <EOL> self . assertIn ( json . loads ( self . no_call_1 ) , result ) <EOL> if filter_ref_matches : <EOL> self . assertIn ( json . loads ( self . expanded_snp_1_filtered ) , result ) <EOL> self . assertIn ( json . loads ( self . expanded_snp_2 ) , result ) <EOL> else : <EOL> self . assertIn ( json . loads ( self . expanded_snp_1 ) , result ) <EOL> self . assertIn ( json . loads ( self . expanded_snp_2 ) , result ) <EOL> for i in range ( <NUM_LIT:7> , <NUM_LIT:11> ) : <EOL> result = expander . reduce ( pairs [ i ] ) <EOL> self . assertEqual ( <NUM_LIT:0> , len ( result ) ) <EOL> def test_same_start ( self ) : <EOL> expander = GvcfExpander ( emit_ref_blocks = False ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> expander . reduce ( expander . map ( fields = json . loads ( '<STR_LIT>' ) ) [ <NUM_LIT:0> ] ) <EOL> result = expander . finalize ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , len ( result ) ) <EOL> self . assertEqual ( <NUM_LIT> , len ( result [ <NUM_LIT:0> ] [ '<STR_LIT>' ] ) ) <EOL> def setUp ( self ) : <EOL> self . maxDiff = None <EOL> self . ref_a = """<STR_LIT>""" <EOL> self . ref_b = """<STR_LIT>""" <EOL> self . ref_ambiguous = """<STR_LIT>""" <EOL> self . ref_c = """<STR_LIT>""" <EOL> self . ref_d = """<STR_LIT>""" <EOL> self . no_call_1 = """<STR_LIT>""" <EOL> self . snp_1 = """<STR_LIT>""" <EOL> self . expanded_snp_1 = """<STR_LIT>""" <EOL> self . expanded_snp_1_filtered = """<STR_LIT>""" <EOL> self . snp_2 = """<STR_LIT>""" <EOL> self . expanded_snp_2 = """<STR_LIT>""" <EOL> self . insertion_1 = """<STR_LIT>""" <EOL> self . deletion_1 = """<STR_LIT>""" <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
