Mercurial > repos > thondeboer > neat_genreads

annotate py/biopython_modified_bgzf.py @ 10:7d10b55965c9 draft default tip

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
author thondeboer
date Wed, 16 May 2018 17:02:51 -0400
parents 6e75a84e9338
children
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
0
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
1 #!/usr/bin/env python
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
2 # Copyright 2010-2013 by Peter Cock.
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
3 # All rights reserved.
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
4 # This code is part of the Biopython distribution and governed by its
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
5 # license. Please see the LICENSE file that should have been included
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
6 # as part of this package.
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
7
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
8 """ ############################################################################
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
9 ####### #######
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
10 ####### 06/02/2015: #######
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
11 ####### - I picked out the bits and pieces of code needed #######
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
12 ####### to write BAM files, removed python 3.0 compatibility #######
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
13 ####### #######
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
14 ############################################################################ """
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
15
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
16 import zlib
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
17 import struct
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
18
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
19 _bgzf_header = b"\x1f\x8b\x08\x04\x00\x00\x00\x00\x00\xff\x06\x00\x42\x43\x02\x00"
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
20 _bgzf_eof = b"\x1f\x8b\x08\x04\x00\x00\x00\x00\x00\xff\x06\x00\x42\x43\x02\x00\x1b\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00"
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
21
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
22 class BgzfWriter(object):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
23
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
24 def __init__(self, filename=None, mode="w", fileobj=None, compresslevel=6):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
25 if fileobj:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
26 assert filename is None
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
27 handle = fileobj
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
28 else:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
29 if "w" not in mode.lower() \
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
30 and "a" not in mode.lower():
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
31 raise ValueError("Must use write or append mode, not %r" % mode)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
32 if "a" in mode.lower():
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
33 handle = open(filename, "ab")
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
34 else:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
35 handle = open(filename, "wb")
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
36 self._text = "b" not in mode.lower()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
37 self._handle = handle
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
38 self._buffer = b""
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
39 self.compresslevel = compresslevel
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
40
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
41 def _write_block(self, block):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
42 start_offset = self._handle.tell()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
43 assert len(block) <= 65536
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
44 # Giving a negative window bits means no gzip/zlib headers, -15 used in samtools
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
45 c = zlib.compressobj(self.compresslevel,
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
46 zlib.DEFLATED,
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
47 -15,
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
48 zlib.DEF_MEM_LEVEL,
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
49 0)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
50 compressed = c.compress(block) + c.flush()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
51 del c
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
52 assert len(compressed) < 65536, "TODO - Didn't compress enough, try less data in this block"
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
53 crc = zlib.crc32(block)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
54 # Should cope with a mix of Python platforms...
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
55 if crc < 0:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
56 crc = struct.pack("<i", crc)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
57 else:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
58 crc = struct.pack("<I", crc)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
59 bsize = struct.pack("<H", len(compressed) + 25) # includes -1
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
60 crc = struct.pack("<I", zlib.crc32(block) & 0xffffffff)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
61 uncompressed_length = struct.pack("<I", len(block))
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
62 data = _bgzf_header + bsize + compressed + crc + uncompressed_length
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
63 self._handle.write(data)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
64
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
65 def write(self, data):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
66 data_len = len(data)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
67 if len(self._buffer) + data_len < 65536:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
68 self._buffer += data
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
69 return
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
70 else:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
71 self._buffer += data
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
72 while len(self._buffer) >= 65536:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
73 self._write_block(self._buffer[:65536])
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
74 self._buffer = self._buffer[65536:]
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
75
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
76 def flush(self):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
77 while len(self._buffer) >= 65536:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
78 self._write_block(self._buffer[:65535])
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
79 self._buffer = self._buffer[65535:]
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
80 self._write_block(self._buffer)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
81 self._buffer = b""
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
82 self._handle.flush()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
83
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
84 def close(self):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
85 """Flush data, write 28 bytes empty BGZF EOF marker, and close the BGZF file."""
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
86 if self._buffer:
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
87 self.flush()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
88 # samtools will look for a magic EOF marker, just a 28 byte empty BGZF block,
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
89 # and if it is missing warns the BAM file may be truncated. In addition to
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
90 # samtools writing this block, so too does bgzip - so we should too.
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
91 self._handle.write(_bgzf_eof)
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
92 self._handle.flush()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
93 self._handle.close()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
94
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
95 def __enter__(self):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
96 return self
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
97
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
98 def __exit__(self, type, value, traceback):
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
99 self.close()
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
100
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
101
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
102 if __name__ == "__main__":
6e75a84e9338 planemo upload commit e96b43f96afce6a7b7dfd4499933aad7d05c955e-dirty
thondeboer
parents:
diff changeset
103 pass