-
Notifications
You must be signed in to change notification settings - Fork 55
/
genotype.py
70 lines (62 loc) · 2.2 KB
/
genotype.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import sys
class Genotype(object):
'''
This class stores information about each sample.
'''
def __init__(self, variant, value_list):
'''
Initialize the class. All instances have a GT field,
but that is enforced in the Variant class.
'''
self.value_list = value_list
self.variant = variant
def __eq__(self, other):
return self.get_gt_string() == other.get_gt_string()
def set_format(self, field, value):
'''
Set information for an individual format field.
'''
if field in self.variant.format_set:
if field in self.variant.format_dict:
i = self.variant.format_dict[field]
self._set_value(i, value)
else:
num_fields = len(self.variant.format_dict)
self.variant.format_dict[field] = num_fields
self._set_value(num_fields, value)
else:
sys.stderr.write('\nError: invalid FORMAT field, \"' + field + '\"\n')
sys.exit(1)
def _set_value(self, index, value):
try:
self.value_list[index] = value
except IndexError:
self.value_list.extend(['.'] * (index - len(self.value_list)))
self.value_list.append(value)
def get_format(self, field):
'''
Get value of particular field key
'''
try:
return self.value_list[self.variant.format_dict[field]]
except IndexError:
if field != 'GT':
return '.'
else:
return './.'
def get_gt_string(self):
'''
Convert object back to string.
If some values are missing (at the end for example) they are printed out as
all format fields present in any Genotype instance in the Variant line
are tracked.
'''
g_list = list()
for f in self.variant.format_list:
if f.id in self.variant.format_dict:
value = self.get_format(f.id)
if type(value) == float:
g_list.append('%0.2f' % value)
else:
g_list.append(str(value))
return ':'.join(g_list)