From e6d536e6c128190a2cce3c03643c8ec069da56e1 Mon Sep 17 00:00:00 2001
From: Wenjie Du <wenjay.du@gmail.com>
Date: Fri, 28 Jun 2024 14:07:40 +0800
Subject: [PATCH] feat: add the arg check_exists check_path();

---
 tsdb/utils/file.py | 31 +++++++++++++++++++++----------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/tsdb/utils/file.py b/tsdb/utils/file.py
index 76c7170..818b25b 100644
--- a/tsdb/utils/file.py
+++ b/tsdb/utils/file.py
@@ -15,7 +15,10 @@
 from .logging import logger
 
 
-def check_path(path: str) -> str:
+def check_path(
+    path: str,
+    check_exists: bool = False,
+) -> str:
     """Check the given path and return the absolute path.
 
     Parameters
@@ -23,10 +26,14 @@ def check_path(path: str) -> str:
     path :
         The path to be checked.
 
+    check_exists :
+        If True, check if the path exists, and will raise an AssertionError if the path does not exist.
+
     Returns
     -------
     checked_path:
         The absolute path of the given path.
+
     """
     # expand the home dir if the path starts with "~"
     if path.startswith("~"):
@@ -35,6 +42,12 @@ def check_path(path: str) -> str:
         checked_path = path
 
     checked_path = os.path.abspath(checked_path)
+
+    if check_exists:
+        assert os.path.exists(
+            checked_path
+        ), f"The given path {checked_path} does not exists"
+
     return checked_path
 
 
@@ -54,6 +67,9 @@ def pickle_dump(data: object, path: str) -> Optional[str]:
     `path` if succeed else None
 
     """
+    # check the given path
+    path = check_path(path)
+
     try:
         with open(path, "wb") as f:
             pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
@@ -78,6 +94,8 @@ def pickle_load(path: str) -> object:
         Pickled object.
 
     """
+    # check the given path
+    path = check_path(path, check_exists=True)
     try:
         with open(path, "rb") as f:
             data = pickle.load(f)
@@ -99,11 +117,8 @@ def purge_path(path: str, ignore_errors: bool = True) -> None:
         Errors are ignored if ignore_errors is set.
 
     """
-    # check the path
+    # check the given path, no need to check if the path exists because ignore_errors is set
     path = check_path(path)
-    assert os.path.exists(
-        path
-    ), f"The given path {path} does not exists. Operation aborted."
 
     try:
         if os.path.isdir(path):
@@ -181,13 +196,9 @@ def migrate(old_path: str, new_path: str) -> None:
 
     """
     # check both old_path and new_path
-    old_path = check_path(old_path)
+    old_path = check_path(old_path, check_exists=True)
     new_path = check_path(new_path)
 
-    # check if old_path exists
-    if not os.path.exists(old_path):
-        raise FileNotFoundError(f"Given old_path {old_path} does not exist.")
-
     # create new_path if not exists
     if not os.path.exists(new_path):
         os.makedirs(new_path, exist_ok=True)