Adding _safe_open_handle. (#608)

Narsil · web-flow · commit 48933f3822e6 · 2025-08-01T17:25:35.000+02:00
* [WIP]. Adding safe_handle.

* Adding S3.

* File handle becomes private for merge.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -38,12 +38,11 @@ repos:
         name: "Python (black)"
         args: ["--line-length", "119", "--target-version", "py35"]
         types: ["python"]
-  - repo: https://github.com/pycqa/flake8
-    rev: 7.2.0
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.11.11
     hooks:
-      - id: flake8
-        args: ["--config", "bindings/python/setup.cfg"]
-  - repo: https://github.com/pre-commit/mirrors-isort
-    rev: v5.7.0 # Use the revision sha / tag you want to point at
-    hooks:
-      - id: isort
+      # Run the linter.
+      - id: ruff-check
+      # Run the formatter.
+      - id: ruff-format
diff --git a/bindings/python/py_src/safetensors/__init__.py b/bindings/python/py_src/safetensors/__init__.py
@@ -4,6 +4,7 @@
     __version__,
     deserialize,
     safe_open,
+    _safe_open_handle,
     serialize,
     serialize_file,
 )
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
@@ -1258,13 +1258,126 @@ pyo3::create_exception!(
     "Custom Python Exception for Safetensor errors."
 );
 
+#[pyclass]
+#[allow(non_camel_case_types)]
+struct _safe_open_handle {
+    inner: Option<Open>,
+}
+
+impl _safe_open_handle {
+    fn inner(&self) -> PyResult<&Open> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| SafetensorError::new_err("File is closed".to_string()))?;
+        Ok(inner)
+    }
+}
+
+#[pymethods]
+impl _safe_open_handle {
+    #[new]
+    #[pyo3(signature = (f, framework, device=Some(Device::Cpu)))]
+    fn new(f: PyObject, framework: Framework, device: Option<Device>) -> PyResult<Self> {
+        let filename = Python::with_gil(|py| -> PyResult<PathBuf> {
+            let _ = f.getattr(py, "fileno")?;
+            let filename = f.getattr(py, "name")?;
+            let filename: PathBuf = filename.extract(py)?;
+            Ok(filename)
+        })?;
+        let inner = Some(Open::new(filename, framework, device)?);
+        Ok(Self { inner })
+    }
+
+    /// Return the special non tensor information in the header
+    ///
+    /// Returns:
+    ///     (`Dict[str, str]`):
+    ///         The freeform metadata.
+    pub fn metadata(&self) -> PyResult<Option<HashMap<String, String>>> {
+        Ok(self.inner()?.metadata())
+    }
+
+    /// Returns the names of the tensors in the file.
+    ///
+    /// Returns:
+    ///     (`List[str]`):
+    ///         The name of the tensors contained in that file
+    pub fn keys(&self) -> PyResult<Vec<String>> {
+        self.inner()?.keys()
+    }
+
+    /// Returns the names of the tensors in the file, ordered by offset.
+    ///
+    /// Returns:
+    ///     (`List[str]`):
+    ///         The name of the tensors contained in that file
+    pub fn offset_keys(&self) -> PyResult<Vec<String>> {
+        self.inner()?.offset_keys()
+    }
+
+    /// Returns a full tensor
+    ///
+    /// Args:
+    ///     name (`str`):
+    ///         The name of the tensor you want
+    ///
+    /// Returns:
+    ///     (`Tensor`):
+    ///         The tensor in the framework you opened the file for.
+    ///
+    /// Example:
+    /// ```python
+    /// from safetensors import safe_open
+    ///
+    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
+    ///     tensor = f.get_tensor("embedding")
+    ///
+    /// ```
+    pub fn get_tensor(&self, name: &str) -> PyResult<PyObject> {
+        self.inner()?.get_tensor(name)
+    }
+
+    /// Returns a full slice view object
+    ///
+    /// Args:
+    ///     name (`str`):
+    ///         The name of the tensor you want
+    ///
+    /// Returns:
+    ///     (`PySafeSlice`):
+    ///         A dummy object you can slice into to get a real tensor
+    /// Example:
+    /// ```python
+    /// from safetensors import safe_open
+    ///
+    /// with safe_open("model.safetensors", framework="pt", device=0) as f:
+    ///     tensor_part = f.get_slice("embedding")[:, ::8]
+    ///
+    /// ```
+    pub fn get_slice(&self, name: &str) -> PyResult<PySafeSlice> {
+        self.inner()?.get_slice(name)
+    }
+
+    /// Start the context manager
+    pub fn __enter__(slf: Py<Self>) -> Py<Self> {
+        slf
+    }
+
+    /// Exits the context manager
+    pub fn __exit__(&mut self, _exc_type: PyObject, _exc_value: PyObject, _traceback: PyObject) {
+        self.inner = None;
+    }
+}
+
 /// A Python module implemented in Rust.
 #[pymodule(gil_used = false)]
 fn _safetensors_rust(m: &PyBound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(serialize, m)?)?;
     m.add_function(wrap_pyfunction!(serialize_file, m)?)?;
     m.add_function(wrap_pyfunction!(deserialize, m)?)?;
     m.add_class::<safe_open>()?;
