pybam_old.py 52KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744
  1. '''
  2. Awesome people who have directly contributed to the project:
  3. Jon Palmer - Bug finder & advice on project direction
  4. Mahmut Uludag - Bug finder
  5. Help: print pybam.wat
  6. Github: http://github.com/JohnLonginotto/pybam
  7. This code was written by John Longinotto, a PhD student of the Pospisilik Lab at the Max Planck Institute of Immunbiology & Epigenetics, Freiburg.
  8. My PhD is funded by the Deutsches Epigenom Programm (DEEP), and the Max Planck IMPRS Program.
  9. I study Adipose Biology and Circadian Rhythm in mice, although it seems these days I spend most of my time at the computer :-)
  10. Ported to Python3 by Thomas Forest (thomas.forest@college-de-france.fr)
  11. '''
  12. import os
  13. import sys
  14. import zlib
  15. import time
  16. import tempfile
  17. import subprocess
  18. from array import array
  19. from struct import unpack
  20. CtoPy = { 'A':'<c', 'c':'<b', 'C':'<B', 's':'<h', 'S':'<H', 'i':'<i', 'I':'<I', 'f':'<f' }
  21. py4py = { 'A': 1 , 'c': 1 , 'C': 1 , 's': 2 , 'S': 2 , 'i': 4 , 'I': 4 , 'f': 4 }
  22. dna_codes = '=ACMGRSVTWYHKDBN'
  23. cigar_codes = 'MIDNSHP=X'
  24. parse_codes = {
  25. 'sam': ' The current alignment in SAM format.',
  26. 'bam': ' All the bytes that make up the current alignment ("read"),\n still in binary just as it was in the BAM file. Useful\n when creating a new BAM file of filtered alignments.',
  27. 'sam_qname': ' [1st column in SAM] The QNAME (fragment ID) of the alignment.',
  28. 'bam_qname': ' The original bytes before decoding to sam_qname.',
  29. 'sam_flag': ' [2nd column in SAM] The FLAG number of the alignment.',
  30. 'bam_flag': ' The original bytes before decoding to sam_flag.',
  31. 'sam_refID': ' The chromosome ID (not the same as the name!).\n Chromosome names are stored in the BAM header (file_chromosomes),\n so to convert refIDs to chromsome names one needs to do:\n "my_bam.file_chromosomes[read.sam_refID]" (or use sam_rname)\n But for comparisons, using the refID is much faster that using\n the actual chromosome name (for example, when reading through a\n sorted BAM file and looking for where last_refID != this_refID)\n Note that when negative the alignment is not aligned, and thus one\n must not perform my_bam.file_chromosomes[read.sam_refID]\n without checking that the value is positive first.',
  32. 'sam_rname': ' [3rd column in SAM] The actual chromosome/contig name for the\n alignment. Will return "*" if refID is negative.',
  33. 'bam_refID': ' The original bytes before decoding to sam_refID.',
  34. 'sam_pos1': ' [4th column in SAM] The 1-based position of the alignment. Note\n that in SAM format values less than 1 are converted to "0" for\n "no data" and sam_pos1 will also do this.',
  35. 'sam_pos0': ' The 0-based position of the alignment. Note that in SAM all\n positions are 1-based, but in BAM they are stored as 0-based.\n Unlike sam_pos1, negative values are kept as negative values,\n essentially giving one the decoded value as it was stored.',
  36. 'bam_pos': ' The original bytes before decoding to sam_pos*.',
  37. 'sam_mapq': ' [5th column in SAM] The Mapping Quality of the current alignment.',
  38. 'bam_mapq': ' The original bytes before decoding to sam_mapq.',
  39. 'sam_cigar_string': ' [6th column in SAM] The CIGAR string, as per the SAM format.\n Allowed values are "MIDNSHP=X".',
  40. 'sam_cigar_list': ' A list of tuples with 2 values per tuple:\n the number of bases, and the CIGAR operation applied to those\n bases. Faster to calculate than sam_cigar_string.',
  41. 'bam_cigar': ' The original bytes before decoding to sam_cigar_*.',
  42. 'sam_next_refID': ' The sam_refID of the alignment\'s mate (if any). Note that as per\n sam_refID, this value can be negative and is not the actual\n chromosome name (see sam_pnext1).',
  43. 'sam_rnext': ' [7th column in SAM] The chromosome name of the alignment\'s mate.\n Value is "*" if unmapped. Note that in a SAM file this value\n is "=" if it is the same as the sam_rname, however pybam will\n only do this if the user prints the whole SAM entry with "sam".',
  44. 'bam_next_refID': ' The original bytes before decoding to sam_next_refID.',
  45. 'sam_pnext1': ' [8th column in SAM] The 1-based position of the alignment\'s mate.\n Note that in SAM format values less than 1 are converted to "0"\n for "no data", and sam_pnext1 will also do this.',
  46. 'sam_pnext0': ' The 0-based position of the alignment\'s mate. Note that in SAM all\n positions are 1-based, but in BAM they are stored as 0-based.\n Unlike sam_pnext1, negative values are kept as negative values\n here, essentially giving you the value as it was stored in BAM.',
  47. 'bam_pnext': ' The original bytes before decoding to sam_pnext0.',
  48. 'sam_tlen': ' [9th column in SAM] The TLEN value.',
  49. 'bam_tlen': ' The original bytes before decoding to sam_tlen.',
  50. 'sam_seq': ' [10th column in SAM] The SEQ value (DNA sequence of the alignment).\n Allowed values are "ACGTMRSVWYHKDBN and =".',
  51. 'bam_seq': ' The original bytes before decoding to sam_seq.',
  52. 'sam_qual': ' [11th column in SAM] The QUAL value (quality scores per DNA base\n in SEQ) of the alignment.',
  53. 'bam_qual': ' The original bytes before decoding to sam_qual.',
  54. 'sam_tags_list': ' A list of tuples with 3 values per tuple: a two-letter TAG ID, the\n type code used to describe the data in the TAG value (see SAM spec.\n for details), and the value of the TAG. Note that the BAM format\n has type codes like "c" for a number in the range -127 to +127,\n and "C" for a number in the range of 0 to 255.\n In a SAM file however, all numerical codes appear to just be stored\n using "i", which is a number in the range -2147483647 to +2147483647.\n sam_tags_list will therefore return the code used in the BAM file,\n and not "i" for all numbers.',
  55. 'sam_tags_string': ' [12th column a SAM] Returns the TAGs in the same format as would be found \n in a SAM file (with all numbers having a signed 32bit code of "i").',
  56. 'bam_tags': ' The original bytes before decoding to sam_tags_*.',
  57. 'sam_bin': ' The bin value of the alignment (used for indexing reads).\n Please refer to section 5.3 of the SAM spec for how this\n value is calculated.',
  58. 'bam_bin': ' The original bytes before decoding to sam_bin.',
  59. 'sam_block_size': ' The number of bytes the current alignment takes up in the BAM\n file minus the four bytes used to store the block_size value\n itself. Essentially sam_block_size +4 is bytes needed to store\n the current alignment.',
  60. 'bam_block_size': ' The original bytes before decoding to sam_block_size.',
  61. 'sam_l_read_name': ' The length of the QNAME plus 1 because the QNAME is terminated\n with a NUL byte.',
  62. 'bam_l_read_name': ' The original bytes before decoding to sam_l_read_name.',
  63. 'sam_l_seq': ' The number of bases in the seq. Useful if you just want to know\n how many bases are in the SEQ but do not need to know what those\n bases are (which requires more decoding effort).',
  64. 'bam_l_seq': ' The original bytes before decoding to sam_l_seq.',
  65. 'sam_n_cigar_op': ' The number of CIGAR operations in the CIGAR field. Useful if one\n wants to know how many CIGAR operations there are, but does not\n need to know what they are.',
  66. 'bam_n_cigar_op': ' The original bytes before decoding to sam_n_cigar_op.',
  67. 'file_alignments_read': ' A running counter of the number of alignments ("reads"),\n processed thus far. Note the BAM format does not store\n how many reads are in a file, so the usefulness of this\n metric is somewhat limited unless one already knows how\n many reads are in the file.',
  68. 'file_binary_header': ' From the first byte in the file, until the first byte of\n the first read. The original binary header.',
  69. 'file_bytes_read': ' A running counter of the bytes read from the file. Note\n that as data is read in arbitary chunks, this is literally\n the amount of data read from the file/pipe by pybam.',
  70. 'file_chromosome_lengths': ' The binary header of the BAM file includes chromosome names\n and chromosome lengths. This is a dictionary of chromosome-name\n keys and chromosome-length values.',
  71. 'file_chromosomes': ' A list of chromosomes from the binary header.',
  72. 'file_decompressor': ' BAM files are compressed with bgzip. The value here reflects\n the decompressor used. "internal" if pybam\'s internal\n decompressor is being used, "gzip" or "pigz" if the system\n has these binaries installed and pybam can find them.\n Any other value reflects a custom decompression command.',
  73. 'file_directory': ' The directory the input BAM file can be found in. This will be\n correct if the input file is specified via a string or python\n file object, however if the input is a pipe such as sys.stdin, \n then the current working directory will be used.',
  74. 'file_header': ' The ASCII portion of the BAM header. This is the typical header\n users of samtools will be familiar with.',
  75. 'file_name': ' The file name (base name) of input file if input is a string or\n python file object. If input is via stdin this will be "<stdin>"'
  76. }
  77. wat = '''
  78. Main class: pybam.read
  79. Github: http://github.com/JohnLonginotto/pybam
  80. [ Dynamic Parser Example ]
  81. for alignment in pybam.read('/my/data.bam'):
  82. print alignment.sam_seq
  83. [ Static Parser Example ]
  84. for seq,mapq in pybam.read('/my/data.bam',['sam_seq','sam_mapq']):
  85. print seq
  86. print mapq
  87. [ Mixed Parser Example ]
  88. my_bam = pybam.read('/my/data.bam',['sam_seq','sam_mapq'])
  89. print my_bam._static_parser_code
  90. for seq,mapq in my_bam:
  91. if seq.startswith('ACGT') and mapq > 10:
  92. print my_bam.sam
  93. [ Custom Decompressor (from file path) Example ]
  94. my_bam = pybam.read('/my/data.bam.lzma',decompressor='lzma --decompress --stdout /my/data.bam.lzma')
  95. [ Custom Decompressor (from file object) Example ]
  96. my_bam = pybam.read(sys.stdin,decompressor='lzma --decompress --stdout') # data given to lzma via stdin
  97. [ Force Internal bgzip Decompressor ]
  98. my_bam = pybam.read('/my/data.bam',decompressor='internal')
  99. [ Parse Words (hah) ]'''
  100. wat += '\n'+''.join([('\n===============================================================================================\n\n ' if code == 'file_alignments_read' or code == 'sam' else ' ')+(code+' ').ljust(25,'-')+description+'\n' for code,description in sorted(parse_codes.items())]) + '\n'
  101. class read:
  102. '''
  103. [ Dynamic Parser Example ]
  104. for alignment in pybam.read('/my/data.bam'):
  105. print alignment.sam_seq
  106. [ Static Parser Example ]
  107. for seq,mapq in pybam.read('/my/data.bam',['sam_seq','sam_mapq']):
  108. print seq
  109. print mapq
  110. [ Mixed Parser Example ]
  111. my_bam = pybam.read('/my/data.bam',['sam_seq','sam_mapq'])
  112. print my_bam._static_parser_code
  113. for seq,mapq in my_bam:
  114. if seq.startswith('ACGT') and mapq > 10:
  115. print my_bam.sam
  116. [ Custom Decompressor (from file path) Example ]
  117. my_bam = pybam.read('/my/data.bam.lzma',decompressor='lzma --decompress --stdout /my/data.bam.lzma')
  118. [ Custom Decompressor (from file object) Example ]
  119. my_bam = pybam.read(sys.stdin,decompressor='lzma --decompress --stdout') # data given to lzma via stdin
  120. [ Force Internal bgzip Decompressor ]
  121. my_bam = pybam.read('/my/data.bam',decompressor='internal')
  122. "print pybam.wat" in the python terminal to see the possible parsable values,
  123. or v==it http://github.com/JohnLonginotto/pybam for the latest info.
  124. '''
  125. def __init__(self,f,fields=False,decompressor=False):
  126. self.file_bytes_read = 0
  127. self.file_chromosomes = []
  128. self.file_alignments_read = 0
  129. self.file_chromosome_lengths = {}
  130. if fields != False:
  131. if type(fields) != list or len(fields) == 0:
  132. raise PybamError('\n\nFields for the static parser must be provided as a non-empty list. You gave a ' + str(type(fields)) + '\n')
  133. else:
  134. for field in fields:
  135. if field.startswith('sam') or field.startswith('bam'):
  136. if field not in parse_codes.keys():
  137. raise PybamError('\n\nStatic parser field "' + str(field) + '" from fields ' + str(fields) + ' is not known to this version of pybam!\nPrint "pybam.wat" to see available field names with explinations.\n')
  138. else:
  139. raise PybamError('\n\nStatic parser field "' + str(field) + '" from fields ' + str(fields) + ' does not start with "sam" or "bam" and thus is an avaliable field for the static parsing.\nPrint "pybam.wat" in interactive python to see available field names with explinations.\n')
  140. if decompressor:
  141. if type(decompressor) == str:
  142. if decompressor != 'internal' and '{}' not in decompressor: raise PybamError('\n\nWhen a custom decompressor is used and the input file is a string, the decompressor string must contain at least one occurence of "{}" to be substituted with a filepath by pybam.\n')
  143. else: raise PybamError('\n\nUser-supplied decompressor must be a string that when run on the command line decompresses a named file (or stdin), to stdout:\ne.g. "lzma --decompress --stdout {}" if pybam is provided a path as input file, where {} is substituted for that path.\nor just "lzma --decompress --stdout" if pybam is provided a file object instead of a file path, as data from that file object will be piped via stdin to the decompression program.\n')
  144. ## First we make a generator that will return chunks of uncompressed data, regardless of how we choose to decompress:
  145. def generator():
  146. DEVNULL = open(os.devnull, 'wb')
  147. # First we need to figure out what sort of file we have - whether it's gzip compressed, uncompressed, or something else entirely!
  148. if type(f) == str:
  149. try: self._file = open(f,'rb')
  150. except: raise PybamError('\n\nCould not open "' + str(self._file.name) + '" for reading!\n')
  151. try: magic = os.read(self._file.fileno(),4)
  152. except: raise PybamError('\n\nCould not read from "' + str(self._file.name) + '"!\n')
  153. elif type(f) == file:
  154. self._file = f
  155. try: magic = os.read(self._file.fileno(),4)
  156. except: raise PybamError('\n\nCould not read from "' + str(self._file.name) + '"!\n')
  157. else: raise PybamError('\n\nInput file was not a string or a file object. It was: "' + str(f) + '"\n')
  158. self.file_name = os.path.basename(os.path.realpath(self._file.name))
  159. self.file_directory = os.path.dirname(os.path.realpath(self._file.name))
  160. if magic == 'BAM\1':
  161. # The user has passed us already unzipped BAM data! Job done :)
  162. data = 'BAM\1' + self._file.read(35536)
  163. self.file_bytes_read += len(data)
  164. self.file_decompressor = 'None'
  165. while data:
  166. yield data
  167. data = self._file.read(35536)
  168. self.file_bytes_read += len(data)
  169. self._file.close()
  170. DEVNULL.close()
  171. raise StopIteration
  172. elif magic == "\x1f\x8b\x08\x04": # The user has passed us compressed gzip/bgzip data, which is typical for a BAM file
  173. # use custom decompressor if provided:
  174. if decompressor != False and decompressor != 'internal':
  175. if type(f) == str: self._subprocess = subprocess.Popen( decompressor.replace('{}',f), shell=True, stdout=subprocess.PIPE, stderr=DEVNULL)
  176. else: self._subprocess = subprocess.Popen('{ printf "'+magic+'"; cat; } | ' + decompressor, stdin=self._file, shell=True, stdout=subprocess.PIPE, stderr=DEVNULL)
  177. self.file_decompressor = decompressor
  178. data = self._subprocess.stdout.read(35536)
  179. self.file_bytes_read += len(data)
  180. while data:
  181. yield data
  182. data = self._subprocess.stdout.read(35536)
  183. self.file_bytes_read += len(data)
  184. self._file.close()
  185. DEVNULL.close()
  186. raise StopIteration
  187. # else look for pigz or gzip:
  188. else:
  189. try:
  190. self._subprocess = subprocess.Popen(["pigz"],stdin=DEVNULL,stdout=DEVNULL,stderr=DEVNULL)
  191. if self._subprocess.returncode == None: self._subprocess.kill()
  192. use = 'pigz'
  193. except OSError:
  194. try:
  195. self._subprocess = subprocess.Popen(["gzip"],stdin=DEVNULL,stdout=DEVNULL,stderr=DEVNULL)
  196. if self._subprocess.returncode == None: self._subprocess.kill()
  197. use = 'gzip'
  198. except OSError: use = 'internal'
  199. if use != 'internal' and decompressor != 'internal':
  200. if type(f) == str: self._subprocess = subprocess.Popen([ use , '--decompress','--stdout', f ], stdout=subprocess.PIPE, stderr=DEVNULL)
  201. else: self._subprocess = subprocess.Popen('{ printf "'+magic+'"; cat; } | ' + use + ' --decompress --stdout', stdin=f, shell=True, stdout=subprocess.PIPE, stderr=DEVNULL)
  202. time.sleep(1)
  203. if self._subprocess.poll() == None:
  204. data = self._subprocess.stdout.read(35536)
  205. self.file_decompressor = use
  206. self.file_bytes_read += len(data)
  207. while data:
  208. yield data
  209. data = self._subprocess.stdout.read(35536)
  210. self.file_bytes_read += len(data)
  211. self._file.close()
  212. DEVNULL.close()
  213. raise StopIteration
  214. # Python's gzip module can't read from a stream that doesn't support seek(), and the zlib module cannot read the bgzip format without a lot of help:
  215. self.file_decompressor = 'internal'
  216. raw_data = magic + self._file.read(65536)
  217. self.file_bytes_read = len(raw_data)
  218. internal_cache = []
  219. blocks_left_to_grab = 50
  220. bs = 0
  221. checkpoint = 0
  222. decompress = zlib.decompress
  223. while raw_data:
  224. if len(raw_data) - bs < 35536:
  225. raw_data = raw_data[bs:] + self._file.read(65536)
  226. self.file_bytes_read += len(raw_data) - bs
  227. bs = 0
  228. magic = raw_data[bs:bs+4]
  229. if not magic: break # a child's heart
  230. if magic != "\x1f\x8b\x08\x04": raise PybamError('\n\nThe input file is not in a format I understand. First four bytes: ' + repr(magic) + '\n')
  231. try:
  232. more_bs = bs + unpack("<H", raw_data[bs+16:bs+18])[0] +1
  233. internal_cache.append(decompress(raw_data[bs+18:more_bs-8],-15))
  234. bs = more_bs
  235. except: ## zlib doesnt have a nice exception for when things go wrong. just "error"
  236. header_data = magic + raw_data[bs+4:bs+12]
  237. header_size = 12
  238. extra_len = unpack("<H", header_data[-2:])[0]
  239. while header_size-12 < extra_len:
  240. header_data += raw_data[bs+12:bs+16]
  241. subfield_id = header_data[-4:-2]
  242. subfield_len = unpack("<H", header_data[-2:])[0]
  243. subfield_data = raw_data[bs+16:bs+16+subfield_len]
  244. header_data += subfield_data
  245. header_size += subfield_len + 4
  246. if subfield_id == 'BC': block_size = unpack("<H", subfield_data)[0]
  247. raw_data = raw_data[bs+16+subfield_len:bs+16+subfield_len+block_size-extra_len-19]
  248. crc_data = raw_data[bs+16+subfield_len+block_size-extra_len-19:bs+16+subfield_len+block_size-extra_len-19+8] # I have left the numbers in verbose, because the above try is the optimised code.
  249. bs = bs+16+subfield_len+block_size-extra_len-19+8
  250. zipped_data = header_data + raw_data + crc_data
  251. internal_cache.append(decompress(zipped_data,47)) # 31 works the same as 47.
  252. # Although the following in the bgzip code from biopython, its not needed if you let zlib decompress the whole zipped_data, header and crc, because it checks anyway (in C land)
  253. # I've left the manual crc checks in for documentation purposes:
  254. '''
  255. expected_crc = crc_data[:4]
  256. expected_size = unpack("<I", crc_data[4:])[0]
  257. if len(unzipped_data) != expected_size: print 'ERROR: Failed to unpack due to a Type 1 CRC error. Could the BAM be corrupted?'; exit()
  258. crc = zlib.crc32(unzipped_data)
  259. if crc < 0: crc = pack("<i", crc)
  260. else: crc = pack("<I", crc)
  261. if expected_crc != crc: print 'ERROR: Failed to unpack due to a Type 2 CRC error. Could the BAM be corrupted?'; exit()
  262. '''
  263. blocks_left_to_grab -= 1
  264. if blocks_left_to_grab == 0:
  265. yield ''.join(internal_cache)
  266. internal_cache = []
  267. blocks_left_to_grab = 50
  268. self._file.close()
  269. DEVNULL.close()
  270. if internal_cache != '': yield ''.join(internal_cache)
  271. raise StopIteration
  272. elif decompressor != False and decompressor != 'internal':
  273. # It wouldn't be safe to just print to the shell four random bytes from the beginning of a file, so instead it's
  274. # written to a temp file and cat'd. The idea here being that we trust the decompressor string as it was written by
  275. # someone with access to python, so it has system access anyway. The file/data, however, should not be trusted.
  276. magic_file = os.path.join(tempfile.mkdtemp(),'magic')
  277. with open(magic_file,'wb') as mf: mf.write(magic)
  278. if type(f) is str: self._subprocess = subprocess.Popen( decompressor.replace('{}',f), shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  279. else: self._subprocess = subprocess.Popen('{ cat "'+magic_file+'"; cat; } | ' + decompressor, stdin=self._file, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  280. self.file_decompressor = decompressor
  281. data = self._subprocess.stdout.read(35536)
  282. self.file_bytes_read += len(data)
  283. while data:
  284. yield data
  285. data = self._subprocess.stdout.read(35536)
  286. self.file_bytes_read += len(data)
  287. self._file.close()
  288. DEVNULL.close()
  289. raise StopIteration
  290. else:
  291. raise PybamError('\n\nThe input file is not in a format I understand. First four bytes: ' + repr(magic) + '\n')
  292. ## At this point, we know that whatever decompression method was used, a call to self._generator will return some uncompressed data.
  293. self._generator = generator()
  294. ## So lets parse the BAM header:
  295. header_cache = ''
  296. while len(header_cache) < 8: header_cache += next(self._generator)
  297. p_from = 0; p_to = 4
  298. if header_cache[p_from:p_to] != 'BAM\1':
  299. raise PybamError('\n\nInput file ' + self.file_name + ' does not appear to be a BAM file.\n')
  300. ## Parse the BAM header:
  301. p_from = p_to; p_to += 4
  302. length_of_header = unpack('<i',header_cache[p_from:p_to])[0]
  303. p_from = p_to; p_to += length_of_header
  304. while len(header_cache) < p_to: header_cache += next(self._generator)
  305. self.file_header = header_cache[p_from:p_to]
  306. p_from = p_to; p_to += 4
  307. while len(header_cache) < p_to: header_cache += next(self._generator)
  308. number_of_reference_sequences = unpack('<i',header_cache[p_from:p_to])[0]
  309. for _ in range(number_of_reference_sequences):
  310. p_from = p_to; p_to += 4
  311. while len(header_cache) < p_to: header_cache += next(self._generator)
  312. l_name = unpack('<l',header_cache[p_from:p_to])[0]
  313. p_from = p_to; p_to += l_name
  314. while len(header_cache) < p_to: header_cache += next(self._generator)
  315. self.file_chromosomes.append(header_cache[p_from:p_to -1])
  316. p_from = p_to; p_to += 4
  317. while len(header_cache) < p_to: header_cache += next(self._generator)
  318. self.file_chromosome_lengths[self.file_chromosomes[-1]] = unpack('<l',header_cache[p_from:p_to])[0]
  319. self.file_bytes_read = p_to
  320. self.file_binary_header = buffer(header_cache[:p_to])
  321. header_cache = header_cache[p_to:]
  322. # A quick check to make sure the header of this BAM file makes sense:
  323. chromosomes_from_header = []
  324. for line in self.file_header.split('\n'):
  325. if line.startswith('@SQ\tSN:'):
  326. chromosomes_from_header.append(line.split('\t')[1][3:])
  327. if chromosomes_from_header != self.file_chromosomes:
  328. raise PybamWarn('For some reason the BAM format stores the chromosome names in two locations,\n the ASCII text header we all know and love, viewable with samtools view -H, and another special binary header\n which is used to translate the chromosome refID (a number) into a chromosome RNAME when you do bam -> sam.\n\nThese two headers should always be the same, but apparently they are not:\nThe ASCII header looks like: ' + self.file_header + '\nWhile the binary header has the following chromosomes: ' + self.file_chromosomes + '\n')
  329. ## Variable parsing:
  330. def new_entry(header_cache):
  331. cache = header_cache # we keep a small cache of X bytes of decompressed BAM data, to smoothen out disk access.
  332. p = 0 # where the next alignment/entry starts in the cache
  333. while True:
  334. try:
  335. while len(cache) < p + 4: cache = cache[p:] + next(self._generator); p = 0 # Grab enough bytes to parse blocksize
  336. self.sam_block_size = unpack('<i',cache[p:p+4])[0]
  337. self.file_alignments_read += 1
  338. while len(cache) < p + 4 + self.sam_block_size:
  339. cache = cache[p:] + next(self._generator); p = 0 # Grab enough bytes to parse entry
  340. except StopIteration: break
  341. self.bam = cache[p:p + 4 + self.sam_block_size]
  342. p = p + 4 + self.sam_block_size
  343. yield self
  344. self._new_entry = new_entry(header_cache)
  345. def compile_parser(self,fields):
  346. temp_code = ''
  347. end_of_qname = False
  348. end_of_cigar = False
  349. end_of_seq = False
  350. end_of_qual = False
  351. dependencies = set(fields)
  352. if 'bam' in fields:
  353. fields[fields.index('bam')] = 'self.bam'
  354. if 'sam_block_size' in fields:
  355. fields[fields.index('sam_block_size')] = 'self.sam_block_size'
  356. if 'sam' in dependencies:
  357. dependencies.update(['sam_qname','sam_flag','sam_rname','sam_pos1','sam_mapq','sam_cigar_string','bam_refID','bam_next_refID','sam_rnext','sam_pnext1','sam_tlen','sam_seq','sam_qual','sam_tags_string'])
  358. if 'sam_tags_string' in dependencies:
  359. dependencies.update(['sam_tags_list'])
  360. if 'sam_pos1' in dependencies:
  361. temp_code += "\n sam_pos1 = (0 if sam_pos0 < 0 else sam_pos0 + 1)"
  362. dependencies.update(['sam_pos0'])
  363. if 'sam_pnext1' in dependencies:
  364. temp_code += "\n sam_pnext1 = (0 if sam_pnext0 < 0 else sam_pnext0 + 1)"
  365. dependencies.update(['sam_pnext0'])
  366. if 'sam_qname' in dependencies or 'bam_qname' in dependencies:
  367. temp_code += "\n _end_of_qname = 36 + sam_l_read_name"
  368. dependencies.update(['sam_l_read_name'])
  369. end_of_qname = True
  370. if 'sam_cigar_string' in dependencies or 'sam_cigar_list' in dependencies or 'bam_cigar' in dependencies:
  371. if end_of_qname:
  372. pass
  373. else:
  374. temp_code += "\n _end_of_qname = 36 + sam_l_read_name"
  375. temp_code += "\n _end_of_cigar = _end_of_qname + (4*sam_n_cigar_op)"
  376. dependencies.update(['sam_l_read_name','sam_n_cigar_op'])
  377. end_of_cigar = True
  378. if 'sam_seq' in dependencies or 'bam_seq' in dependencies:
  379. if end_of_cigar:
  380. pass
  381. elif end_of_qname:
  382. temp_code += "\n _end_of_cigar = _end_of_qname + (4*sam_n_cigar_op)"
  383. else:
  384. temp_code += "\n _end_of_cigar = 36 + sam_l_read_name + (4*sam_n_cigar_op)"
  385. temp_code += "\n _end_of_seq = _end_of_cigar + (-((-sam_l_seq)//2))"
  386. dependencies.update(['sam_l_seq','sam_n_cigar_op','sam_l_read_name'])
  387. end_of_seq = True
  388. if 'sam_qual' in dependencies or 'bam_qual' in dependencies:
  389. if end_of_seq:
  390. pass
  391. elif end_of_cigar:
  392. temp_code += "\n _end_of_seq = _end_of_cigar + (-((-sam_l_seq)//2))"
  393. elif end_of_qname:
  394. temp_code += "\n _end_of_seq = _end_of_qname + (4*sam_n_cigar_op) + (-((-sam_l_seq)//2))"
  395. else:
  396. temp_code += "\n _end_of_seq = 36 + sam_l_read_name + (4*sam_n_cigar_op) + (-((-sam_l_seq)//2))"
  397. temp_code += "\n _end_of_qual = _end_of_seq + sam_l_seq"
  398. dependencies.update(['sam_l_seq','sam_n_cigar_op','sam_l_read_name'])
  399. end_of_qual = True
  400. if 'sam_tags_list' in dependencies or 'bam_tags' in dependencies:
  401. if end_of_qual:
  402. pass
  403. elif end_of_seq:
  404. temp_code += "\n _end_of_qual = _end_of_seq + sam_l_seq"
  405. elif end_of_cigar:
  406. temp_code += "\n _end_of_qual = _end_of_cigar + (-((-sam_l_seq)//2)) + sam_l_seq"
  407. elif end_of_qname:
  408. temp_code += "\n _end_of_qual = _end_of_qname + (4*sam_n_cigar_op) + (-((-sam_l_seq)//2)) + sam_l_seq"
  409. else:
  410. temp_code += "\n _end_of_qual = 36 + sam_l_read_name + (4*sam_n_cigar_op) + (-((-sam_l_seq)//2)) + sam_l_seq"
  411. dependencies.update(['sam_l_seq','sam_n_cigar_op','sam_l_read_name'])
  412. if 'sam_rname' in dependencies:
  413. temp_code += "\n sam_rname = '*' if sam_refID < 0 else self.file_chromosomes[sam_refID]"
  414. dependencies.update(['sam_refID'])
  415. if 'sam_rnext' in dependencies:
  416. temp_code += "\n sam_rnext = '*' if sam_next_refID < 0 else self.file_chromosomes[sam_next_refID]"
  417. dependencies.update(['sam_next_refID'])
  418. ## First we figure out what data from the static portion of the BAM entry we'll need:
  419. tmp = {}
  420. tmp['code'] = 'def parser(self):\n from array import array\n from struct import unpack\n for _ in self._new_entry:'
  421. tmp['last_start'] = None
  422. tmp['name_list'] = []
  423. tmp['dtype_list'] = []
  424. def pack_up(name,dtype,length,end,tmp):
  425. if name in dependencies:
  426. if tmp['last_start'] == None:
  427. tmp['last_start'] = end - length
  428. tmp['name_list'].append(name)
  429. tmp['dtype_list'].append(dtype)
  430. elif tmp['last_start'] != None:
  431. tmp['code'] += '\n ' + ', '.join(tmp['name_list']) + ' = unpack("<' + ''.join(tmp['dtype_list']) + '",self.bam[' + str(tmp['last_start']) + ':' + str(end-length) + '])'
  432. if len(tmp['dtype_list']) == 1:
  433. tmp['code'] += '[0]'
  434. tmp['last_start'] = None
  435. tmp['name_list'] = []
  436. tmp['dtype_list'] = []
  437. pack_up('sam_refID', 'i',4, 8,tmp)
  438. pack_up('sam_pos0', 'i',4,12,tmp)
  439. pack_up('sam_l_read_name', 'B',1,13,tmp)
  440. pack_up('sam_mapq', 'B',1,14,tmp)
  441. pack_up('sam_bin', 'H',2,16,tmp)
  442. pack_up('sam_n_cigar_op', 'H',2,18,tmp)
  443. pack_up('sam_flag', 'H',2,20,tmp)
  444. pack_up('sam_l_seq', 'i',4,24,tmp)
  445. pack_up('sam_next_refID', 'i',4,28,tmp)
  446. pack_up('sam_pnext0', 'i',4,32,tmp)
  447. pack_up('sam_tlen', 'i',4,36,tmp)
  448. pack_up( None, None,0,36,tmp) # To add anything not yet added.
  449. code = tmp['code']
  450. del tmp
  451. code += temp_code
  452. # Fixed-length BAM data (where we just grab the bytes, we dont unpack) can, however, be grabbed individually.
  453. if 'bam_block_size' in dependencies: code += "\n bam_block_size = self.bam[0 : 4 ]"
  454. if 'bam_refID' in dependencies: code += "\n bam_refID = self.bam[4 : 8 ]"
  455. if 'bam_pos' in dependencies: code += "\n bam_pos = self.bam[8 : 12 ]"
  456. if 'bam_l_read_name' in dependencies: code += "\n bam_l_read_name = self.bam[12 : 13 ]"
  457. if 'bam_mapq' in dependencies: code += "\n bam_mapq = self.bam[13 : 14 ]"
  458. if 'bam_bin' in dependencies: code += "\n bam_bin = self.bam[14 : 16 ]"
  459. if 'bam_n_cigar_op' in dependencies: code += "\n bam_n_cigar_op = self.bam[16 : 18 ]"
  460. if 'bam_flag' in dependencies: code += "\n bam_flag = self.bam[18 : 20 ]"
  461. if 'bam_l_seq' in dependencies: code += "\n bam_l_seq = self.bam[20 : 24 ]"
  462. if 'bam_next_refID' in dependencies: code += "\n bam_next_refID = self.bam[24 : 28 ]"
  463. if 'bam_pnext' in dependencies: code += "\n bam_pnext = self.bam[28 : 32 ]"
  464. if 'bam_tlen' in dependencies: code += "\n bam_tlen = self.bam[32 : 36 ]"
  465. if 'bam_qname' in dependencies: code += "\n bam_qname = self.bam[36 : _end_of_qname ]"
  466. if 'bam_cigar' in dependencies: code += "\n bam_cigar = self.bam[_end_of_qname : _end_of_cigar ]"
  467. if 'bam_seq' in dependencies: code += "\n bam_seq = self.bam[_end_of_cigar : _end_of_seq ]"
  468. if 'bam_qual' in dependencies: code += "\n bam_qual = self.bam[_end_of_seq : _end_of_qual ]"
  469. if 'bam_tags' in dependencies: code += "\n bam_tags = self.bam[_end_of_qual : ]"
  470. if 'sam_qname' in dependencies:
  471. if 'bam_qname' in dependencies: code += "\n sam_qname = bam_qname[:-1]"
  472. else: code += "\n sam_qname = self.bam[36 : _end_of_qname -1 ]"
  473. if 'sam_cigar_list' in dependencies:
  474. if 'bam_cigar' in dependencies: code += "\n sam_cigar_list = [( cig >> 4 , cigar_codes[cig & 0b1111]) for cig in array('I', bam_cigar) ]"
  475. else: code += "\n sam_cigar_list = [( cig >> 4 , cigar_codes[cig & 0b1111]) for cig in array('I', self.bam[_end_of_qname : _end_of_cigar]) ]"
  476. if 'sam_cigar_string'in dependencies:
  477. if 'bam_cigar' in dependencies: code += "\n sam_cigar_string = ''.join([ str(cig >> 4) + cigar_codes[cig & 0b1111] for cig in array('I', bam_cigar)])"
  478. else: code += "\n sam_cigar_string = ''.join([ str(cig >> 4) + cigar_codes[cig & 0b1111] for cig in array('I', self.bam[_end_of_qname : _end_of_cigar]) ])"
  479. if 'sam_seq' in dependencies:
  480. if 'bam_seq' in dependencies: code += "\n sam_seq = ''.join( [ dna_codes[dna >> 4] + dna_codes[dna & 0b1111] for dna in array('B', bam_seq)])[:sam_l_seq]"
  481. else: code += "\n sam_seq = ''.join( [ dna_codes[dna >> 4] + dna_codes[dna & 0b1111] for dna in array('B', self.bam[_end_of_cigar : _end_of_seq])])[:sam_l_seq]"
  482. if 'sam_qual' in dependencies:
  483. if 'bam_qual' in dependencies: code += "\n sam_qual = ''.join( [ chr(ord(quality) + 33) for quality in bam_qual ])"
  484. else: code += "\n sam_qual = ''.join( [ chr(ord(quality) + 33) for quality in self.bam[_end_of_seq : _end_of_qual ]])"
  485. if 'sam_tags_list' in dependencies:
  486. code += '''
  487. sam_tags_list = []
  488. offset = _end_of_qual
  489. while offset != len(self.bam):
  490. tag_name = self.bam[offset:offset+2]
  491. tag_type = self.bam[offset+2]
  492. if tag_type == 'Z':
  493. offset_end = self.bam.index('\\0',offset+3)+1
  494. tag_data = self.bam[offset+3:offset_end-1]
  495. elif tag_type in CtoPy:
  496. offset_end = offset+3+py4py[tag_type]
  497. tag_data = unpack(CtoPy[tag_type],self.bam[offset+3:offset_end])[0]
  498. elif tag_type == 'B':
  499. offset_end = offset+8+(unpack('<i',self.bam[offset+4:offset+8])[0]*py4py[self.bam[offset+3]])
  500. tag_data = array(self.bam[offset+3] , self.bam[offset+8:offset_end] )
  501. else:
  502. print 'PYBAM ERROR: I dont know how to parse BAM tags in this format: ',repr(tag_type)
  503. print ' This is simply because I never saw this kind of tag during development.'
  504. print ' If you could mail the following chunk of text to john at john.uk.com, i will fix this up for everyone :)'
  505. print repr(tag_type),repr(self.bam[offset+3:end])
  506. exit()
  507. sam_tags_list.append((tag_name,tag_type,tag_data))
  508. offset = offset_end'''
  509. if 'sam_tags_string' in dependencies:
  510. code += "\n sam_tags_string = '\t'.join(A + ':' + ('i' if B in 'cCsSI' else B) + ':' + ((C.typecode + ',' + ','.join(map(str,C))) if type(C)isarray else str(C)) for A,B,C in self.sam_tags_list)"
  511. if 'sam' in dependencies:
  512. code += "\n sam = sam_qname + '\t' + str(sam_flag) + '\t' + sam_rname + '\t' + str(sam_pos1) + '\t' + str(sam_mapq) + '\t' + ('*' if sam_cigar_string is '' else sam_cigar_string) + '\t' + ('=' if bam_refID == bam_next_refID else sam_rnext) + '\t' + str(sam_pnext1) + '\t' + str(sam_tlen) + '\t' + sam_seq + '\t' + sam_qual + '\t' + sam_tags_string"
  513. code += '\n yield ' + ','.join([x for x in fields]) + '\n'
  514. self._static_parser_code = code # "code" is the static parser's code as a string (a function called "parser")
  515. exec_dict = { # This dictionary stores things the exec'd code needs to know about, and will store the compiled function after exec()
  516. 'unpack':unpack,
  517. 'array':array,
  518. 'dna_codes':dna_codes,
  519. 'CtoPy':CtoPy,
  520. 'py4py':py4py,
  521. 'cigar_codes':cigar_codes
  522. }
  523. exec(code in exec_dict) # exec() compiles "code" to real code, creating the "parser" function and adding it to exec_dict['parser']
  524. return exec_dict['parser']
  525. if fields:
  526. static_parser = compile_parser(self,fields)(self)
  527. def next_read(): return next(static_parser)
  528. else:
  529. def next_read(): return next(self._new_entry)
  530. self.next = next_read
  531. def __iter__(self): return self
  532. def __str__(self): return self.sam
  533. ## Methods to pull out raw bam data from entry (so still in its binary encoding). This can be helpful in some scenarios.
  534. @property
  535. def bam_block_size(self): return self.bam[ : 4 ]
  536. @property
  537. def bam_refID(self): return self.bam[ 4 : 8 ]
  538. @property
  539. def bam_pos(self): return self.bam[ 8 : 12 ]
  540. @property
  541. def bam_l_read_name(self): return self.bam[ 12 : 13 ]
  542. @property
  543. def bam_mapq(self): return self.bam[ 13 : 14 ]
  544. @property
  545. def bam_bin(self): return self.bam[ 14 : 16 ]
  546. @property
  547. def bam_n_cigar_op(self): return self.bam[ 16 : 18 ]
  548. @property
  549. def bam_flag(self): return self.bam[ 18 : 20 ]
  550. @property
  551. def bam_l_seq(self): return self.bam[ 20 : 24 ]
  552. @property
  553. def bam_next_refID(self): return self.bam[ 24 : 28 ]
  554. @property
  555. def bam_pnext(self): return self.bam[ 28 : 32 ]
  556. @property
  557. def bam_tlen(self): return self.bam[ 32 : 36 ]
  558. @property
  559. def bam_qname(self): return self.bam[ 36 : self._end_of_qname ]
  560. @property
  561. def bam_cigar(self): return self.bam[ self._end_of_qname : self._end_of_cigar ]
  562. @property
  563. def bam_seq(self): return self.bam[ self._end_of_cigar : self._end_of_seq ]
  564. @property
  565. def bam_qual(self): return self.bam[ self._end_of_seq : self._end_of_qual ]
  566. @property
  567. def bam_tags(self): return self.bam[ self._end_of_qual : ]
  568. @property
  569. def sam_refID(self): return unpack( '<i', self.bam[ 4 : 8 ] )[0]
  570. @property
  571. def sam_pos0(self): return unpack( '<i', self.bam[ 8 : 12 ] )[0]
  572. @property
  573. def sam_l_read_name(self): return unpack( '<B', self.bam[ 12 : 13 ] )[0]
  574. @property
  575. def sam_mapq(self): return unpack( '<B', self.bam[ 13 : 14 ] )[0]
  576. @property
  577. def sam_bin(self): return unpack( '<H', self.bam[ 14 : 16 ] )[0]
  578. @property
  579. def sam_n_cigar_op(self): return unpack( '<H', self.bam[ 16 : 18 ] )[0]
  580. @property
  581. def sam_flag(self): return unpack( '<H', self.bam[ 18 : 20 ] )[0]
  582. @property
  583. def sam_l_seq(self): return unpack( '<i', self.bam[ 20 : 24 ] )[0]
  584. @property
  585. def sam_next_refID(self): return unpack( '<i', self.bam[ 24 : 28 ] )[0]
  586. @property
  587. def sam_pnext0(self): return unpack( '<i', self.bam[ 28 : 32 ] )[0]
  588. @property
  589. def sam_tlen(self): return unpack( '<i', self.bam[ 32 : 36 ] )[0]
  590. @property
  591. def sam_qname(self): return self.bam[ 36 : self._end_of_qname -1 ] # -1 to remove trailing NUL byte
  592. @property
  593. def sam_cigar_list(self): return [ (cig >> 4 , cigar_codes[cig & 0b1111] ) for cig in array('I', self.bam[self._end_of_qname : self._end_of_cigar ])]
  594. @property
  595. def sam_cigar_string(self): return ''.join( [ str(cig >> 4) + cigar_codes[cig & 0b1111] for cig in array('I', self.bam[self._end_of_qname : self._end_of_cigar ])])
  596. @property
  597. def sam_seq(self): return ''.join( [ dna_codes[dna >> 4] + dna_codes[dna & 0b1111] for dna in array('B', self.bam[self._end_of_cigar : self._end_of_seq ])])[:self.sam_l_seq] # As DNA is 4 bits packed 2-per-byte, there might be a trailing '0000', so we can either
  598. @property
  599. def sam_qual(self): return ''.join( [ chr(ord(quality) + 33) for quality in self.bam[self._end_of_seq : self._end_of_qual ]])
  600. @property
  601. def sam_tags_list(self):
  602. result = []
  603. offset = self._end_of_qual
  604. while offset != len(self.bam):
  605. tag_name = self.bam[offset:offset+2]
  606. tag_type = self.bam[offset+2]
  607. if tag_type == 'Z':
  608. offset_end = self.bam.index('\x00',offset+3)+1
  609. tag_data = self.bam[offset+3:offset_end-1]
  610. elif tag_type in CtoPy:
  611. offset_end = offset+3+py4py[tag_type]
  612. tag_data = unpack(CtoPy[tag_type],self.bam[offset+3:offset_end])[0]
  613. elif tag_type == 'B':
  614. offset_end = offset+8+(unpack('<i',self.bam[offset+4:offset+8])[0]*py4py[self.bam[offset+3]])
  615. tag_data = array(self.bam[offset+3] , self.bam[offset+8:offset_end] )
  616. else:
  617. print('PYBAM ERROR: I dont know how to parse BAM tags in this format: ',repr(tag_type))
  618. print(' This is simply because I never saw this kind of tag during development.')
  619. print(' If you could mail the following chunk of text to john at john.uk.com, ill fix this up :)')
  620. print(repr(tag_type),repr(self.bam[offset+3:end]))
  621. exit()
  622. result.append((tag_name,tag_type,tag_data))
  623. offset = offset_end
  624. return result
  625. @property
  626. def sam_tags_string(self):
  627. return '\t'.join(A + ':' + ('i' if B in 'cCsSI' else B) + ':' + ((C.typecode + ',' + ','.join(map(str,C))) if type(C)==array else str(C)) for A,B,C in self.sam_tags_list)
  628. ## BONUS methods - methods that mimic how samtools works.
  629. @property
  630. def sam_pos1(self): return 0 if self.sam_pos0 < 0 else self.sam_pos0 + 1
  631. @property
  632. def sam_pnext1(self): return 0 if self.sam_pnext0 < 0 else self.sam_pnext0 + 1
  633. @property
  634. def sam_rname(self): return '*' if self.sam_refID < 0 else self.file_chromosomes[self.sam_refID ]
  635. @property
  636. def sam_rnext(self): return '*' if self.sam_next_refID < 0 else self.file_chromosomes[self.sam_next_refID]
  637. @property
  638. def sam(self): return (
  639. self.sam_qname + '\t' +
  640. str(self.sam_flag) + '\t' +
  641. self.sam_rname + '\t' +
  642. str(self.sam_pos1) + '\t' +
  643. str(self.sam_mapq) + '\t' +
  644. ('*' if self.sam_cigar_string == '' else self.sam_cigar_string) + '\t' +
  645. ('=' if self.bam_refID == self.bam_next_refID else self.sam_rnext) + '\t' +
  646. str(self.sam_pnext1) + '\t' +
  647. str(self.sam_tlen) + '\t' +
  648. self.sam_seq + '\t' +
  649. self.sam_qual + '\t' +
  650. self.sam_tags_string
  651. )
  652. ## Internal methods - methods used to calculate where variable-length blocks start/end
  653. @property
  654. def _end_of_qname(self): return self.sam_l_read_name + 36 # fixed-length stuff at the beginning takes up 36 bytes.
  655. @property
  656. def _end_of_cigar(self): return self._end_of_qname + (4*self.sam_n_cigar_op) # 4 bytes per n_cigar_op
  657. @property
  658. def _end_of_seq(self): return self._end_of_cigar + (-((-self.sam_l_seq)//2)) # {blurgh}
  659. @property
  660. def _end_of_qual(self): return self._end_of_seq + self.sam_l_seq # qual has the same length as seq
  661. def __del__(self):
  662. if self._subprocess.returncode == None: self._subprocess.kill()
  663. self._file.close()
  664. class PybamWarn(Exception): pass
  665. class PybamError(Exception): pass