Return self from database __enter__ method

microfluidica · Apr 4, 2024 · 29d0eca · 29d0eca
1 parent 59902ec
commit 29d0eca
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/electrolytes/__init__.py b/electrolytes/__init__.py
@@ -188,11 +188,13 @@ def _save_user_constituents(self) -> None:
             self._user_constituents_file.write_bytes(data)
         self._user_constituents_dirty = False
 
-    def __enter__(self) -> None:
+    def __enter__(self) -> "_Database":
         if not self._user_constituents_lock.is_locked:
             self._invalidate_user_constituents()
         self._user_constituents_lock.acquire()
 
+        return self
+
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         try:
             if (