+    m.add_class::<_safe_open_handle>()?;
     m.add("SafetensorError", m.py().get_type::<SafetensorError>())?;
     m.add("__version__", env!("CARGO_PKG_VERSION"))?;
     Ok(())
diff --git a/bindings/python/tests/test_handle.py b/bindings/python/tests/test_handle.py
@@ -0,0 +1,65 @@
+import unittest
+
+import numpy as np
+
+from safetensors import _safe_open_handle
+from safetensors.numpy import save_file, save
+
+
+class ReadmeTestCase(unittest.TestCase):
+    def assertTensorEqual(self, tensors1, tensors2, equality_fn):
+        self.assertEqual(tensors1.keys(), tensors2.keys(), "tensor keys don't match")
+
+        for k, v1 in tensors1.items():
+            v2 = tensors2[k]
+
+            self.assertTrue(equality_fn(v1, v2), f"{k} tensors are different")
+
+    def test_numpy_example(self):
+        tensors = {"a": np.zeros((2, 2)), "b": np.zeros((2, 3), dtype=np.uint8)}
+
+        save_file(tensors, "./out_np.safetensors")
+
+        # Now loading
+        loaded = {}
+        with open("./out_np.safetensors", "r") as f:
+            with safe_open_handle(f, framework="np", device="cpu") as g:
+                for key in g.keys():
+                    loaded[key] = g.get_tensor(key)
+        self.assertTensorEqual(tensors, loaded, np.allclose)
+
+    def test_fsspec(self):
+        import fsspec
+
+        tensors = {"a": np.zeros((2, 2)), "b": np.zeros((2, 3), dtype=np.uint8)}
+
+        fs = fsspec.filesystem("file")
+        byts = save(tensors)
+        with fs.open("fs.safetensors", "wb") as f:
+            f.write(byts)
+        # Now loading
+        loaded = {}
+        with fs.open("fs.safetensors", "rb") as f:
+            with safe_open_handle(f, framework="np", device="cpu") as g:
+                for key in g.keys():
+                    loaded[key] = g.get_tensor(key)
+        self.assertTensorEqual(tensors, loaded, np.allclose)
+
+    @unittest.skip("Will not work without s3 access")
+    def test_fsspec_s3(self):
+        import s3fs
+
+        tensors = {"a": np.zeros((2, 2)), "b": np.zeros((2, 3), dtype=np.uint8)}
+
+        s3 = s3fs.S3FileSystem(anon=True)
+        byts = save(tensors)
+        print(s3.ls("my-bucket"))
+        with s3.open("out/fs.safetensors", "wb") as f:
+            f.write(byts)
+        # Now loading
+        loaded = {}
+        with s3.open("out/fs.safetensors", "rb") as f:
+            with safe_open_handle(f, framework="np", device="cpu") as g:
+                for key in g.keys():
+                    loaded[key] = g.get_tensor(key)
+        self.assertTensorEqual(tensors, loaded, np.allclose)

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`	`__version__,`
`5`	`5`	`deserialize,`
`6`	`6`	`safe_open,`
	`7`	`+ _safe_open_handle,`
`7`	`8`	`serialize,`
`8`	`9`	`serialize_file,`
`9`	`10`	`)`