feat: extend OutputStream, AvroOutputFile and ManifestWriter with __len__ method

Author: felixscherz

pyiceberg/avro/file.py

@@ -228,6 +228,7 @@ class AvroOutputFile(Generic[D]):
     encoder: BinaryEncoder
     sync_bytes: bytes
     writer: Writer
+    closed: bool
 
     def __init__(
         self,
@@ -247,6 +248,7 @@ def __init__(
             else resolve_writer(record_schema=record_schema, file_schema=self.file_schema)
         )
         self.metadata = metadata
+        self.closed = False
 
     def __enter__(self) -> AvroOutputFile[D]:
         """
@@ -267,6 +269,7 @@ def __exit__(
     ) -> None:
         """Perform cleanup when exiting the scope of a 'with' statement."""
         self.output_stream.close()
+        self.closed = True
 
     def _write_header(self) -> None:
         json_schema = json.dumps(AvroSchemaConversion().iceberg_to_avro(self.file_schema, schema_name=self.schema_name))
@@ -285,3 +288,9 @@ def write_block(self, objects: List[D]) -> None:
         self.encoder.write_int(len(block_content))
         self.encoder.write(block_content)
         self.encoder.write(self.sync_bytes)
+
+    def __len__(self) -> int:
+        """Returns the total number number of bytes written."""
+        if self.closed:
+            return len(self.output_file)
+        return len(self.output_stream)

pyiceberg/io/__init__.py

@@ -124,6 +124,10 @@ def __exit__(
     ) -> None:
         """Perform cleanup when exiting the scope of a 'with' statement."""
 
+    @abstractmethod
+    def __len__(self) -> int:
+        """Returns the total number number of bytes written to the stream."""
+
 
 class InputFile(ABC):
     """A base class for InputFile implementations.

pyiceberg/manifest.py

@@ -777,6 +777,11 @@ def add_entry(self, entry: ManifestEntry) -> ManifestWriter:
         self._writer.write_block([self.prepare_entry(entry)])
         return self
 
+    def __len__(self) -> int:
+        """Returns the total number number of bytes written."""
+        return len(self._writer)
+
+
 
 class RollingManifestWriter:
     closed: bool
@@ -829,7 +834,7 @@ def _should_roll_to_new_file(self) -> bool:
             return False
         return (
             self._current_file_rows >= self._target_number_of_rows
-            or len(self._current_writer._output_file) >= self._target_file_size_in_bytes
+            or len(self._current_writer) >= self._target_file_size_in_bytes
         )
 
     def _close_current_writer(self) -> None:

tests/utils/test_manifest.py

@@ -484,8 +484,8 @@ def test_write_manifest(
     [
         (19514, 388873, 1),  # should not roll over
         (19513, 388873, 2),  # should roll over due to target_rows
-        (19514, 388872, 2),  # should roll over due target_bytes
-        (19513, 388872, 2),  # should roll over due to target_rows and target_bytes
+        (4000, 388872, 2),  # should roll over due target_bytes
+        (4000, 388872, 2),  # should roll over due to target_rows and target_bytes
     ],
 )
 def test_rolling_manifest_writer(