ztensor 0.1.2__py3-none-win32.whl → 0.1.4__py3-none-win32.whl

ztensor/__init__.py

@@ -92,12 +92,23 @@ class TensorMetadata:
     def __init__(self, meta_ptr):
         self._ptr = ffi.gc(meta_ptr, lib.ztensor_metadata_free)
         _check_ptr(self._ptr, "TensorMetadata constructor")
+        # Cache for properties to avoid repeated FFI calls
         self._name = None
         self._dtype_str = None
         self._shape = None
+        self._offset = None
+        self._size = None
+        self._layout = None
+        self._encoding = None
+        self._endianness = "not_checked"
+        self._checksum = "not_checked"
+
+    def __repr__(self):
+        return f"<TensorMetadata name='{self.name}' shape={self.shape} dtype='{self.dtype_str}'>"
 
     @property
     def name(self):
+        """The name of the tensor."""
         if self._name is None:
             name_ptr = lib.ztensor_metadata_get_name(self._ptr)
             _check_ptr(name_ptr, "get_name")
@@ -107,6 +118,7 @@ class TensorMetadata:
 
     @property
     def dtype_str(self):
+        """The zTensor dtype string (e.g., 'float32')."""
         if self._dtype_str is None:
             dtype_ptr = lib.ztensor_metadata_get_dtype_str(self._ptr)
             _check_ptr(dtype_ptr, "get_dtype_str")
@@ -116,7 +128,7 @@ class TensorMetadata:
 
     @property
     def dtype(self):
-        """Returns the numpy dtype for this tensor."""
+        """The numpy dtype for this tensor."""
         dtype_str = self.dtype_str
         dt = DTYPE_ZT_TO_NP.get(dtype_str)
         if dt is None:
@@ -130,6 +142,7 @@ class TensorMetadata:
 
     @property
     def shape(self):
+        """The shape of the tensor as a tuple."""
         if self._shape is None:
             shape_len = lib.ztensor_metadata_get_shape_len(self._ptr)
             if shape_len > 0:
@@ -141,6 +154,64 @@ class TensorMetadata:
                 self._shape = tuple()
         return self._shape
 
+    @property
+    def offset(self):
+        """The on-disk offset of the tensor data in bytes."""
+        if self._offset is None:
+            self._offset = lib.ztensor_metadata_get_offset(self._ptr)
+        return self._offset
+
+    @property
+    def size(self):
+        """The on-disk size of the tensor data in bytes (can be compressed size)."""
+        if self._size is None:
+            self._size = lib.ztensor_metadata_get_size(self._ptr)
+        return self._size
+
+    @property
+    def layout(self):
+        """The tensor layout as a string (e.g., 'dense')."""
+        if self._layout is None:
+            layout_ptr = lib.ztensor_metadata_get_layout_str(self._ptr)
+            _check_ptr(layout_ptr, "get_layout_str")
+            self._layout = ffi.string(layout_ptr).decode('utf-8')
+            lib.ztensor_free_string(layout_ptr)
+        return self._layout
+
+    @property
+    def encoding(self):
+        """The tensor encoding as a string (e.g., 'raw', 'zstd')."""
+        if self._encoding is None:
+            encoding_ptr = lib.ztensor_metadata_get_encoding_str(self._ptr)
+            _check_ptr(encoding_ptr, "get_encoding_str")
+            self._encoding = ffi.string(encoding_ptr).decode('utf-8')
+            lib.ztensor_free_string(encoding_ptr)
+        return self._encoding
+
+    @property
+    def endianness(self):
+        """The data endianness ('little', 'big') if applicable, else None."""
+        if self._endianness == "not_checked":
+            endian_ptr = lib.ztensor_metadata_get_data_endianness_str(self._ptr)
+            if endian_ptr == ffi.NULL:
+                self._endianness = None
+            else:
+                self._endianness = ffi.string(endian_ptr).decode('utf-8')
+                lib.ztensor_free_string(endian_ptr)
+        return self._endianness
+
+    @property
+    def checksum(self):
+        """The checksum string if present, else None."""
+        if self._checksum == "not_checked":
+            checksum_ptr = lib.ztensor_metadata_get_checksum_str(self._ptr)
+            if checksum_ptr == ffi.NULL:
+                self._checksum = None
+            else:
+                self._checksum = ffi.string(checksum_ptr).decode('utf-8')
+                lib.ztensor_free_string(checksum_ptr)
+        return self._checksum
+
 
 class Reader:
     """A Pythonic context manager for reading zTensor files."""
@@ -157,6 +228,39 @@ class Reader:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self._ptr = None
 
