Py3: Return bytes from FileGenerator

and use __getitem__ instead of __getslice__
This commit is contained in:
Shadab Zafar 2016-06-02 13:00:44 +05:30
parent a04d7fd166
commit b510616c69
2 changed files with 11 additions and 9 deletions

View File

@ -77,17 +77,17 @@ class FileGenerator(object):
def __init__(self, path): def __init__(self, path):
self.path = path self.path = path
self.fp = file(path, "rb") self.fp = open(path, "rb")
self.map = mmap.mmap(self.fp.fileno(), 0, access=mmap.ACCESS_READ) self.map = mmap.mmap(self.fp.fileno(), 0, access=mmap.ACCESS_READ)
def __len__(self): def __len__(self):
return len(self.map) return len(self.map)
def __getitem__(self, x): def __getitem__(self, x):
if isinstance(x, slice):
return self.map.__getitem__(x) return self.map.__getitem__(x)
# A slice of length 1 returns a byte object (not an integer)
def __getslice__(self, a, b): return self.map.__getitem__(slice(x, x+1 or self.map.size()))
return self.map.__getslice__(a, b)
def __repr__(self): def __repr__(self):
return "<%s" % self.path return "<%s" % self.path

View File

@ -19,13 +19,15 @@ def test_filegenerator():
with tutils.tmpdir() as t: with tutils.tmpdir() as t:
path = os.path.join(t, "foo") path = os.path.join(t, "foo")
f = open(path, "wb") f = open(path, "wb")
f.write("x" * 10000) f.write(b"x" * 10000)
f.close() f.close()
g = generators.FileGenerator(path) g = generators.FileGenerator(path)
assert len(g) == 10000 assert len(g) == 10000
assert g[0] == "x" assert g[0] == b"x"
assert g[-1] == "x" assert g[-1] == b"x"
assert g[0:5] == "xxxxx" assert g[0:5] == b"xxxxx"
assert len(g[1:10]) == 9
assert len(g[10000:10001]) == 0
assert repr(g) assert repr(g)
# remove all references to FileGenerator instance to close the file # remove all references to FileGenerator instance to close the file
# handle. # handle.