index samples by sample name for easier access.

vcf.py

@@ -216,7 +216,6 @@ def __doc__(self):
 class VCFReader(object):
     '''Read and parse a VCF v 4.0 file'''
     def __init__(self, fsock, aggressive=False):
-        super(VCFReader, self).__init__()
         self.aggro = aggressive
         self._metadata = None
         self._infos = None
@@ -351,9 +350,9 @@ def _parse_info(self, info_str):
     def _parse_samples(self, samples, samp_fmt):
         '''Parse a sample entry according to the format specified in the FORMAT
         column.'''
-        samp_data = []
+        samp_data = {}
         samp_fmt = samp_fmt.split(':')
-        for sample in samples:
+        for i, sample in enumerate(samples):
             sampdict = dict(zip(samp_fmt, sample.split(':')))
             for fmt in sampdict:
                 vals = sampdict[fmt].split(',')
@@ -372,13 +371,32 @@ def _parse_samples(self, samples, samp_fmt):
                 elif sampdict[fmt] == './.' and self.aggro:
                     sampdict[fmt] = None
 
-            samp_data.append(sampdict)
-
-        for name, data in zip(self.samples, samp_data):
-            data['name'] = name
+            sample_name = self.samples[i]
+            sampdict["name"] = sample_name
+            samp_data[sample_name] = sampdict
 
         return samp_data
 
+    def genotypes(self, rec):
+        """
+        return the actual allele from the encoded genotypes
+        """
+        # create an array so that the 0/1/2 etc. just index
+        # ref/first-ALT/2nd-ALT respectively...
+        gts = [rec.REF] + rec.ALT
+
+        genotypes = {}
+        for name in self.samples:
+            sample = rec.samples[name]
+            sample_genotype = sample['GT']
+            if sample_genotype in (None, "."):
+                genotypes[name] = None
+                continue
+            # split on '|' or '/';
+            vals = map(int, re.split('\||/', sample_genotype))
+            genotypes[name] = [gts[i] for i in vals]
+        return genotypes
+
     def next(self):
         '''Return the next record in the file.'''
         if self._samples is None:
@@ -449,7 +467,8 @@ def main():
         vcf_file = VCFReader(sock, aggressive=True)
         for record in vcf_file:
             print record
-
+        print record
+        print vcf_file.genotypes(record)
 
 if __name__ == '__main__':
     main()