+    def __len__(self):
+        """Returns the number of tensors in the file."""
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
+        return lib.ztensor_reader_get_metadata_count(self._ptr)
+
+    def __iter__(self):
+        """Iterates over the metadata of all tensors in the file."""
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
+        for i in range(len(self)):
+            yield self[i]
+
+    def __getitem__(self, index: int) -> TensorMetadata:
+        """Retrieves metadata for a tensor by its index."""
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
+        if index >= len(self):
+            raise IndexError("Tensor index out of range")
+        meta_ptr = lib.ztensor_reader_get_metadata_by_index(self._ptr, index)
+        _check_ptr(meta_ptr, f"get_metadata_by_index: {index}")
+        return TensorMetadata(meta_ptr)
+
+    def list_tensors(self) -> list[TensorMetadata]:
+        """Returns a list of all TensorMetadata objects in the file."""
+        return list(self)
+
+    def get_tensor_names(self) -> list[str]:
+        """Returns a list of all tensor names in the file."""
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
+        c_array_ptr = lib.ztensor_reader_get_all_tensor_names(self._ptr)
+        _check_ptr(c_array_ptr, "get_all_tensor_names")
+        c_array_ptr = ffi.gc(c_array_ptr, lib.ztensor_free_string_array)
+
+        return [ffi.string(c_array_ptr.strings[i]).decode('utf-8') for i in range(c_array_ptr.len)]
+
     def get_metadata(self, name: str) -> TensorMetadata:
         """Retrieves metadata for a tensor by its name."""
         if self._ptr is None: raise ZTensorError("Reader is closed.")
