Skip to content

Commit

Permalink
fix: zip stream repeats initial bytes
Browse files Browse the repository at this point in the history
  • Loading branch information
dnlbauer committed Jan 24, 2025
1 parent c034c66 commit 27457d8
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 4 deletions.
28 changes: 28 additions & 0 deletions rocrate/memory_buffer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
from io import RawIOBase


class MemoryBuffer(RawIOBase):
"""
A buffer class that supports reading and writing binary data.
The buffer automatically resets upon reading to make sure all data is read only once.
"""

def __init__(self):
self._buffer = b''

def write(self, data):
if self.closed:
raise ValueError('write to closed file')
self._buffer += data
return len(data)

def read(self, size=-1):
if self.closed:
raise ValueError('read from closed file')
if size < 0:
data = self._buffer
self._buffer = b''
else:
data = self._buffer[:size]
self._buffer = self._buffer[size:]
return data
7 changes: 3 additions & 4 deletions rocrate/rocrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@
import warnings

from collections import OrderedDict
from io import BytesIO
from pathlib import Path
from urllib.parse import urljoin

from .memory_buffer import MemoryBuffer
from .model import (
ComputationalWorkflow,
ComputerLanguage,
Expand Down Expand Up @@ -482,7 +482,7 @@ def write_zip(self, out_path):

def stream_zip(self):
""" Create a stream of bytes representing the RO-Crate as a ZIP file. """
buffer = BytesIO()
buffer = MemoryBuffer()
with zipfile.ZipFile(buffer, mode='w', compression=zipfile.ZIP_DEFLATED) as archive:
for writeable_entity in self.data_entities + self.default_entities:
current_file_path, current_out_file = None, None
Expand All @@ -493,11 +493,10 @@ def stream_zip(self):
current_file_path = path
current_out_file = archive.open(path, mode='w')
current_out_file.write(chunk)
buffer.seek(0)
yield buffer.read()
if current_out_file:
current_out_file.close()
buffer.seek(0)

yield buffer.read()
buffer.close()

Expand Down

0 comments on commit 27457d8

Please sign in to comment.