From e90e58a0dc8418a5e99e21466a597ac8e322dbbe Mon Sep 17 00:00:00 2001
From: vinoyang <yanghua1127@gmail.com>
Date: Mon, 26 Aug 2024 15:24:17 +0800
Subject: [PATCH] Core: Implement isdir api (#34)

---
 tosfs/core.py             | 48 +++++++++++++++++++++++++++++++++++++++
 tosfs/tests/test_tosfs.py | 17 ++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/tosfs/core.py b/tosfs/core.py
index 2d0d20f..9d5c3ec 100644
--- a/tosfs/core.py
+++ b/tosfs/core.py
@@ -275,6 +275,54 @@ def touch(self, path: str, truncate: bool = True, **kwargs: Any) -> None:
         except Exception as e:
             raise TosfsError(f"Tosfs failed with unknown error: {e}") from e
 
+    def isdir(self, path: str) -> bool:
+        """Check if the path is a directory.
+
+        Parameters
+        ----------
+        path : str
+            The path to check.
+
+        Returns
+        -------
+        bool
+            True if the path is a directory, False otherwise.
+
+        Raises
+        ------
+        TosClientError
+            If there is a client error while accessing the path.
+        TosServerError
+            If there is a server error while accessing the path.
+        TosfsError
+            If there is an unknown error while accessing the path.
+
+        Examples
+        --------
+        >>> fs = TosFileSystem()
+        >>> fs.isdir("tos://mybucket/mydir/")
+
+        """
+        path = self._strip_protocol(path).rstrip("/") + "/"
+        bucket, key, _ = self._split_path(path)
+        if not key:
+            return False
+
+        key = key.rstrip("/") + "/"
+
+        try:
+            self.tos_client.head_object(bucket, key)
+            return True
+        except tos.exceptions.TosClientError as e:
+            raise e
+        except tos.exceptions.TosServerError as e:
+            if e.status_code == SERVER_RESPONSE_CODE_NOT_FOUND:
+                return False
+            else:
+                raise e
+        except Exception as e:
+            raise TosfsError(f"Tosfs failed with unknown error: {e}") from e
+
     def _bucket_info(self, bucket: str) -> dict:
         """Get the information of a bucket.
 
diff --git a/tosfs/tests/test_tosfs.py b/tosfs/tests/test_tosfs.py
index 488c1b2..4e3a3e3 100644
--- a/tosfs/tests/test_tosfs.py
+++ b/tosfs/tests/test_tosfs.py
@@ -133,3 +133,20 @@ def test_touch(tosfs: TosFileSystem, bucket: str, temporary_workspace: str) -> N
     tosfs.touch(f"{bucket}/{temporary_workspace}/{file_name}", truncate=True)
     assert tosfs.info(f"{bucket}/{temporary_workspace}/{file_name}")["size"] == 0
     tosfs.rm_file(f"{bucket}/{temporary_workspace}/{file_name}")
+
+
+def test_isdir(tosfs: TosFileSystem, bucket: str, temporary_workspace: str) -> None:
+    assert not tosfs.isdir("")
+    assert not tosfs.isdir("/")
+    assert not tosfs.isdir(bucket)
+    assert tosfs.isdir(f"{bucket}/{temporary_workspace}")
+    assert tosfs.isdir(f"{bucket}/{temporary_workspace}/")
+    assert not tosfs.isdir(f"{bucket}/{temporary_workspace}/nonexistent")
+    assert not tosfs.isdir(f"{bucket}/{temporary_workspace}/nonexistent/")
+
+    file_name = random_path()
+    tosfs.tos_client.put_object(bucket=bucket, key=f"{temporary_workspace}/{file_name}")
+    assert not tosfs.isdir(f"{bucket}/{temporary_workspace}/{file_name}")
+    assert not tosfs.isdir(f"{bucket}/{temporary_workspace}/{file_name}/")
+
+    tosfs._rm(f"{bucket}/{temporary_workspace}/{file_name}")