tabix.py

   1 #!/usr/bin/env python
   2
   3 # Author: Heng Li and Aaron Quinlan
   4 # License: MIT/X11
   5
   6 import sys
   7 from ctypes import *
   8 from ctypes.util import find_library
   9 import glob, platform
  10
  11 def load_shared_library(lib, _path='.', ver='*'):
  12     """Search for and load the tabix library. The
  13     expectation is that the library is located in
  14     the current directory (ie. "./")
  15     """
  16     # find from the system path
  17     path = find_library(lib)
  18     if (path == None): # if fail, search in the custom directory
  19         s = platform.system()
  20         if (s == 'Darwin'): suf = ver+'.dylib'
  21         elif (s == 'Linux'): suf = '.so'+ver
  22         candidates = glob.glob(_path+'/lib'+lib+suf);
  23         if (len(candidates) == 1): path = candidates[0]
  24         else: return None
  25     cdll.LoadLibrary(path)
  26     return CDLL(path)
  27
  28 def tabix_init():
  29     """Initialize and return a tabix reader object
  30     for subsequent tabix_get() calls.
  31     """
  32     tabix = load_shared_library('tabix')
  33     if (tabix == None): return None
  34     tabix.ti_read.restype = c_char_p
  35     # on Mac OS X 10.6, the following declarations are required.
  36     tabix.ti_open.restype = c_void_p
  37     tabix.ti_querys.argtypes = [c_void_p, c_char_p]
  38     tabix.ti_querys.restype = c_void_p
  39     tabix.ti_query.argtypes = [c_void_p, c_char_p, c_int, c_int]
  40     tabix.ti_query.restype = c_void_p
  41     tabix.ti_read.argtypes = [c_void_p, c_void_p, c_void_p]
  42     tabix.ti_iter_destroy.argtypes = [c_void_p]
  43     tabix.ti_close.argtypes = [c_void_p]
  44     # FIXME: explicit declarations for APIs not used in this script
  45     return tabix
  46
  47 # OOP interface
  48 class Tabix:
  49     def __init__(self, fn, fnidx=0):
  50         self.tabix = tabix_init();
  51         if (self.tabix == None):
  52             sys.stderr.write("[Tabix] Please make sure the shared library is compiled and available.\n")
  53             return
  54         self.fp = self.tabix.ti_open(fn, fnidx);
  55
  56     def __del__(self):
  57         if (self.tabix): self.tabix.ti_close(self.fp)
  58
  59     def fetch(self, chr, start=-1, end=-1):
  60         """Generator function that will yield each interval
  61         within the requested range from the requested file.
  62         """
  63         if (self.tabix == None): return
  64         if (start < 0): iter = self.tabix.ti_querys(self.fp, chr) # chr looks like: "chr2:1,000-2,000" or "chr2"
  65         else: iter = self.tabix.ti_query(self.fp, chr, start, end) # chr must be a sequence name
  66         if (iter == None):
  67             sys.stderr.write("[Tabix] Malformatted query or wrong sequence name.\n")
  68             return
  69         while (1): # iterate
  70             s = self.tabix.ti_read(self.fp, iter, 0)
  71             if (s == None): break
  72             yield s
  73         self.tabix.ti_iter_destroy(iter)
  74
  75 # command-line interface
  76 def main():
  77     if (len(sys.argv) < 3):
  78         sys.stderr.write("Usage: tabix.py <in.gz> <reg>\n")
  79         sys.exit(1)
  80
  81     # report the features in the requested interval
  82     tabix = Tabix(sys.argv[1])
  83     for line in tabix.fetch(sys.argv[2]):
  84         print line
  85
  86 if __name__ == '__main__':
  87     main()