diff --git a/bunker/files/bunkerfile.py b/bunker/files/bunkerfile.py index 1bd211d..dfe3217 100644 --- a/bunker/files/bunkerfile.py +++ b/bunker/files/bunkerfile.py @@ -78,7 +78,8 @@ class BunkeredFile(io.RawIOBase): """ size = None if(file_.seekable()): - size = file_.seek(0, 2) + file_.seek(0, 2) + size = file_.tell() file_.seek(0, 0) return cls(file_, name, size=size, isvirtual=False, ismem=False) @@ -96,9 +97,10 @@ class BunkeredFile(io.RawIOBase): @classmethod def from_BytesIO(cls, bytes_, name): size = None - if(file_.seekable()): - size = file_.seek(0, 2) - file_.seek(0, 0) + if(bytes_.seekable()): + bytes_.seek(0, 2) + size = bytes_.tell() + bytes_.seek(0, 0) return cls(bytes_, name, size=size, isvirtual=True, ismem=True) @classmethod diff --git a/bunker/files/tarfile.py b/bunker/files/tarfile.py index 34e5776..084c6ea 100644 --- a/bunker/files/tarfile.py +++ b/bunker/files/tarfile.py @@ -73,7 +73,7 @@ class RewriteableTarFile(object): for member in handle.getmembers(): if(member.name == membername): continue - new_handle.add_file(member, handle.extractfile(member)) + new_handle.addfile(member, handle.extractfile(member)) new_handle.close() def close(self): diff --git a/test/test_files_bunkerfile.py b/test/test_files_bunkerfile.py index 1b1c80b..780dca2 100644 --- a/test/test_files_bunkerfile.py +++ b/test/test_files_bunkerfile.py @@ -1,5 +1,6 @@ import os import tarfile +from io import BytesIO from bunker.files.bunkerfile import BunkeredFile def test_load_from_tar(tmpdir): @@ -27,3 +28,12 @@ def test_load_from_tar(tmpdir): assert a.read() == b"abcdefg" assert b.read() == b"foobar" + +def test_from_BytesIO(): + b = BytesIO(b"foobar") + f = BunkeredFile.from_BytesIO(b, "test") + + assert len(f) == len(b"foobar") + assert f.read() == b"foobar" + + diff --git a/test/test_files_tarfile.py b/test/test_files_tarfile.py index 62ff48b..18a4d00 100644 --- a/test/test_files_tarfile.py +++ b/test/test_files_tarfile.py @@ -1,5 +1,6 @@ import os import tarfile +import pytest from bunker.files.tarfile import RewriteableTarFile from bunker.files.bunkerfile import BunkeredFile @@ -21,3 +22,75 @@ def test_create(tmpdir): f = tf.get_file("__bunker_main__") assert f.read() == b"foobar" + +def test_rewrite(tmpdir): + tmpdname = str(tmpdir) + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = BunkeredFile.empty("__bunker_main__") + tf.add_file(f) + tf.add_file(BunkeredFile.empty("test.tx")) + + f = tf.get_file("__bunker_main__") + f.write(b"foobar") + + tf.close() + + assert f.closed == True + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = tf.get_file("__bunker_main__") + f.write(b"foobaz") + tf.close() + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = tf.get_file("__bunker_main__") + assert f.read() == b"foobaz" + +def test_writeback(tmpdir): + tmpdname = str(tmpdir) + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = BunkeredFile.empty("__bunker_main__") + tf.add_file(f) + tf.add_file(BunkeredFile.empty("test.tx")) + + f = tf.get_file("__bunker_main__") + f.write(b"foobar") + + tf.close() + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = tf.get_file("__bunker_main__") + f.write(b"foobaz") + f.close() + tf.close() + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = tf.get_file("__bunker_main__") + assert f.read() == b"foobaz" + + + +def test_delete(tmpdir): + tmpdname = str(tmpdir) + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + f = BunkeredFile.empty("__bunker_main__") + tf.add_file(f) + tf.add_file(BunkeredFile.empty("test.tx")) + + f = tf.get_file("__bunker_main__") + f.write(b"foobar") + + tf.close() + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + tf.delete_file("__bunker_main__") + tf.close() + + tf = RewriteableTarFile.open(os.path.join(tmpdname, "test.bunker")) + + with pytest.raises(KeyError): + f = tf.get_file("__bunker_main__") +