Imported Upstream version 0.7
[pysam.git] / pysam / VCF.py
1 #
2 # Code to read, write and edit VCF files
3 #
4 # VCF lines are encoded as a dictionary with these keys (note: all lowercase):
5 # 'chrom':  string
6 # 'pos':    integer
7 # 'id':     string
8 # 'ref':    string
9 # 'alt':    list of strings
10 # 'qual':   integer
11 # 'filter': None (missing value), or list of keys (strings); empty list parsed as ["PASS"]
12 # 'info':   dictionary of values (see below)
13 # 'format': list of keys (strings)
14 # sample keys: dictionary of values (see below)
15 #
16 # The sample keys are accessible through vcf.getsamples()
17 #
18 # A dictionary of values contains value keys (defined in ##INFO or ##FORMAT lines) which map
19 # to a list, containign integers, floats, strings, or characters.  Missing values are replaced 
20 # by a particular value, often -1 or .
21 #
22 # Genotypes are not stored as a string, but as a list of 1 or 3 elements (for haploid and diploid samples),
23 # the first (and last) the integer representing an allele, and the second the separation character.
24 # Note that there is just one genotype per sample, but for consistency the single element is stored in a list.
25 #
26 # Header lines other than ##INFO, ##FORMAT and ##FILTER are stored as (key, value) pairs and are accessible
27 # through getheader()
28 #
29 # The VCF class can be instantiated with a 'regions' variable consisting of tuples (chrom,start,end) encoding
30 # 0-based half-open segments.  Only variants with a position inside the segment will be parsed.  A regions
31 # parser is available under parse_regions.
32 #
33 # When instantiated, a reference can be passed to the VCF class.  This may be any class that supports a
34 # fetch(chrom, start, end) method.
35 #
36 #
37 #
38 # NOTE: the position that is returned to Python is 0-based, NOT 1-based as in the VCF file.
39 #
40 #
41 #
42 # TODO:
43 #  only v4.0 writing is complete; alleles are not converted to v3.3 format
44 #
45
46 from collections import namedtuple, defaultdict
47 from operator import itemgetter
48 import sys, re, copy, bisect
49
50 import pysam
51
52 gtsRegEx = re.compile("[|/\\\\]")
53 alleleRegEx = re.compile('^[ACGTN]+$')
54
55 # Utility function.  Uses 0-based coordinates
56 def get_sequence(chrom, start, end, fa):
57     # obtain sequence from .fa file, without truncation
58     if end<=start: return ""
59     if not fa: return "N"*(end-start)
60     if start<0: return "N"*(-start) + get_sequence(chrom, 0, end, fa).upper()
61     sequence = fa.fetch(chrom, start, end).upper()
62     if len(sequence) < end-start: sequence += "N"*(end-start-len(sequence))
63     return sequence
64
65 # Utility function.  Parses a region string
66 def parse_regions( string ):
67     result = []
68     for r in string.split(','):
69         elts = r.split(':')
70         chrom, start, end = elts[0], 0, 3000000000
71         if len(elts)==1: pass
72         elif len(elts)==2:
73             if len(elts[1])>0:
74                 ielts = elts[1].split('-')
75                 if len(ielts) != 2: ValueError("Don't understand region string '%s'" % r)
76                 try:    start, end = int(ielts[0])-1, int(ielts[1])
77                 except: raise ValueError("Don't understand region string '%s'" % r)
78         else:
79             raise ValueError("Don't understand region string '%s'" % r)
80         result.append( (chrom,start,end) )
81     return result
82             
83
84 FORMAT = namedtuple('FORMAT','id numbertype number type description missingvalue')
85
86 ###########################################################################################################
87
88 # New class
89
90 ###########################################################################################################
91
92 class VCFRecord:
93     '''vcf record.
94
95     initialized from data and vcf meta 
96     '''
97     
98     data = None
99     vcf = None
100
101     def __init__(self, data, vcf):
102         self.data, self.vcf = data, vcf
103
104         if len(data) != len(self.vcf._samples):
105             self.error(str(data),
106                        self.BAD_NUMBER_OF_COLUMNS, 
107                        "expected %s for %s samples (%s), got %s" % \
108                            (len(self.vcf._samples), 
109                             len(self.vcf._samples), 
110                             self.vcf._samples, 
111                             len(data)))
112         
113     property contig:
114     def contig( self ): return self.data[0]
115
116     property pos:
117     def __get__( self ): 
118             return self.data.pos
119         
120     property id:
121         def __get__( self ): return self.data[2]
122
123     property ref:
124         def __get__(self ): 
125             # note: gerton substitutes reference if it can be fixed.
126             return self.data[3].upper()
127
128     property alt:
129         def __get__(self):
130             # convert v3.3 to v4.0 alleles below
131             alt = self.data[4] 
132             if alt == ".": alt = []
133             else: alt = alt.upper().split(',')
134             return alt
135
136     property qual:
137         def __get__(self):
138             qual = self.data[5]
139             if qual == ".": qual = -1
140             else: 
141                 try:    qual = float(qual)
142                 except: self.error(line,self.QUAL_NOT_NUMERICAL)
143
144     property filter:
145         def __get__(self):
146             # postpone checking that filters exist.  Encode missing filter or no filtering as empty list
147             if cols[6] == "." or cols[6] == "PASS" or cols[6] == "0": filter = []
148             else: filter = cols[6].split(';')
149
150             return filter
151
152     property info:
153         def __get__(self):
154             col = self.data[7]
155             # dictionary of keys, and list of values
156             info = {}
157             if col != ".":
158                 for blurp in col.split(';'):
159                     elts = blurp.split('=')
160                     if len(elts) == 1: v = None
161                     elif len(elts) == 2: v = elts[1]
162                     else: self.error(str(self.data),self.ERROR_INFO_STRING)
163                     info[elts[0]] = self.parse_formatdata(elts[0], v, self.vcf._info, line)
164             return info
165
166     property format:
167          def __get__(self):
168              return self.data[8].split(':')
169
170     def __getitem__(self, key):
171         
172         # parse sample columns
173         values = self.data[self.vcf._sample2column[key]].split(':')
174         alt = self.alt
175         format = self.format
176
177         if len(values) > len(format):
178             self.error(line,self.BAD_NUMBER_OF_VALUES,"(found %s values in element %s; expected %s)" % (len(values),sample,len(format)))
179
180         result = {}
181         for idx in range(len(format)):
182             expected = self.vcf.get_expected(format[idx], self.vcf._format, alt)
183             if idx < len(values): value = values[idx]
184             else:
185                 if expected == -1: value = "."
186                 else: value = ",".join(["."]*expected)
187
188             result[format[idx]] = self.vcf.parse_formatdata(format[idx], value, self.vcf._format, line)
189             if expected != -1 and len(result[format[idx]]) != expected:
190                 self.error(str(self.data),self.BAD_NUMBER_OF_PARAMETERS,
191                            "id=%s, expected %s parameters, got %s" % (format[idx],expected,result[format[idx]]))
192                 if len(result[format[idx]] ) < expected: result[format[idx]] += [result[format[idx]][-1]]*(expected-len(result[format[idx]]))
193                 result[format[idx]] = result[format[idx]][:expected]
194
195         return result
196  
197     def __str__(self):
198         return str(self.data)
199
200 class VCF:
201
202     # types
203     NT_UNKNOWN = 0
204     NT_NUMBER = 1
205     NT_ALLELES = 2
206     NT_NR_ALLELES = 3
207     NT_GENOTYPES = 4
208     NT_PHASED_GENOTYPES = 5
209
210     _errors = { 0:"UNKNOWN_FORMAT_STRING:Unknown file format identifier",
211                 1:"BADLY_FORMATTED_FORMAT_STRING:Formatting error in the format string",
212                 2:"BADLY_FORMATTED_HEADING:Did not find 9 required headings (CHROM, POS, ..., FORMAT) %s",
213                 3:"BAD_NUMBER_OF_COLUMNS:Wrong number of columns found (%s)",
214                 4:"POS_NOT_NUMERICAL:Position column is not numerical",
215                 5:"UNKNOWN_CHAR_IN_REF:Unknown character in reference field",
216                 6:"V33_BAD_REF:Reference should be single-character in v3.3 VCF",
217                 7:"V33_BAD_ALLELE:Cannot interpret allele for v3.3 VCF",
218                 8:"POS_NOT_POSITIVE:Position field must be >0",
219                 9:"QUAL_NOT_NUMERICAL:Quality field must be numerical, or '.'",
220                10:"ERROR_INFO_STRING:Error while parsing info field",
221                11:"ERROR_UNKNOWN_KEY:Unknown key (%s) found in formatted field (info; format; or filter)",
222                12:"ERROR_FORMAT_NOT_NUMERICAL:Expected integer or float in formatted field; got %s",
223                13:"ERROR_FORMAT_NOT_CHAR:Eexpected character in formatted field; got string",
224                14:"FILTER_NOT_DEFINED:Identifier (%s) in filter found which was not defined in header",
225                15:"FORMAT_NOT_DEFINED:Identifier (%s) in format found which was not defined in header",
226                16:"BAD_NUMBER_OF_VALUES:Found too many of values in sample column (%s)",
227                17:"BAD_NUMBER_OF_PARAMETERS:Found unexpected number of parameters (%s)",
228                18:"BAD_GENOTYPE:Cannot parse genotype (%s)",
229                19:"V40_BAD_ALLELE:Bad allele found for v4.0 VCF (%s)",
230                20:"MISSING_REF:Reference allele missing",
231                21:"V33_UNMATCHED_DELETION:Deleted sequence does not match reference (%s)",
232                22:"V40_MISSING_ANGLE_BRACKETS:Format definition is not deliminted by angular brackets",
233                23:"FORMAT_MISSING_QUOTES:Description field in format definition is not surrounded by quotes",
234                24:"V40_FORMAT_MUST_HAVE_NAMED_FIELDS:Fields in v4.0 VCF format definition must have named fields",
235                25:"HEADING_NOT_SEPARATED_BY_TABS:Heading line appears separated by spaces, not tabs",
236                26:"WRONG_REF:Wrong reference %s",
237                27:"ERROR_TRAILING_DATA:Numerical field ('%s') has semicolon-separated trailing data",
238                28:"BAD_CHR_TAG:Error calculating chr tag for %s",
239                29:"ZERO_LENGTH_ALLELE:Found zero-length allele",
240                30:"MISSING_INDEL_ALLELE_REF_BASE:Indel alleles must begin with single reference base"
241                 }
242
243     # tag-value pairs; tags are not unique; does not include fileformat, INFO, FILTER or FORMAT fields
244     _header = []
245
246     # version number; 33=v3.3; 40=v4.0
247     _version = 40
248
249     # info, filter and format data
250     _info = {}
251     _filter = {}
252     _format = {}
253
254     # header; and required columns
255     _required = ["CHROM","POS","ID","REF","ALT","QUAL","FILTER","INFO","FORMAT"]
256     _samples = []
257
258     # control behaviour
259     _ignored_errors = set([11])   # ERROR_UNKNOWN_KEY
260     _warn_errors = set([])
261     _leftalign = False
262
263     # reference sequence
264     _reference = None
265
266     # regions to include; None includes everything
267     _regions = None
268
269     # statefull stuff
270     _lineno = -1
271     _line = None
272     _lines = None
273
274     def __init__(self, _copy=None, reference=None, regions=None, lines=None, leftalign=False):
275         # make error identifiers accessible by name
276         for id in self._errors.keys(): self.__dict__[self._errors[id].split(':')[0]] = id
277         if _copy != None:
278             self._leftalign = _copy._leftalign
279             self._header = _copy._header[:]
280             self._version = _copy._version
281             self._info = copy.deepcopy(_copy._info)
282             self._filter = copy.deepcopy(_copy._filter)
283             self._format = copy.deepcopy(_copy._format)
284             self._samples = _copy._samples[:]
285             self._sample2column = copy.deepcopy(_copy._sample2column)
286             self._ignored_errors = copy.deepcopy(_copy._ignored_errors)
287             self._warn_errors = copy.deepcopy(_copy._warn_errors)
288             self._reference = _copy._reference
289             self._regions = _copy._regions
290         if reference: self._reference = reference
291         if regions: self._regions = regions
292         if leftalign: self._leftalign = leftalign
293         self._lines = lines
294
295     def error(self,line,error,opt=None):
296         if error in self._ignored_errors: return
297         errorlabel, errorstring = self._errors[error].split(':')
298         if opt: errorstring = errorstring % opt
299         errwarn = ["Error","Warning"][error in self._warn_errors]
300         sys.stderr.write("Line %s: '%s'\n%s %s: %s\n" % (self._lineno,line,errwarn,errorlabel,errorstring))
301         if error in self._warn_errors: return
302         raise ValueError(errorstring)
303
304     def parse_format(self,line,format,filter=False):
305         if self._version >= 40:
306             if not format.startswith('<'): 
307                 self.error(line,self.V40_MISSING_ANGLE_BRACKETS)
308                 format = "<"+format
309             if not format.endswith('>'): 
310                 self.error(line,self.V40_MISSING_ANGLE_BRACKETS)
311                 format += ">"
312             format = format[1:-1]
313         data = {'id':None,'number':None,'type':None,'descr':None}
314         idx = 0
315         while len(format.strip())>0:
316             elts = format.strip().split(',')
317             first, rest = elts[0], ','.join(elts[1:])
318             if first.find('=') == -1 or (first.find('"')>=0 and first.find('=') > first.find('"')):
319                 if self._version >= 40: self.error(line,self.V40_FORMAT_MUST_HAVE_NAMED_FIELDS)
320                 if idx == 4: self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
321                 first = ["ID=","Number=","Type=","Description="][idx] + first
322             if first.startswith('ID='):            data['id'] = first.split('=')[1]
323             elif first.startswith('Number='):      data['number'] = first.split('=')[1]
324             elif first.startswith('Type='):        data['type'] = first.split('=')[1]
325             elif first.startswith('Description='):
326                 elts = format.split('"')
327                 if len(elts)<3: 
328                     self.error(line,self.FORMAT_MISSING_QUOTES)
329                     elts = first.split('=') + [rest] 
330                 data['descr'] = elts[1]
331                 rest = '"'.join(elts[2:])
332                 if rest.startswith(','): rest = rest[1:]
333             else:
334                 self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
335             format = rest
336             idx += 1
337             if filter and idx==1: idx=3  # skip number and type fields for FILTER format strings
338         if not data['id']: self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
339         if not data['descr']: 
340             self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
341             data['descr'] = '<none>'
342         if not data['type'] and not data['number']:
343             # fine, ##filter format
344             return FORMAT(data['id'],self.NT_NUMBER,0,"Flag",data['descr'],'.')
345         if not data['type'] in ["Integer","Float","Character","String","Flag"]:
346             self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
347         # I would like a missing-value field, but it isn't there
348         if data['type'] in ['Integer','Float']: data['missing'] = None    # Do NOT use arbitrary int/float as missing value
349         else:                                   data['missing'] = '.'
350         if not data['number']: self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
351         try:
352             n = int(data['number'])
353             t = self.NT_NUMBER
354         except ValueError:
355             n = -1
356             if data['number'] == '.':                   t = self.NT_UNKNOWN
357             elif data['number'] == '#alleles':          t = self.NT_ALLELES
358             elif data['number'] == '#nonref_alleles':   t = self.NT_NR_ALLELES
359             elif data['number'] == '#genotypes':        t = self.NT_GENOTYPES
360             elif data['number'] == '#phased_genotypes': t = self.NT_PHASED_GENOTYPES
361             else:
362                 self.error(line,self.BADLY_FORMATTED_FORMAT_STRING)
363         return FORMAT(data['id'],t,n,data['type'],data['descr'],data['missing'])
364     
365
366     def format_format( self, fmt, filter=False ):
367         values = [('ID',fmt.id)]
368         if fmt.number != None and not filter:
369             if fmt.numbertype == self.NT_UNKNOWN: nmb = "."
370             elif fmt.numbertype == self.NT_NUMBER: nmb = str(fmt.number)
371             elif fmt.numbertype == self.NT_ALLELES: nmb = "#alleles"
372             elif fmt.numbertype == self.NT_NR_ALLELES: nmb = "#nonref_alleles"
373             elif fmt.numbertype == self.NT_GENOTYPES: nmb = "#genotypes"
374             elif fmt.numbertype == self.NT_PHASED_GENOTYPES: nmb = "#phased_genotypes"
375             else:
376                 raise ValueError("Unknown number type encountered: %s" % fmt.numbertype)
377             values.append( ('Number',nmb) )
378             values.append( ('Type', fmt.type) )
379         values.append( ('Description', '"' + fmt.description + '"') )
380         if self._version == 33:
381             format = ",".join(v for k,v in values)
382         else:
383             format = "<" + (",".join( "%s=%s" % (k,v) for (k,v) in values )) + ">"
384         return format
385
386     def get_expected(self, format, formatdict, alt):
387         fmt = formatdict[format]
388         if fmt.numbertype == self.NT_UNKNOWN: return -1
389         if fmt.numbertype == self.NT_NUMBER: return fmt.number
390         if fmt.numbertype == self.NT_ALLELES: return len(alt)+1
391         if fmt.numbertype == self.NT_NR_ALLELES: return len(alt)
392         if fmt.numbertype == self.NT_GENOTYPES: return ((len(alt)+1)*(len(alt)+2)) // 2
393         if fmt.numbertype == self.NT_PHASED_GENOTYPES: return (len(alt)+1)*(len(alt)+1)
394         return 0
395
396
397     def _add_definition(self, formatdict, key, data, line ):
398         if key in formatdict: return
399         self.error(line,self.ERROR_UNKNOWN_KEY,key)
400         if data == None:
401             formatdict[key] = FORMAT(key,self.NT_NUMBER,0,"Flag","(Undefined tag)",".")
402             return
403         if data == []: data = [""]             # unsure what type -- say string
404         if type(data[0]) == type(0.0):
405             formatdict[key] = FORMAT(key,self.NT_UNKNOWN,-1,"Float","(Undefined tag)",None)
406             return
407         if type(data[0]) == type(0):
408             formatdict[key] = FORMAT(key,self.NT_UNKNOWN,-1,"Integer","(Undefined tag)",None)
409             return
410         formatdict[key] = FORMAT(key,self.NT_UNKNOWN,-1,"String","(Undefined tag)",".")
411
412
413     # todo: trim trailing missing values
414     def format_formatdata( self, data, format, key=True, value=True, separator=":" ):
415         output, sdata = [], []
416         if type(data) == type([]): # for FORMAT field, make data with dummy values
417             d = {}
418             for k in data: d[k] = []
419             data = d
420         # convert missing values; and silently add definitions if required
421         for k in data:
422             self._add_definition( format, k, data[k], "(output)" )
423             for idx,v in enumerate(data[k]):
424                 if v == format[k].missingvalue: data[k][idx] = "."
425         # make sure GT comes first; and ensure fixed ordering; also convert GT data back to string
426         for k in data: 
427             if k != 'GT': sdata.append( (k,data[k]) )
428         sdata.sort()
429         if 'GT' in data:
430             sdata = [('GT',map(self.convertGTback,data['GT']))] + sdata
431         for k,v in sdata:
432             if v == []: v = None
433             if key and value:
434                 if v != None: output.append( k+"="+','.join(map(str,v)) )
435                 else: output.append( k )
436             elif key: output.append(k)
437             elif value:
438                 if v != None: output.append( ','.join(map(str,v)) )
439                 else: output.append( "." )                    # should not happen
440         # snip off trailing missing data
441         while len(output) > 1:
442             last = output[-1].replace(',','').replace('.','')
443             if len(last)>0: break
444             output = output[:-1]
445         return separator.join(output)
446
447
448     def enter_default_format(self):
449         for f in [FORMAT('GT',self.NT_NUMBER,1,'String','Genotype','.'),
450                   FORMAT('GQ',self.NT_NUMBER,1,'Integer','Genotype Quality',-1),
451                   FORMAT('DP',self.NT_NUMBER,1,'Integer','Read depth at this position for this sample',-1),
452                   FORMAT('HQ',self.NT_UNKNOWN,-1,'Integer','Haplotype Quality',-1),    # unknown number, since may be haploid
453                   FORMAT('FT',self.NT_NUMBER,1,'String','Sample Genotype Filter','.')]:
454             if f.id not in self._format:
455                 self._format[f.id] = f
456
457     def parse_header( self, line ):
458         assert line.startswith('##')
459         elts = line[2:].split('=')
460         key = elts[0].strip()
461         value = '='.join(elts[1:]).strip()
462         if key == "fileformat":
463             if value == "VCFv3.3":
464                 self._version = 33
465             elif value == "VCFv4.0":
466                 self._version = 40
467             elif value == "VCFv4.1":
468                 self._version = 41
469             else:
470                 self.error(line,self.UNKNOWN_FORMAT_STRING)
471         elif key == "INFO":
472             f = self.parse_format(line, value)
473             self._info[ f.id ] = f
474         elif key == "FILTER":
475             f = self.parse_format(line, value, filter=True)
476             self._filter[ f.id ] = f
477         elif key == "FORMAT":
478             f = self.parse_format(line, value)
479             self._format[ f.id ] = f
480         else:
481             # keep other keys in the header field
482             self._header.append( (key,value) )
483
484
485     def write_header( self, stream ):
486         stream.write("##fileformat=VCFv%s.%s\n" % (self._version // 10, self._version % 10))
487         for key,value in self._header: stream.write("##%s=%s\n" % (key,value))
488         for var,label in [(self._info,"INFO"),(self._filter,"FILTER"),(self._format,"FORMAT")]:
489             for f in var.itervalues(): stream.write("##%s=%s\n" % (label,self.format_format(f,filter=(label=="FILTER"))))
490         
491
492     def parse_heading( self, line ):
493         assert line.startswith('#')
494         assert not line.startswith('##')
495         headings = line[1:].split('\t')
496         if len(headings)==1 and len(line[1:].split()) >= 9:
497             self.error(line,self.HEADING_NOT_SEPARATED_BY_TABS)
498             headings = line[1:].split()
499
500         for i,s in enumerate(self._required):
501
502             if len(headings)<=i or headings[i] != s:
503
504                 if len(headings) <= i: 
505                     err = "(%sth entry not found)" % (i+1)
506                 else:
507                     err = "(found %s, expected %s)" % (headings[i],s)
508
509                 #self.error(line,self.BADLY_FORMATTED_HEADING,err)
510
511                 # allow FORMAT column to be absent
512                 if len(headings) == 8:
513                     headings.append("FORMAT")
514                 else:
515                     self.error(line,self.BADLY_FORMATTED_HEADING,err)
516
517         self._samples = headings[9:]
518         self._sample2column = dict( [(y,x) for x,y in enumerate( self._samples ) ] )
519                            
520     def write_heading( self, stream ):
521         stream.write("#" + "\t".join(self._required + self._samples) + "\n")
522
523     def convertGT(self, GTstring):
524         if GTstring == ".": return ["."]
525         try:
526             gts = gtsRegEx.split(GTstring)
527             if len(gts) == 1: return [int(gts[0])]
528             if len(gts) != 2: raise ValueError()
529             if gts[0] == "." and gts[1] == ".": return [gts[0],GTstring[len(gts[0]):-len(gts[1])],gts[1]]
530             return [int(gts[0]),GTstring[len(gts[0]):-len(gts[1])],int(gts[1])]
531         except ValueError:
532             self.error(self._line,self.BAD_GENOTYPE,GTstring)
533             return [".","|","."]
534
535
536     def convertGTback(self, GTdata):
537         return ''.join(map(str,GTdata))
538
539     def parse_formatdata( self, key, value, formatdict, line ):
540         # To do: check that the right number of values is present
541         f = formatdict.get(key,None)
542         if f == None:
543             self._add_definition(formatdict, key, value, line )
544             f = formatdict[key]
545         if f.type == "Flag":
546             if value is not None: self.error(line,self.ERROR_FLAG_HAS_VALUE)
547             return []
548         values = value.split(',')
549         # deal with trailing data in some early VCF files
550         if f.type in ["Float","Integer"] and len(values)>0 and values[-1].find(';') > -1:
551             self.error(line,self.ERROR_TRAILING_DATA,values[-1])
552             values[-1] = values[-1].split(';')[0]
553         if f.type == "Integer": 
554             for idx,v in enumerate(values):
555                 try:
556                     if v == ".": values[idx] = f.missingvalue
557                     else:        values[idx] = int(v)
558                 except: 
559                     self.error(line,self.ERROR_FORMAT_NOT_NUMERICAL,values)
560                     return [0] * len(values)
561             return values
562         elif f.type == "String":
563             self._line = line
564             if f.id == "GT": values = map( self.convertGT, values )
565             return values
566         elif f.type == "Character":
567             for v in values: 
568                 if len(v) != 1: self.error(line,self.ERROR_FORMAT_NOT_CHAR)
569             return values
570         elif f.type == "Float":
571             for idx,v in enumerate(values):
572                 if v == ".": values[idx] = f.missingvalue
573             try: return map(float,values)
574             except:
575                 self.error(line,self.ERROR_FORMAT_NOT_NUMERICAL,values)
576                 return [0.0] * len(values)
577         else:
578             # can't happen
579             self.error(line,self.ERROR_INFO_STRING)
580
581
582     def inregion(self, chrom, pos):
583         if not self._regions: return True
584         for r in self._regions:
585             if r[0] == chrom and r[1] <= pos < r[2]: return True
586         return False
587         
588
589     def parse_data( self, line, lineparse=False ):
590         cols = line.split('\t')
591         if len(cols) != len(self._samples)+9:
592             # gracefully deal with absent FORMAT column
593             if len(cols) == 8 and len(self._samples)==0:
594                 cols.append("")
595             else:
596                 self.error(line,
597                            self.BAD_NUMBER_OF_COLUMNS, 
598                            "expected %s for %s samples (%s), got %s" % (len(self._samples)+9, len(self._samples), self._samples, len(cols)))
599
600         chrom = cols[0]
601
602         # get 0-based position
603         try:    pos = int(cols[1])-1
604         except: self.error(line,self.POS_NOT_NUMERICAL)
605         if pos < 0: self.error(line,self.POS_NOT_POSITIVE)
606
607         # implement filtering
608         if not self.inregion(chrom,pos): return None
609
610         # end of first-pass parse for sortedVCF
611         if lineparse: return chrom, pos, line
612         
613         id = cols[2]
614
615         ref = cols[3].upper()
616         if ref == ".":
617             self.error(line,self.MISSING_REF)
618             if self._version == 33: ref = get_sequence(chrom,pos,pos+1,self._reference)
619             else:                   ref = ""
620         else:
621             for c in ref:
622                 if c not in "ACGTN": self.error(line,self.UNKNOWN_CHAR_IN_REF)
623             if "N" in ref: ref = get_sequence(chrom,pos,pos+len(ref),self._reference)
624
625         # make sure reference is sane
626         if self._reference:
627             left = max(0,pos-100)
628             faref_leftflank = get_sequence(chrom,left,pos+len(ref),self._reference)
629             faref = faref_leftflank[pos-left:]
630             if faref != ref: self.error(line,self.WRONG_REF,"(reference is %s, VCF says %s)" % (faref,ref))
631             ref = faref
632
633         # convert v3.3 to v4.0 alleles below
634         if cols[4] == ".": alt = []
635         else: alt = cols[4].upper().split(',')
636
637         if cols[5] == ".": qual = -1
638         else: 
639             try:    qual = float(cols[5])
640             except: self.error(line,self.QUAL_NOT_NUMERICAL)
641
642         # postpone checking that filters exist.  Encode missing filter or no filtering as empty list
643         if cols[6] == "." or cols[6] == "PASS" or cols[6] == "0": filter = []
644         else: filter = cols[6].split(';')
645
646         # dictionary of keys, and list of values
647         info = {}
648         if cols[7] != ".":
649             for blurp in cols[7].split(';'):
650                 elts = blurp.split('=')
651                 if len(elts) == 1: v = None
652                 elif len(elts) == 2: v = elts[1]
653                 else: self.error(line,self.ERROR_INFO_STRING)
654                 info[elts[0]] = self.parse_formatdata(elts[0], v, self._info, line)
655
656         # Gracefully deal with absent FORMAT column
657         if cols[8] == "": format = []
658         else: format = cols[8].split(':')
659
660         # check: all filters are defined
661         for f in filter:
662             if f not in self._filter: self.error(line,self.FILTER_NOT_DEFINED, f)
663             
664         # check: format fields are defined
665         for f in format:
666             if f not in self._format: self.error(line,self.FORMAT_NOT_DEFINED, f)
667
668         # convert v3.3 alleles
669         if self._version == 33:
670             if len(ref) != 1: self.error(line,self.V33_BAD_REF)
671             newalts = []
672             have_deletions = False
673             for a in alt:
674                 if len(a) == 1: a = a + ref[1:]                       # SNP; add trailing reference
675                 elif a.startswith('I'): a = ref[0] + a[1:] + ref[1:]  # insertion just beyond pos; add first and trailing reference
676                 elif a.startswith('D'): # allow D<seq> and D<num>
677                     have_deletions = True
678                     try:
679                         l = int(a[1:])          # throws ValueError if sequence
680                         if len(ref) < l:        # add to reference if necessary
681                             addns = get_sequence(chrom,pos+len(ref),pos+l,self._reference)
682                             ref += addns
683                             for i,na in enumerate(newalts): newalts[i] = na+addns
684                         a = ref[l:]             # new deletion, deleting pos...pos+l
685                     except ValueError:
686                         s = a[1:]
687                         if len(ref) < len(s):   # add Ns to reference if necessary
688                             addns = get_sequence(chrom,pos+len(ref),pos+len(s),self._reference)
689                             if not s.endswith(addns) and addns != 'N'*len(addns):
690                                 self.error(line,self.V33_UNMATCHED_DELETION,
691                                            "(deletion is %s, reference is %s)" % (a,get_sequence(chrom,pos,pos+len(s),self._reference)))
692                             ref += addns
693                             for i,na in enumerate(newalts): newalts[i] = na+addns
694                         a = ref[len(s):]        # new deletion, deleting from pos
695                 else:
696                     self.error(line,self.V33_BAD_ALLELE)
697                 newalts.append(a)
698             alt = newalts
699             # deletion alleles exist, add dummy 1st reference allele, and account for leading base
700             if have_deletions:
701                 if pos == 0:
702                     # Petr Danacek's: we can't have a leading nucleotide at (1-based) position 1
703                     addn = get_sequence(chrom,pos+len(ref),pos+len(ref)+1,self._reference)
704                     ref += addn
705                     alt = [allele+addn for allele in alt]
706                 else:
707                     addn = get_sequence(chrom,pos-1,pos,self._reference)
708                     ref = addn + ref
709                     alt = [addn + allele for allele in alt]
710                     pos -= 1
711         else:
712             # format v4.0 -- just check for nucleotides
713             for allele in alt:
714                 if not alleleRegEx.match(allele):
715                     self.error(line,self.V40_BAD_ALLELE,allele)
716
717         # check for leading nucleotide in indel calls
718         for allele in alt:
719             if len(allele) != len(ref):
720                 if len(allele) == 0: self.error(line,self.ZERO_LENGTH_ALLELE)
721                 if ref[0].upper() != allele[0].upper() and "N" not in (ref[0]+allele[0]).upper():
722                     self.error(line,self.MISSING_INDEL_ALLELE_REF_BASE)
723
724         # trim trailing bases in alleles
725         for i in range(1,min(len(ref),min(map(len,alt)))):
726             if len(set(allele[-1].upper() for allele in alt)) > 1 or ref[-1].upper() != alt[0][-1].upper():
727                 break
728             ref, alt = ref[:-1], [allele[:-1] for allele in alt]
729
730         # left-align alleles, if a reference is available
731         if self._leftalign and self._reference:
732             while left < pos:
733                 movable = True
734                 for allele in alt:
735                     if len(allele) > len(ref):
736                         longest, shortest = allele, ref
737                     else:
738                         longest, shortest = ref, allele
739                     if len(longest) == len(shortest) or longest[:len(shortest)].upper() != shortest.upper():
740                         movable = False
741                     if longest[-1].upper() != longest[len(shortest)-1].upper():
742                         movable = False
743                 if not movable:
744                     break
745                 ref = ref[:-1]
746                 alt = [allele[:-1] for allele in alt]
747                 if min(len(allele) for allele in alt) == 0 or len(ref) == 0:
748                     ref = faref_leftflank[pos-left-1] + ref
749                     alt = [faref_leftflank[pos-left-1] + allele for allele in alt]
750                     pos -= 1
751
752         # parse sample columns
753         samples = []
754         for sample in cols[9:]:
755             dict = {}
756             values = sample.split(':')
757             if len(values) > len(format):
758                 self.error(line,self.BAD_NUMBER_OF_VALUES,"(found %s values in element %s; expected %s)" % (len(values),sample,len(format)))
759             for idx in range(len(format)):
760                 expected = self.get_expected(format[idx], self._format, alt)
761                 if idx < len(values): value = values[idx]
762                 else:
763                     if expected == -1: value = "."
764                     else: value = ",".join(["."]*expected)
765                 dict[format[idx]] = self.parse_formatdata(format[idx], value, self._format, line)
766                 if expected != -1 and len(dict[format[idx]]) != expected:
767                     self.error(line,self.BAD_NUMBER_OF_PARAMETERS,
768                                "id=%s, expected %s parameters, got %s" % (format[idx],expected,dict[format[idx]]))
769                     if len(dict[format[idx]] ) < expected: dict[format[idx]] += [dict[format[idx]][-1]]*(expected-len(dict[format[idx]]))
770                     dict[format[idx]] = dict[format[idx]][:expected]
771             samples.append( dict )
772
773         # done
774         d = {'chrom':chrom,
775              'pos':pos,      # return 0-based position
776              'id':id,
777              'ref':ref,
778              'alt':alt,
779              'qual':qual,
780              'filter':filter,
781              'info':info,
782              'format':format}
783         for key,value in zip(self._samples,samples):
784             d[key] = value
785         
786         return d
787
788
789     def write_data(self, stream, data):
790         required = ['chrom','pos','id','ref','alt','qual','filter','info','format'] + self._samples
791         for k in required:
792             if k not in data: raise ValueError("Required key %s not found in data" % str(k))
793         if data['alt'] == []: alt = "."
794         else: alt = ",".join(data['alt'])
795         if data['filter'] == None: filter = "."
796         elif data['filter'] == []: 
797             if self._version == 33: filter = "0"
798             else: filter = "PASS"
799         else: filter = ';'.join(data['filter'])
800         if data['qual'] == -1: qual = "."
801         else: qual = str(data['qual'])
802
803         output = [data['chrom'], 
804                   str(data['pos']+1),   # change to 1-based position
805                   data['id'],
806                   data['ref'],
807                   alt,
808                   qual,
809                   filter,
810                   self.format_formatdata( data['info'], self._info, separator=";" ),
811                   self.format_formatdata( data['format'], self._format, value=False ) ]
812         
813         for s in self._samples:
814             output.append( self.format_formatdata( data[s], self._format, key=False ) )
815         
816         stream.write( "\t".join(output) + "\n" )
817
818     def _parse_header(self, stream):
819         self._lineno = 0
820         for line in stream:
821             self._lineno += 1
822             if line.startswith('##'):
823                 self.parse_header( line.strip() )
824             elif line.startswith('#'):
825                 self.parse_heading( line.strip() )
826                 self.enter_default_format()
827             else:
828                 break
829         return line
830
831     def _parse(self, line, stream):
832         if len(line.strip()) > 0:
833             d = self.parse_data( line.strip() )
834             if d: yield d
835         for line in stream:
836             self._lineno += 1
837             if self._lines and self._lineno > self._lines: raise StopIteration
838             d = self.parse_data( line.strip() )
839             if d: yield d
840
841     ######################################################################################################
842     #
843     # API follows
844     #
845     ######################################################################################################
846
847     def getsamples(self):
848         """ List of samples in VCF file """
849         return self._samples
850
851     def setsamples(self,samples):
852         """ List of samples in VCF file """
853         self._samples = samples
854
855     def getheader(self):
856         """ List of header key-value pairs (strings) """
857         return self._header
858
859     def setheader(self,header):
860         """ List of header key-value pairs (strings) """
861         self._header = header
862
863     def getinfo(self):
864         """ Dictionary of ##INFO tags, as VCF.FORMAT values """
865         return self._info
866
867     def setinfo(self,info):
868         """ Dictionary of ##INFO tags, as VCF.FORMAT values """
869         self._info = info
870
871     def getformat(self):
872         """ Dictionary of ##FORMAT tags, as VCF.FORMAT values """
873         return self._format
874
875     def setformat(self,format):
876         """ Dictionary of ##FORMAT tags, as VCF.FORMAT values """
877         self._format = format
878
879     def getfilter(self):
880         """ Dictionary of ##FILTER tags, as VCF.FORMAT values """
881         return self._filter
882
883     def setfilter(self,filter):
884         """ Dictionary of ##FILTER tags, as VCF.FORMAT values """
885         self._filter = filter
886
887     def setversion(self, version):
888         if version not in [33,40,41]: raise ValueError("Can only handle v3.3, v4.0 and v4.1 VCF files")
889         self._version = version
890
891     def setregions(self, regions):
892         self._regions = regions
893
894     def setreference(self, ref):
895         """ Provide a reference sequence; a Python class supporting a fetch(chromosome, start, end) method, e.g. PySam.FastaFile """
896         self._reference = ref
897
898     def ignoreerror(self, errorstring):
899         try:             self._ignored_errors.add(self.__dict__[errorstring])
900         except KeyError: raise ValueError("Invalid error string: %s" % errorstring)
901
902     def warnerror(self, errorstring):
903         try:             self._warn_errors.add(self.__dict__[errorstring])
904         except KeyError: raise ValueError("Invalid error string: %s" % errorstring)
905
906     def parse(self, stream):
907         """ Parse a stream of VCF-formatted lines.  Initializes class instance and return generator """
908         last_line = self._parse_header(stream)
909         # now return a generator that does the actual work.  In this way the pre-processing is done
910         # before the first piece of data is yielded
911         return self._parse(last_line, stream)
912
913     def write(self, stream, datagenerator):
914         """ Writes a VCF file to a stream, using a data generator (or list) """
915         self.write_header(stream)
916         self.write_heading(stream)
917         for data in datagenerator: self.write_data(stream,data)
918
919     def writeheader(self, stream):
920         """ Writes a VCF header """
921         self.write_header(stream)
922         self.write_heading(stream)
923
924     def compare_calls(self, pos1, ref1, alt1, pos2, ref2, alt2):
925         """ Utility function: compares two calls for equality """
926         # a variant should always be assigned to a unique position, one base before
927         # the leftmost position of the alignment gap.  If this rule is implemented
928         # correctly, the two positions must be equal for the calls to be identical.
929         if pos1 != pos2: return False
930         # from both calls, trim rightmost bases when identical.  Do this safely, i.e.
931         # only when the reference bases are not Ns
932         while len(ref1)>0 and len(alt1)>0 and ref1[-1] == alt1[-1]:
933             ref1 = ref1[:-1]
934             alt1 = alt1[:-1]
935         while len(ref2)>0 and len(alt2)>0 and ref2[-1] == alt2[-1]:
936             ref2 = ref2[:-1]
937             alt2 = alt2[:-1]
938         # now, the alternative alleles must be identical
939         return alt1 == alt2
940
941 ###########################################################################################################
942 ###########################################################################################################
943 ## API functions added by Andreas
944 ###########################################################################################################
945
946     def connect( self, filename ):
947         '''connect to tabix file.'''
948         self.tabixfile = pysam.Tabixfile( filename )
949         self._parse_header(self.tabixfile.header)
950         
951     def fetch(self,
952               reference = None,
953               start = None, 
954               end = None, 
955               region = None ):
956         """ Parse a stream of VCF-formatted lines.  Initializes class instance and return generator """
957
958         iter = self.tabixfile.fetch( reference, start, end, region, parser = pysam.asVCF() )
959         for x in iter:
960             yield VCFRecord( x, self )
961
962     def validate( self, record ):
963         '''validate vcf record.
964
965         returns a validated record.
966         '''
967
968         chrom, pos = record.chrom, record.pos
969
970         # check reference
971         ref = record.ref
972         if ref == ".":
973             self.error(str(record),self.MISSING_REF)
974             if self._version == 33: ref = get_sequence(chrom,pos,pos+1,self._reference)
975             else:                   ref = ""
976         else:
977             for c in ref:
978                 if c not in "ACGTN": self.error(str(record),self.UNKNOWN_CHAR_IN_REF)
979                 if "N" in ref: ref = get_sequence(chrom,
980                                                   pos,
981                                                   pos+len(ref),
982                                                   self._reference)
983
984         # make sure reference is sane
985         if self._reference:
986             left = max(0,self.pos-100)
987             faref_leftflank = get_sequence(chrom,left,self.pos+len(ref),self._reference)
988             faref = faref_leftflank[pos-left:]
989             if faref != ref: self.error(line,self.WRONG_REF,"(reference is %s, VCF says %s)" % (faref,ref))
990             ref = faref
991             
992         # check: format fields are defined
993         for f in record.format:
994             if f not in self._format: self.error(str(record),self.FORMAT_NOT_DEFINED, f)
995             
996         # check: all filters are defined
997         for f in record.filter:
998             if f not in self._filter: self.error(str(record),self.FILTER_NOT_DEFINED, f)
999
1000         # convert v3.3 alleles
1001         if self._version == 33:
1002             if len(ref) != 1: self.error(line,self.V33_BAD_REF)
1003             newalts = []
1004             have_deletions = False
1005             for a in alt:
1006                 if len(a) == 1: a = a + ref[1:]                       # SNP; add trailing reference
1007                 elif a.startswith('I'): a = ref[0] + a[1:] + ref[1:]  # insertion just beyond pos; add first and trailing reference
1008                 elif a.startswith('D'): # allow D<seq> and D<num>
1009                     have_deletions = True
1010                     try:
1011                         l = int(a[1:])          # throws ValueError if sequence
1012                         if len(ref) < l:        # add to reference if necessary
1013                             addns = get_sequence(chrom,pos+len(ref),pos+l,self._reference)
1014                             ref += addns
1015                             for i,na in enumerate(newalts): newalts[i] = na+addns
1016                         a = ref[l:]             # new deletion, deleting pos...pos+l
1017                     except ValueError:
1018                         s = a[1:]
1019                         if len(ref) < len(s):   # add Ns to reference if necessary
1020                             addns = get_sequence(chrom,pos+len(ref),pos+len(s),self._reference)
1021                             if not s.endswith(addns) and addns != 'N'*len(addns):
1022                                 self.error(line,self.V33_UNMATCHED_DELETION,
1023                                            "(deletion is %s, reference is %s)" % (a,get_sequence(chrom,pos,pos+len(s),self._reference)))
1024                             ref += addns
1025                             for i,na in enumerate(newalts): newalts[i] = na+addns
1026                         a = ref[len(s):]        # new deletion, deleting from pos
1027                 else:
1028                     self.error(line,self.V33_BAD_ALLELE)
1029                 newalts.append(a)
1030             alt = newalts
1031             # deletion alleles exist, add dummy 1st reference allele, and account for leading base
1032             if have_deletions:
1033                 if pos == 0:
1034                     # Petr Danacek's: we can't have a leading nucleotide at (1-based) position 1
1035                     addn = get_sequence(chrom,pos+len(ref),pos+len(ref)+1,self._reference)
1036                     ref += addn
1037                     alt = [allele+addn for allele in alt]
1038                 else:
1039                     addn = get_sequence(chrom,pos-1,pos,self._reference)
1040                     ref = addn + ref
1041                     alt = [addn + allele for allele in alt]
1042                     pos -= 1
1043         else:
1044             # format v4.0 -- just check for nucleotides
1045             for allele in alt:
1046                 if not alleleRegEx.match(allele):
1047                     self.error(line,self.V40_BAD_ALLELE,allele)
1048                     
1049
1050         # check for leading nucleotide in indel calls
1051         for allele in alt:
1052             if len(allele) != len(ref):
1053                 if len(allele) == 0: self.error(line,self.ZERO_LENGTH_ALLELE)
1054                 if ref[0].upper() != allele[0].upper() and "N" not in (ref[0]+allele[0]).upper():
1055                     self.error(line,self.MISSING_INDEL_ALLELE_REF_BASE)
1056
1057         # trim trailing bases in alleles
1058         for i in range(1,min(len(ref),min(map(len,alt)))):
1059             if len(set(allele[-1].upper() for allele in alt)) > 1 or ref[-1].upper() != alt[0][-1].upper():
1060                 break
1061             ref, alt = ref[:-1], [allele[:-1] for allele in alt]
1062
1063         # left-align alleles, if a reference is available
1064         if self._leftalign and self._reference:
1065             while left < pos:
1066                 movable = True
1067                 for allele in alt:
1068                     if len(allele) > len(ref):
1069                         longest, shortest = allele, ref
1070                     else:
1071                         longest, shortest = ref, allele
1072                     if len(longest) == len(shortest) or longest[:len(shortest)].upper() != shortest.upper():
1073                         movable = False
1074                     if longest[-1].upper() != longest[len(shortest)-1].upper():
1075                         movable = False
1076                 if not movable:
1077                     break
1078                 ref = ref[:-1]
1079                 alt = [allele[:-1] for allele in alt]
1080                 if min(len(allele) for allele in alt) == 0 or len(ref) == 0:
1081                     ref = faref_leftflank[pos-left-1] + ref
1082                     alt = [faref_leftflank[pos-left-1] + allele for allele in alt]
1083                     pos -= 1
1084
1085
1086
1087