Allow full Vec3Int as input for get_buffered_slice_writer (#438)

Author: fm3

webknossos/webknossos/dataset/_utils/buffered_slice_writer.py

@@ -8,7 +8,7 @@
 import numpy as np
 import psutil
 
-from webknossos.geometry import Vec3Int
+from webknossos.geometry import Vec3Int, Vec3IntLike
 
 if TYPE_CHECKING:
     from webknossos.dataset import View
@@ -32,7 +32,7 @@ class BufferedSliceWriter(object):
     def __init__(
         self,
         view: "View",
-        offset: Vec3Int,
+        offset: Vec3IntLike,
         # buffer_size specifies, how many slices should be aggregated until they are flushed.
         buffer_size: int = 32,
         dimension: int = 2,  # z
@@ -48,7 +48,7 @@ def __init__(
         self.view = view
         self.buffer_size = buffer_size
         self.dtype = self.view.get_dtype()
-        self.offset = offset
+        self.offset = Vec3Int(offset)
         self.dimension = dimension
 
         assert 0 <= dimension <= 2

webknossos/webknossos/dataset/view.py

@@ -273,7 +273,10 @@ def get_view(
         )
 
     def get_buffered_slice_writer(
-        self, offset: Vec3Int = None, buffer_size: int = 32, dimension: int = 2  # z
+        self,
+        offset: Vec3IntLike = Vec3Int(0, 0, 0),
+        buffer_size: int = 32,
+        dimension: int = 2,  # z
     ) -> "BufferedSliceWriter":
         """
         The BufferedSliceWriter buffers multiple slices before they are written to disk.
@@ -299,15 +302,15 @@ def get_buffered_slice_writer(
 
         return BufferedSliceWriter(
             view=self,
-            offset=offset if offset is not None else Vec3Int(0, 0, 0),
+            offset=Vec3Int(offset),
             buffer_size=buffer_size,
             dimension=dimension,
         )
 
     def get_buffered_slice_reader(
         self,
-        offset: Vec3Int = None,
-        size: Vec3Int = None,
+        offset: Vec3IntLike = Vec3Int(0, 0, 0),
+        size: Optional[Vec3IntLike] = None,
         buffer_size: int = 32,
         dimension: int = 2,  # z
     ) -> "BufferedSliceReader":
@@ -333,8 +336,8 @@ def get_buffered_slice_reader(
 
         return BufferedSliceReader(
             view=self,
-            offset=offset if offset is not None else Vec3Int(0, 0, 0),
-            size=size if size is not None else self.size,
+            offset=Vec3Int(offset),
+            size=Vec3Int(size) if size is not None else self.size,
             buffer_size=buffer_size,
             dimension=dimension,
         )