def test_set_info(self, field_name, value, reader, expected):
if reader is not None:
reader = mock.Mock()
reader.field_access_cache.info_field_set_fn.return_value = (
struct_utils.set_string_field)
variant = variants_pb2.Variant()
variant_utils.set_info(variant, field_name, value, reader)
actual = variant.info[field_name].values
self.assertEqual(len(actual), len(expected))
for actual_elem, expected_elem in zip(actual, expected):
self.assertEqual(actual_elem, expected_elem)
def main(argv):
if len(argv) != 3:
print('Usage: %s <input_vcf> <output_vcf>' % argv[0])
sys.exit(-1)
in_vcf = argv[1]
out_vcf = argv[2]
with vcf.VcfReader(in_vcf) as reader:
if 'AD' in [info.id for info in reader.header.infos]:
print('%s already contains AD field.' % in_vcf)
sys.exit(-1)
out_header = reader.header
out_header.infos.extend([vcf_constants.reserved_info_field('AD')])
with vcf.VcfWriter(out_vcf, header=out_header) as writer:
for variant in reader:
variant_utils.set_info(variant, 'AD', get_variant_ad(variant),
writer)
writer.write(variant)
def test_get_info(self, field_name, reader, expected):
if reader is not None:
reader = mock.Mock()
reader.field_access_cache.info_field_get_fn.return_value = (
functools.partial(
struct_utils.get_string_field, is_single_field=True))
variant = variants_pb2.Variant()
variant_utils.set_info(variant, 'AD', [23, 25])
variant_utils.set_info(variant, 'AA', 'C')
variant_utils.set_info(variant, '1000G', True)
variant_utils.set_info(variant, 'DB', False)
actual = variant_utils.get_info(variant, field_name, vcf_object=reader)
self.assertEqual(actual, expected)
