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)