File size: 3,675 Bytes
a65138c | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | # -*- coding: utf-8 -*-
from __future__ import with_statement
import logging
import os.path
from zipfile import ZIP_DEFLATED
from path import split as path_split
from path import get_ancestors as path_ancestors
logger = logging.getLogger(__name__)
def zipfile_nodes(zipfile):
seen = set()
for path in zipfile.namelist():
for anc_path in path_ancestors(path):
if anc_path not in seen:
yield anc_path, ZipFileFolder(zipfile, anc_path)
seen.add(anc_path)
if path not in seen:
if path.endswith('/'):
yield path, ZipFileFolder(zipfile, path)
else:
yield path, ZipFileFile(zipfile, path)
seen.add(path)
class ZipFileNode(object):
def __init__(self, zipfile, path):
self.zipfile = zipfile
self.path = path
class ZipFileFile(ZipFileNode):
def open(self, mode='r', compress_type=ZIP_DEFLATED):
if mode.startswith('r'):
return ZipFileStream(self.zipfile.open(self.path, mode))
elif mode == 'w':
return ZipFileStream(ZipFileWritable(self.zipfile, self.path,
compress_type=compress_type))
def put(self, filesystem_path, compress_type=ZIP_DEFLATED):
self.zipfile.write(filesystem_path, self.path, compress_type)
class ZipFileFolder(ZipFileNode):
def childs(self):
prefix = path_split(self.path)
prefix_len = len(prefix)
for path, node in zipfile_nodes(self.zipfile):
path_segments = path_split(path)
if len(path_segments) == prefix_len + 1:
if path_segments[:prefix_len] == prefix:
yield path_segments[-1], node
def __iter__(self):
for name, node in self.childs():
yield name
def __getitem__(self, name):
for node_name, node in self.childs():
if node_name == name:
return node
raise KeyError(name)
def file(self, name):
path = os.path.join(self.path, name)
return ZipFileFile(self.zipfile, path)
def folder(self, name):
path = os.path.join(self.path, name)
return ZipFileFolder(self.zipfile, path)
class ZipFileStorage(ZipFileFolder):
def __init__(self, *args, **kwargs):
import zipfile
self.zipfile = zipfile.ZipFile(*args, **kwargs)
self.path = ''
def childs(self):
for path, node in zipfile_nodes(self.zipfile):
path_segments = path_split(path)
if len(path_segments) == 1:
yield path_segments[-1], node
def close(self):
self.zipfile.close()
def __enter__(self):
return self
def __exit__(self, *args, **kwargs):
self.close()
class ZipFileStream(object):
def __init__(self, stream):
self.stream = stream
def __getattr__(self, name):
return getattr(self.stream, name)
def __enter__(self):
return self.stream
def __exit__(self, *args, **kwargs):
self.close()
class ZipFileWritable(object):
def __init__(self, zipfile, path, compress_type=ZIP_DEFLATED):
self.zipfile = zipfile
self.path = path
self.compress_type = compress_type
import tempfile
fd, tmp_path = tempfile.mkstemp()
self.tmp_f = os.fdopen(fd, 'w')
self.tmp_path = tmp_path
def __getattr__(self, name):
return getattr(self.tmp_f, name)
def close(self):
self.tmp_f.close()
self.zipfile.write(self.tmp_path, self.path, self.compress_type)
os.unlink(self.tmp_path)
|