forked from gtsystem/python-remotezip
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathremotezip.py
More file actions
172 lines (133 loc) · 5.38 KB
/
remotezip.py
File metadata and controls
172 lines (133 loc) · 5.38 KB
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import io
import zipfile
import requests
__all__ = ['RemoteIOError', 'RemoteZip']
class OutOfBound(Exception):
pass
class RemoteIOError(Exception):
pass
class PartialBuffer:
def __init__(self, buffer, offset, size, stream):
self.buffer = buffer if stream else io.BytesIO(buffer.read())
self.offset = offset
self.size = size
self.position = offset
self.stream = stream
def __repr__(self):
return "<PartialBuffer off=%s size=%s stream=%s>" % (self.offset, self.size, self.stream)
def read(self, size=0):
if size == 0:
size = self.offset + self.size - self.position
content = self.buffer.read(size)
self.position = self.offset + self.buffer.tell()
return content
def close(self):
if not self.buffer.closed:
self.buffer.close()
if hasattr(self.buffer, 'release_conn'):
self.buffer.release_conn()
def seek(self, offset, whence):
if whence == 2:
self.position = self.size + self.offset + offset
elif whence == 0:
self.position = offset
else:
self.position += offset
relative_position = self.position - self.offset
if relative_position < 0 or relative_position >= self.size:
raise OutOfBound("Position out of buffer bound")
if self.stream:
buff_pos = self.buffer.tell()
if relative_position < buff_pos:
raise OutOfBound("Negative seek not supported")
skip_bytes = relative_position - buff_pos
if skip_bytes == 0:
return self.position
self.buffer.read(skip_bytes)
else:
self.buffer.seek(relative_position)
return self.position
class RemoteIO(io.IOBase):
def __init__(self, fetch_fun, initial_buffer_size=64*1024):
self.fetch_fun = fetch_fun
self.initial_buffer_size = initial_buffer_size
self.buffer = None
self.file_size = None
self.position = None
self._seek_succeeded = False
self.member_pos2size = None
def set_pos2size(self, pos2size):
self.member_pos2size = pos2size
def read(self, size=0):
if size == 0:
size = self.file_size - self.buffer.position
if not self._seek_succeeded:
if self.member_pos2size is None:
fetch_size = size
stream = False
else:
fetch_size = self.member_pos2size[self.buffer.position]
stream = True
self._seek_succeeded = True
self.buffer.close()
self.buffer = self.fetch_fun((self.buffer.position, self.buffer.position + fetch_size -1), stream=stream)
return self.buffer.read(size)
def seek(self, offset, whence=0):
if whence == 2 and self.file_size is None:
size = self.initial_buffer_size
self.buffer = self.fetch_fun((-size, None), stream=False)
self.file_size = self.buffer.size + self.buffer.position
try:
pos = self.buffer.seek(offset, whence)
self._seek_succeeded = True
return pos
except OutOfBound:
self._seek_succeeded = False
return self.buffer.position # we ignore the issue here, we will check if buffer is fine during read
def tell(self):
return self.buffer.position
def close(self):
if self.buffer:
self.buffer.close()
self.buffer = None
class RemoteZip(zipfile.ZipFile):
def __init__(self, url, initial_buffer_size=64*1024, **kwargs):
self.kwargs = kwargs
self.url = url
rio = RemoteIO(self.fetch_fun, initial_buffer_size)
super(RemoteZip, self).__init__(rio)
rio.set_pos2size(self.get_position2size())
def get_position2size(self):
ilist = self.infolist()
if len(ilist) == 0:
return {}
position2size = {ilist[-1].header_offset: self.start_dir - ilist[-1].header_offset}
for i in range(len(ilist) - 1):
m1, m2 = ilist[i: i+2]
position2size[m1.header_offset] = m2.header_offset - m1.header_offset
return position2size
@staticmethod
def make_buffer(io_buffer, content_range_header, stream):
range_min, range_max = content_range_header.split("/")[0][6:].split("-")
range_min, range_max = int(range_min), int(range_max)
return PartialBuffer(io_buffer, range_min, range_max - range_min + 1, stream)
@staticmethod
def make_header(range_min, range_max):
if range_max is None:
return "bytes=%s%s" % (range_min, '' if range_min < 0 else '-')
return "bytes=%s-%s" % (range_min, range_max)
@staticmethod
def request(url, range_header, kwargs):
kwargs['headers'] = headers = dict(kwargs.get('headers', {}))
headers['Range'] = range_header
res = requests.get(url, stream=True, **kwargs)
res.raise_for_status()
return res.raw, res.headers
def fetch_fun(self, data_range, stream=False):
range_header = self.make_header(*data_range)
kwargs = dict(self.kwargs)
try:
res, headers = self.request(self.url, range_header, kwargs)
return self.make_buffer(res, headers['Content-Range'], stream=stream)
except IOError as e:
raise RemoteIOError(str(e))