/
compat.py
34 lines (31 loc) · 1.04 KB
/
compat.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import zipfile
try:
from io import BytesIO
except ImportError:
try:
from cStringIO import cStringIO as _BytesIO
except ImportError:
from StringIO import StringIO as _BytesIO
class BytesIO(_BytesIO):
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
if hasattr(zipfile.ZipFile, "__enter__"):
ZipFile = zipfile.ZipFile
else:
# Make compatible with "with_statement"
# and also implement compression argument
class ZipFile(zipfile.ZipFile):
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self.close()
def writestr(self, info, bytes, compression=None):
if compression is not None:
_compression = self.compression
self.compression = compression
zipfile.ZipFile.writestr(self, info, bytes)
if compression is not None:
self.compression = _compression
__all__ = ["BytesIO", "ZipFile"]