enhance(client): use destructor instead of atexit

star-whale · Aug 31, 2022 · d9e2448 · d9e2448
1 parent c6d8c07
commit d9e2448
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/client/starwhale/api/_impl/data_store.py b/client/starwhale/api/_impl/data_store.py
@@ -2,7 +2,6 @@
 import re
 import sys
 import json
-import atexit
 import base64
 import struct
 import urllib
@@ -713,14 +712,16 @@ def get_instance() -> "LocalDataStore":
                 ensure_dir(ds_path)
 
                 LocalDataStore._instance = LocalDataStore(str(ds_path))
-                atexit.register(LocalDataStore._instance.dump)
             return LocalDataStore._instance
 
     def __init__(self, root_path: str) -> None:
         self.root_path = root_path
         self.name_pattern = re.compile(r"^[A-Za-z0-9-_/: ]+$")
         self.tables: Dict[str, MemoryTable] = {}
 
+    def __del__(self) -> None:
+        self.dump()
+
     def update_table(
         self,
         table_name: str,
@@ -1028,7 +1029,6 @@ def __init__(self, table_name: str, key_column: str = "id") -> None:
         self.data_store = get_data_store()
         self.cond = threading.Condition()
         self.setDaemon(True)
-        atexit.register(self.close)
         self.start()
 
     def __enter__(self) -> Any:
@@ -1037,10 +1037,12 @@ def __enter__(self) -> Any:
     def __exit__(self, type: Any, value: Any, tb: Any) -> None:
         self.close()
 
+    def __del__(self) -> None:
+        self.close()
+
     def close(self) -> None:
         with self.cond:
             if not self.stopped:
-                atexit.unregister(self.close)
                 self.stopped = True
                 self.cond.notify()
         self.join()

diff --git a/client/tests/sdk/test_base.py b/client/tests/sdk/test_base.py
@@ -1,7 +1,6 @@
 import os
 import tempfile
 import unittest
-from unittest.mock import patch, MagicMock
 
 from starwhale.utils import config as sw_config
 from starwhale.consts import ENV_SW_CLI_CONFIG, ENV_SW_LOCAL_STORAGE
@@ -18,12 +17,7 @@ def setUp(self) -> None:
         self.datastore_root = str(sw_config.SWCliConfigMixed().datastore_dir)
         ensure_dir(self.datastore_root)
 
-        self.mock_atexit = patch("starwhale.api._impl.data_store.atexit", MagicMock())
-        self.mock_atexit.start()
-
     def tearDown(self) -> None:
         empty_dir(self.local_storage)
         os.environ.pop(ENV_SW_CLI_CONFIG, "")
         os.environ.pop(ENV_SW_LOCAL_STORAGE, "")
-
-        self.mock_atexit.stop()