diff --git a/mar/mar.py b/mar/mar.py index 1d1eaee..6d2a988 100644 --- a/mar/mar.py +++ b/mar/mar.py @@ -27,7 +27,7 @@ log = logging.getLogger(__name__) def read_file(fp, blocksize=8192): """Yields blocks of data from file object fp""" - for block in iter(partial(fp.read, blocksize), ''): + for block in iter(partial(fp.read, blocksize), b''): yield block @@ -188,7 +188,7 @@ class MarInfo: def to_bytes(self): return struct.pack(self._member_fmt, self._offset, self.size, self.flags) + \ - self.name + "\x00" + self.name.encode("ascii") + b"\x00" class MarFile: @@ -236,7 +236,7 @@ class MarFile: self.index_offset += 4 + 8 # Write the magic and placeholder for the index - self.fileobj.write("MAR1" + packint(self.index_offset)) + self.fileobj.write(b"MAR1" + packint(self.index_offset)) # Write placeholder for file size self.fileobj.write(struct.pack(">Q", 0)) diff --git a/tests/test_mar.py b/tests/test_mar.py index 3f5bc7d..dd53489 100644 --- a/tests/test_mar.py +++ b/tests/test_mar.py @@ -4,11 +4,18 @@ import os import tempfile import hashlib -from mar.mar import MarFile, BZ2MarFile +from mar.mar import MarFile, BZ2MarFile, read_file TEST_MAR = os.path.join(os.path.dirname(__file__), 'test.mar') +def test_read_file(): + data = [] + for block in read_file(open(__file__, 'rb')): + data.append(block) + assert b''.join(data) == open(__file__, 'rb').read() + + def sha1sum(b): """Returns the sha1sum of a byte string""" h = hashlib.new('sha1') @@ -26,7 +33,7 @@ def test_list(): assert repr(m.members[1]) == "", m.members[1] -class TestMar(TestCase): +class TestReadingMar(TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() self.marfile = MarFile(TEST_MAR) @@ -48,7 +55,7 @@ class TestMar(TestCase): self.assertEquals("6a7890e740f1e18a425b51fefbde2f6b86f91a12", h) -class TestBZ2Mar(TestCase): +class TestReadingBZ2Mar(TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() self.marfile = BZ2MarFile(TEST_MAR) @@ -72,6 +79,19 @@ class TestBZ2Mar(TestCase): self.assertEquals("5177f5938923e94820d8565a1a0f25d19b4821d1", h) +class TestWritingMar(TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.tmpdir) + + def test_add(self): + marfile = os.path.join(self.tmpdir, 'test.mar') + with MarFile(marfile, 'w') as m: + m.add(__file__) + + class TestExceptions(TestCase): def test_badmar(self): self.assertRaises(ValueError, MarFile, __file__)