Fixing the fact that the context manager wasn't properly cleaning up after itself. (#166)

* Adding failing test for windows.

* This should fail.

* Fixing Windows issue by actually dropping the rust resources.

* Making the rust module private.

Author: Narsil

bindings/python/py_src/safetensors/__init__.py

@@ -1,4 +1,20 @@
 __version__ = "0.2.9"
 
 # Re-export this
-from .safetensors_rust import safe_open  # noqa: F401
+from ._safetensors_rust import safe_open as rust_open, serialize, serialize_file, deserialize, SafetensorError
+
+
+class safe_open:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+    def __getattr__(self, __name: str):
+        return getattr(self.f, __name)
+
+    def __enter__(self):
+        self.f = rust_open(*self.args, **self.kwargs)
+        return self
+
+    def __exit__(self, type, value, traceback):
+        del self.f

bindings/python/py_src/safetensors/numpy.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from .safetensors_rust import deserialize, safe_open, serialize, serialize_file
+from safetensors import deserialize, safe_open, serialize, serialize_file
 
 
 def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None) -> bytes:

bindings/python/py_src/safetensors/torch.py

@@ -5,7 +5,7 @@
 
 import torch
 
-from .safetensors_rust import deserialize, safe_open, serialize, serialize_file
+from safetensors import deserialize, safe_open, serialize, serialize_file
 
 
 def save(tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes:

bindings/python/setup.py

@@ -59,7 +59,7 @@ def deps_list(*pkgs):
     author_email="",
     url="https://github.com/huggingface/safetensors",
     license="Apache License 2.0",
-    rust_extensions=[RustExtension("safetensors.safetensors_rust", binding=Binding.PyO3, debug=False)],
+    rust_extensions=[RustExtension("safetensors._safetensors_rust", binding=Binding.PyO3, debug=False)],
     extras_require=extras,
     classifiers=[
         "Development Status :: 5 - Production/Stable",

bindings/python/src/lib.rs

@@ -1021,7 +1021,7 @@ pyo3::create_exception!(
 
 /// A Python module implemented in Rust.
 #[pymodule]
-fn safetensors_rust(py: Python, m: &PyModule) -> PyResult<()> {
+fn _safetensors_rust(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(serialize, m)?)?;
     m.add_function(wrap_pyfunction!(serialize_file, m)?)?;
     m.add_function(wrap_pyfunction!(deserialize, m)?)?;

bindings/python/tests/test_pt_comparison.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from safetensors.safetensors_rust import safe_open
+from safetensors import safe_open
 from safetensors.torch import load, load_file, save, save_file
 
 

bindings/python/tests/test_simple.py

@@ -6,7 +6,7 @@
 import torch
 
 from safetensors.numpy import load, load_file, save, save_file
-from safetensors.safetensors_rust import SafetensorError, serialize
+from safetensors import safe_open, SafetensorError, serialize
 from safetensors.torch import load_file as load_file_pt
 from safetensors.torch import save_file as save_file_pt
 
@@ -59,6 +59,20 @@ def test_accept_path(self):
         os.remove(Path("./out.safetensors"))
 
 
+class WindowsTestCase(unittest.TestCase):
+    def test_get_correctly_dropped(self):
+        tensors = {
+            "a": torch.zeros((2, 2)),
+            "b": torch.zeros((2, 3), dtype=torch.uint8),
+        }
+        save_file_pt(tensors, "./out.safetensors")
+        with safe_open("./out.safetensors", framework="pt") as f:
+            pass
+
+        with open("./out.safetensors", "w") as g:
+            g.write("something")
+
+
 class ReadmeTestCase(unittest.TestCase):
     def assertTensorEqual(self, tensors1, tensors2, equality_fn):
         self.assertEqual(tensors1.keys(), tensors2.keys(), "tensor keys don't match")