Skip to content

Commit

Permalink
Update safetensors_file.py to support context manager protocol
Browse files Browse the repository at this point in the history
Add __enter__ and __exit__ to allow `with SafeTensorsFile.open_file(tensorsfile, quiet=True) as s:` syntax
  • Loading branch information
duanemoody authored Dec 1, 2023
1 parent 94502e9 commit ae554ca
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions safetensors_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,14 @@ def __init__(self):

def __del__(self):
self.close_file()

def __enter__(self):
return self

def __exit__(self, exc_type, exc_value, traceback):
if self.f is not None:
self.f.close()


def close_file(self):
if self.f is not None:
Expand Down

0 comments on commit ae554ca

Please sign in to comment.