@@ -177,6 +281,7 @@ class Reader:
         Returns:
             np.ndarray or torch.Tensor: The tensor data.
         """
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
         if to not in ['numpy', 'torch']:
             raise ValueError(f"Unsupported format: '{to}'. Choose 'numpy' or 'torch'.")
 
@@ -238,6 +343,7 @@ class Writer:
             if exc_type is None:
                 self.finalize()
             else:
+                # If an error occurred, just free the handle without finalizing
                 lib.ztensor_writer_free(self._ptr)
                 self._ptr = None
 
@@ -294,7 +400,7 @@ class Writer:
         """Finalizes the zTensor file, writing the metadata index."""
         if not self._ptr: raise ZTensorError("Writer is already closed or finalized.")
         status = lib.ztensor_writer_finalize(self._ptr)
-        self._ptr = None
+        self._ptr = None  # The writer is consumed in Rust
         self._finalized = True
         _check_status(status, "finalize")
 

ztensor/ztensor/ffi.py

@@ -3,8 +3,8 @@ import _cffi_backend
 
 ffi = _cffi_backend.FFI('ffi',
     _version = 0x2601,
-    _types = b'\x00\x00\x29\x0D\x00\x00\x3D\x03\x00\x00\x3C\x03\x00\x00\x00\x0F\x00\x00\x2C\x0D\x00\x00\x3D\x03\x00\x00\x3F\x03\x00\x00\x00\x0F\x00\x00\x01\x0D\x00\x00\x06\x11\x00\x00\x00\x0F\x00\x00\x14\x0D\x00\x00\x06\x11\x00\x00\x00\x0F\x00\x00\x35\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x06\x0D\x00\x00\x00\x0F\x00\x00\x40\x0D\x00\x00\x3E\x03\x00\x00\x00\x0F\x00\x00\x40\x0D\x00\x00\x14\x11\x00\x00\x06\x11\x00\x00\x41\x03\x00\x00\x1C\x01\x00\x00\x06\x11\x00\x00\x42\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x1A\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x1A\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x38\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3B\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3C\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x01\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x14\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3F\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x41\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x00\x09\x00\x00\x01\x09\x00\x00\x02\x09\x00\x00\x03\x09\x00\x00\x02\x01\x00\x00\x07\x01\x00\x00\x18\x01\x00\x00\x04\x01\x00\x00\x44\x03\x00\x00\x00\x01',
-    _globals = (b'\xFF\xFF\xFF\x1FALIGNMENT',64,b'\x00\x00\x34\x23ztensor_free_string',0,b'\x00\x00\x28\x23ztensor_free_tensor_view',0,b'\x00\x00\x37\x23ztensor_free_u64_array',0,b'\x00\x00\x11\x23ztensor_last_error_message',0,b'\x00\x00\x2B\x23ztensor_metadata_free',0,b'\x00\x00\x0E\x23ztensor_metadata_get_dtype_str',0,b'\x00\x00\x0E\x23ztensor_metadata_get_name',0,b'\x00\x00\x25\x23ztensor_metadata_get_shape_data',0,b'\x00\x00\x1F\x23ztensor_metadata_get_shape_len',0,b'\x00\x00\x2E\x23ztensor_reader_free',0,b'\x00\x00\x04\x23ztensor_reader_get_metadata_by_name',0,b'\x00\x00\x22\x23ztensor_reader_get_metadata_count',0,b'\x00\x00\x08\x23ztensor_reader_open',0,b'\x00\x00\x00\x23ztensor_reader_read_tensor_view',0,b'\x00\x00\x16\x23ztensor_writer_add_tensor',0,b'\x00\x00\x0B\x23ztensor_writer_create',0,b'\x00\x00\x13\x23ztensor_writer_finalize',0,b'\x00\x00\x31\x23ztensor_writer_free',0),
-    _struct_unions = ((b'\x00\x00\x00\x3B\x00\x00\x00\x02CTensorDataView',b'\x00\x00\x1C\x11data',b'\x00\x00\x1A\x11len',b'\x00\x00\x43\x11_owner'),(b'\x00\x00\x00\x3C\x00\x00\x00\x10TensorMetadata',),(b'\x00\x00\x00\x3D\x00\x00\x00\x10ZTensorReader_BufReader_File',),(b'\x00\x00\x00\x3E\x00\x00\x00\x10ZTensorWriter_BufWriter_File',)),
-    _typenames = (b'\x00\x00\x00\x3BCTensorDataView',b'\x00\x00\x00\x3CCTensorMetadata',b'\x00\x00\x00\x3DCZTensorReader',b'\x00\x00\x00\x3ECZTensorWriter',b'\x00\x00\x00\x3CTensorMetadata',b'\x00\x00\x00\x3DZTensorReader_BufReader_File',b'\x00\x00\x00\x3EZTensorWriter_BufWriter_File'),
+    _types = b'\x00\x00\x33\x0D\x00\x00\x4B\x03\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x4B\x03\x00\x00\x4A\x03\x00\x00\x00\x0F\x00\x00\x39\x0D\x00\x00\x01\x11\x00\x00\x4E\x03\x00\x00\x00\x0F\x00\x00\x39\x0D\x00\x00\x01\x11\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x04\x0D\x00\x00\x09\x11\x00\x00\x00\x0F\x00\x00\x1B\x0D\x00\x00\x09\x11\x00\x00\x00\x0F\x00\x00\x42\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x09\x0D\x00\x00\x00\x0F\x00\x00\x4F\x0D\x00\x00\x4C\x03\x00\x00\x00\x0F\x00\x00\x4F\x0D\x00\x00\x1B\x11\x00\x00\x09\x11\x00\x00\x50\x03\x00\x00\x1C\x01\x00\x00\x09\x11\x00\x00\x51\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x0D\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x0D\x0D\x00\x00\x01\x11\x00\x00\x00\x0F\x00\x00\x45\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x50\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x48\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x49\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x4A\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x1B\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x4E\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x50\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x00\x09\x00\x00\x01\x09\x00\x00\x02\x09\x00\x00\x03\x09\x00\x00\x04\x09\x00\x00\x42\x03\x00\x00\x02\x01\x00\x00\x07\x01\x00\x00\x18\x01\x00\x00\x04\x01\x00\x00\x53\x03\x00\x00\x00\x01',
+    _globals = (b'\xFF\xFF\xFF\x1FALIGNMENT',64,b'\x00\x00\x41\x23ztensor_free_string',0,b'\x00\x00\x32\x23ztensor_free_string_array',0,b'\x00\x00\x35\x23ztensor_free_tensor_view',0,b'\x00\x00\x44\x23ztensor_free_u64_array',0,b'\x00\x00\x18\x23ztensor_last_error_message',0,b'\x00\x00\x38\x23ztensor_metadata_free',0,b'\x00\x00\x15\x23ztensor_metadata_get_checksum_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_data_endianness_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_dtype_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_encoding_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_layout_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_name',0,b'\x00\x00\x2F\x23ztensor_metadata_get_offset',0,b'\x00\x00\x2C\x23ztensor_metadata_get_shape_data',0,b'\x00\x00\x26\x23ztensor_metadata_get_shape_len',0,b'\x00\x00\x2F\x23ztensor_metadata_get_size',0,b'\x00\x00\x3B\x23ztensor_reader_free',0,b'\x00\x00\x00\x23ztensor_reader_get_all_tensor_names',0,b'\x00\x00\x0B\x23ztensor_reader_get_metadata_by_index',0,b'\x00\x00\x07\x23ztensor_reader_get_metadata_by_name',0,b'\x00\x00\x29\x23ztensor_reader_get_metadata_count',0,b'\x00\x00\x0F\x23ztensor_reader_open',0,b'\x00\x00\x03\x23ztensor_reader_read_tensor_view',0,b'\x00\x00\x1D\x23ztensor_writer_add_tensor',0,b'\x00\x00\x12\x23ztensor_writer_create',0,b'\x00\x00\x1A\x23ztensor_writer_finalize',0,b'\x00\x00\x3E\x23ztensor_writer_free',0),
+    _struct_unions = ((b'\x00\x00\x00\x48\x00\x00\x00\x02CStringArray',b'\x00\x00\x4D\x11strings',b'\x00\x00\x0D\x11len'),(b'\x00\x00\x00\x49\x00\x00\x00\x02CTensorDataView',b'\x00\x00\x23\x11data',b'\x00\x00\x0D\x11len',b'\x00\x00\x52\x11_owner'),(b'\x00\x00\x00\x4A\x00\x00\x00\x10TensorMetadata',),(b'\x00\x00\x00\x4B\x00\x00\x00\x10ZTensorReader_BufReader_File',),(b'\x00\x00\x00\x4C\x00\x00\x00\x10ZTensorWriter_BufWriter_File',)),
+    _typenames = (b'\x00\x00\x00\x48CStringArray',b'\x00\x00\x00\x49CTensorDataView',b'\x00\x00\x00\x4ACTensorMetadata',b'\x00\x00\x00\x4BCZTensorReader',b'\x00\x00\x00\x4CCZTensorWriter',b'\x00\x00\x00\x4ATensorMetadata',b'\x00\x00\x00\x4BZTensorReader_BufReader_File',b'\x00\x00\x00\x4CZTensorWriter_BufWriter_File'),
 )

{ztensor-0.1.2.dist-info → ztensor-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ztensor
-Version: 0.1.2
+Version: 0.1.4
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License

ztensor-0.1.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+ztensor-0.1.4.dist-info/METADATA,sha256=Xs4X80I_5Ox-nP5Gx8Ng5vmhEi9lvYgJB8MdLl2pvu8,4530
+ztensor-0.1.4.dist-info/WHEEL,sha256=CiLp0kI_D1SQdHfHt9BWsWFTRNd71UY0p9KVx6QjX7o,89
+ztensor-0.1.4.dist-info/licenses/LICENSE,sha256=qxF7VFxBvMlfiDRJ5oXQuQYaloq0Tcbk95Pn0DFlnss,1084
+ztensor/__init__.py,sha256=qs_Lg29u2jpA4fu01E0FV8rY4tVB-MdtwCAI6_vn6_A,16202
+ztensor/ztensor/__init__.py,sha256=DDVvoEhcXithkluOJ4Dd7H6wIqKcxT6mm6vvPgrQMz4,138
+ztensor/ztensor/ffi.py,sha256=ZBlZW--RXKUdoZZJJjtftY8H8hBmY_FvKcG6ab8jOTA,3611
+ztensor/ztensor/ztensor.dll,sha256=JigXJYuOi0aRBHRKEMCwPceozaHCVgWWEJyu1neRd5M,891904
+ztensor-0.1.4.dist-info/RECORD,,

{ztensor-0.1.2.dist-info → ztensor-0.1.4.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: maturin (1.9.0)
+Generator: maturin (1.9.1)
 Root-Is-Purelib: false
 Tag: py3-none-win32

ztensor-0.1.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-ztensor-0.1.2.dist-info/METADATA,sha256=4UODGfQ_vW4gGHCIunIBOFRA4aB_BrbymnQQNa87vi4,4530
-ztensor-0.1.2.dist-info/WHEEL,sha256=v0atuDHRmBdo6voLOtYJ8NlQfECX5fZNHOXLlNqbhus,89
-ztensor-0.1.2.dist-info/licenses/LICENSE,sha256=qxF7VFxBvMlfiDRJ5oXQuQYaloq0Tcbk95Pn0DFlnss,1084
-ztensor/__init__.py,sha256=MHfJeuyHa9r1dJ-8ezqRX0wpR5EBwhjkMzg4HlnLFL4,11600
-ztensor/ztensor/__init__.py,sha256=DDVvoEhcXithkluOJ4Dd7H6wIqKcxT6mm6vvPgrQMz4,138
-ztensor/ztensor/ffi.py,sha256=5HqR7Szwsn6HmKARaYQqF0nJKEUWZsWpYD4ZiOaoWnk,2756
-ztensor/ztensor/ztensor.dll,sha256=RfrdQFVVG9QGDXU9z4nGxviysq5BCfxrZXTrdQ7wJF4,871936
-ztensor-0.1.2.dist-info/RECORD,,