diff --git a/Lib/mailbox.py b/Lib/mailbox.py --- a/Lib/mailbox.py +++ b/Lib/mailbox.py @@ -1864,97 +1864,137 @@ """Message with MMDF-specific properties.""" -class _ProxyFile: - """A read-only wrapper of a file.""" +class _ProxyFile(io.RawIOBase): + """A io.RawIOBase inheriting read-only wrapper for a seekable file. + It supports __iter__() and the context-manager protocol. + """ + def __init__(self, file, pos=None): + """If pos is not None then the file will keep track of its position.""" + io.RawIOBase.__init__(self) + self._file = file + self._pos = file.tell() if pos is None else pos + self._close = True + self._is_open = True - def __init__(self, f, pos=None): - """Initialize a _ProxyFile.""" - self._file = f - if pos is None: - self._pos = f.tell() + def _set_noclose(self): + """Subclass hook - use to avoid closing internal file object.""" + self._close = False + + def _closed_check(self): + """Raise ValueError if not open.""" + if not self._is_open: + raise ValueError('I/O operation on closed file') + + def close(self): + if self._close: + self._close = False + self._file.close() + del self._file + self._is_open = False + + @property + def closed(self): + return not self._is_open + + def flush(self): + raise io.UnsupportedOperation('flush') + + def _read(self, size, read_method, readinto_arg=None): + if size is None or size < 0: + size = -1 + self._file.seek(self._pos) + if not readinto_arg: + result = read_method(size) else: - self._pos = pos + if size < len(readinto_arg): + del readinto_arg[size:] + result = read_method(readinto_arg) + if result < len(readinto_arg): + del readinto_arg[result:] + self._pos = self._file.tell() + return result - def read(self, size=None): - """Read bytes.""" + def readable(self): + self._closed_check() + return True + + def read(self, size=-1): + self._closed_check() + if size is None or size < 0: + return self.readall() return self._read(size, self._file.read) - def read1(self, size=None): - """Read bytes.""" - return self._read(size, self._file.read1) + def readinto(self, by_arr): + self._closed_check() + return self._read(len(by_arr), self._file.readinto, by_arr) - def readline(self, size=None): - """Read a line.""" + def readall(self): + self._closed_check() + self._file.seek(self._pos) + if hasattr(self._file, 'readall'): + result = self._file.readall() + else: + dl = [] + while 1: + i = self._file.read(8192) + if len(i) == 0: + break + dl.append(i) + result = b''.join(dl) + self._pos = self._file.tell() + return result + + def readline(self, size=-1): + self._closed_check() return self._read(size, self._file.readline) - def readlines(self, sizehint=None): - """Read multiple lines.""" + def readlines(self, sizehint=-1): result = [] for line in self: result.append(line) - if sizehint is not None: + if sizehint >= 0: sizehint -= len(line) if sizehint <= 0: break return result + def seekable(self): + self._closed_check() + return True + + def seek(self, offset, whence=0): + self._closed_check() + if whence == 1: + self._file.seek(self._pos) + self._pos = self._file.seek(offset, whence) + return self._pos + + def tell(self): + self._closed_check() + return self._pos + + def writable(self): + self._closed_check() + return False + + def writelines(self, lines): + raise io.UnsupportedOperation('writelines') + + def write(self, b): + raise io.UnsupportedOperation('write') + def __iter__(self): - """Iterate over lines.""" while True: line = self.readline() if not line: raise StopIteration yield line - def tell(self): - """Return the position.""" - return self._pos - - def seek(self, offset, whence=0): - """Change position.""" - if whence == 1: - self._file.seek(self._pos) - self._file.seek(offset, whence) - self._pos = self._file.tell() - - def close(self): - """Close the file.""" - if hasattr(self._file, 'close'): - self._file.close() - del self._file - - def _read(self, size, read_method): - """Read size bytes using read_method.""" - if size is None: - size = -1 - self._file.seek(self._pos) - result = read_method(size) - self._pos = self._file.tell() - return result - def __enter__(self): - """Context manager protocol support.""" return self - def __exit__(self, *exc): self.close() - def readable(self): - return self._file.readable() - - def writable(self): - return self._file.writable() - - def seekable(self): - return self._file.seekable() - - def flush(self): - return self._file.flush() - - @property - def closed(self): - return self._file.closed - class _PartialFile(_ProxyFile): """A read-only wrapper of part of a file.""" @@ -1962,6 +2002,7 @@ def __init__(self, f, start=None, stop=None): """Initialize a _PartialFile.""" _ProxyFile.__init__(self, f, start) + super()._set_noclose() self._start = start self._stop = stop @@ -1971,6 +2012,7 @@ def seek(self, offset, whence=0): """Change position, possibly with respect to start or stop.""" + self._closed_check() if whence == 0: self._pos = self._start whence = 1 @@ -1979,20 +2021,21 @@ whence = 1 _ProxyFile.seek(self, offset, whence) - def _read(self, size, read_method): + def _read(self, size, read_method, readinto_arg=None): """Read size bytes using read_method, honoring start and stop.""" remaining = self._stop - self._pos if remaining <= 0: return b'' if size is None or size < 0 or size > remaining: size = remaining - return _ProxyFile._read(self, size, read_method) + return _ProxyFile._read(self, size, read_method, readinto_arg) - def close(self): - # do *not* close the underlying file object for partial files, - # since it's global to the mailbox object - del self._file - + def readall(self): + self._closed_check() + remaining = self._stop - self._pos + if remaining <= 0: + return b'' + return _ProxyFile._read(self, remaining, self._file.read) def _lock_file(f, dotlock=True): """Lock file f using lockf and dot locking.""" diff --git a/Lib/test/test_mailbox.py b/Lib/test/test_mailbox.py --- a/Lib/test/test_mailbox.py +++ b/Lib/test/test_mailbox.py @@ -290,12 +290,17 @@ key1 = self._box.add(_sample_message) with self._box.get_file(key0) as file: data0 = file.read() - with self._box.get_file(key1) as file: - data1 = file.read() + file1 = self._box.get_file(key1) + data1 = file1.read() self.assertEqual(data0.decode('ascii').replace(os.linesep, '\n'), self._template % 0) self.assertEqual(data1.decode('ascii').replace(os.linesep, '\n'), _sample_message) + file1.close() + try: + file1.close() + except: + self.fail('.close() doesn\'t handle multiple invocations') def test_iterkeys(self): # Get keys using iterkeys() @@ -1774,6 +1779,55 @@ proxy.seek(2) self.assertEqual(proxy.read(1000), b'r') + def _test_readinto(self, proxy): + # Fill in bytearray + proxy.seek(0) + ba = bytearray(3) + self.assertEqual(proxy.readinto(ba), 3) + self.assertEqual(ba, b'bar') + + proxy.seek(1) + ba = bytearray(2) + self.assertEqual(proxy.readinto(ba), 2) + self.assertEqual(ba, b'ar') + + proxy.seek(0) + ba = bytearray(2) + self.assertEqual(proxy.readinto(ba), 2) + self.assertEqual(ba, b'ba') + + proxy.seek(0) + ba = bytearray(2) + self.assertEqual(proxy.readinto(ba), 2) + self.assertEqual(ba, b'ba') + + proxy.seek(1) + ba = bytearray(1000) + self.assertEqual(proxy.readinto(ba), 2) + self.assertEqual(ba, b'ar') + + proxy.seek(2) + ba = bytearray(1000) + self.assertEqual(proxy.readinto(ba), 1) + self.assertEqual(ba, b'r') + + def _test_readall(self, proxy): + # Read all the data + ls = os.linesep.encode() + lsl = len(ls) + + proxy.seek(0) + x = b'fiesta' + ls + b'mexicana' + ls + self.assertEqual(proxy.readall(), x) + + proxy.seek(6 + lsl) + x = b'mexicana' + ls + self.assertEqual(proxy.readall(), x) + + proxy.seek(6+3 + lsl) + x = b'icana' + ls + self.assertEqual(proxy.readall(), x) + def _test_readline(self, proxy): # Read by line linesep = os.linesep.encode() @@ -1833,10 +1887,36 @@ self.assertFalse(proxy.read()) def _test_close(self, proxy): - # Close a file + self.assertFalse(proxy.closed) + self.assertRaises(io.UnsupportedOperation, proxy.flush) + self.assertTrue(proxy.readable()) + self.assertTrue(proxy.seekable()) + self.assertRaises(io.UnsupportedOperation, proxy.truncate, 0) + self.assertFalse(proxy.writable()) + self.assertRaises(io.UnsupportedOperation, proxy.writelines, ['AU']) + self.assertRaises(io.UnsupportedOperation, proxy.write, 'AU') + proxy.close() - self.assertRaises(AttributeError, lambda: proxy.close()) + self.assertTrue(proxy.closed) + try: + proxy.close() + except: + self.fail('Proxy.close() failure') + self.assertRaises(io.UnsupportedOperation, proxy.flush) + self.assertRaises(ValueError, proxy.readable) + self.assertRaises(ValueError, proxy.read) + self.assertRaises(ValueError, proxy.readinto, bytearray()) + self.assertRaises(ValueError, proxy.readall) + self.assertRaises(ValueError, proxy.readline) + self.assertRaises(ValueError, proxy.readlines) + self.assertRaises(ValueError, proxy.seekable) + self.assertRaises(ValueError, proxy.seek, 0) + self.assertRaises(ValueError, proxy.tell) + self.assertRaises(io.UnsupportedOperation, proxy.truncate, 0) + self.assertRaises(ValueError, proxy.writable) + self.assertRaises(io.UnsupportedOperation, proxy.writelines, ['AU']) + self.assertRaises(io.UnsupportedOperation, proxy.write, 'AU') class TestProxyFile(TestProxyFileBase): @@ -1863,6 +1943,15 @@ self._file.write(b'bar') self._test_read(mailbox._ProxyFile(self._file)) + def test_readinto(self): + self._file.write(b'bar') + self._test_readinto(mailbox._ProxyFile(self._file)) + + def test_readall(self): + ls = os.linesep.encode() + self._file.write(b'fiesta' + ls + b'mexicana' + ls) + self._test_readall(mailbox._ProxyFile(self._file)) + def test_readline(self): self._file.write(bytes('foo%sbar%sfred%sbob' % (os.linesep, os.linesep, os.linesep), 'ascii')) @@ -1886,6 +1975,13 @@ self._file.write(bytes('foo%sbar%s' % (os.linesep, os.linesep), 'ascii')) self._test_close(mailbox._ProxyFile(self._file)) + def test_ctxman(self): + self._file.write(b'foo') + proxy = mailbox._ProxyFile(self._file) + with proxy as p: + pass + self.assertTrue(proxy.closed) + class TestPartialFile(TestProxyFileBase): @@ -1909,6 +2005,16 @@ self._file.write(bytes('***bar***', 'ascii')) self._test_read(mailbox._PartialFile(self._file, 3, 6)) + def test_readinto(self): + self._file.write(b'***bar***') + self._test_readinto(mailbox._PartialFile(self._file, 3, 6)) + + def test_readall(self): + ls = os.linesep.encode() + lsl = len(ls) + self._file.write(b'***fiesta' + ls + b'mexicana' + ls + b'***') + self._test_readall(mailbox._PartialFile(self._file, 3, 3+14+2*lsl)) + def test_readline(self): self._file.write(bytes('!!!!!foo%sbar%sfred%sbob!!!!!' % (os.linesep, os.linesep, os.linesep), 'ascii')) @@ -1937,6 +2043,14 @@ self._test_close(mailbox._PartialFile(self._file, 1, 6 + 3 * len(os.linesep))) + def test_ctxman(self): + self._file.write(bytes('foo' + os.linesep + 'bar', 'ascii')) + pos = self._file.tell() + proxy = mailbox._PartialFile(self._file, 2, 5) + with proxy as p: + pass + self.assertTrue(proxy.closed) + ## Start: tests from the original module (for backward compatibility).