feat: add DB status

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

Author: ktro2828

t4_devkit/cli/sanity.py

@@ -24,13 +24,10 @@ def _run_sanity_check(
     revision: str | None = None,
     include_warning: bool = False,
 ) -> list[DBException]:
-    exceptions: list[DBException] = []
-
-    for db_root in tqdm(Path(db_parent).glob("*"), desc=">>>Sanity checking..."):
-        result = sanity_check(db_root, revision=revision, include_warning=include_warning)
-        if result:
-            exceptions.append(result)
-    return exceptions
+    return [
+        sanity_check(db_root, revision=revision, include_warning=include_warning)
+        for db_root in tqdm(Path(db_parent).glob("*"), desc=">>>Sanity checking...")
+    ]
 
 
 @cli.command()
@@ -53,10 +50,10 @@ def main(
 ) -> None:
     exceptions = _run_sanity_check(db_parent, revision=revision, include_warning=include_warning)
 
-    if not exceptions:
+    if all(e.is_ok() for e in exceptions):
         print("✅ No exceptions occurred!!")
     else:
         print("⚠️  Encountered some exceptions!!")
-        headers = ["DatasetID", "Version", "Message"]
-        table = [[e.dataset_id, e.version, e.message] for e in exceptions]
+        headers = ["DatasetID", "Version", "status", "Message"]
+        table = [[e.dataset_id, e.version, e.status, e.message] for e in exceptions]
         print(tabulate(table, headers=headers, tablefmt="pretty"))

t4_devkit/common/sanity.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from enum import Enum, unique
 from pathlib import Path
 
 from attrs import define
@@ -12,19 +13,38 @@
 
 @define
 class DBException:
-    """A dataclass to store error message of the corresponding dataset."""
+    """A dataclass to store error message of the corresponding dataset.
+
+    Attributes:
+        dataset_id (str): Dataset ID.
+        version (str | None): Dataset version.
+        status (DBStatus): Status of the dataset.
+        message (str): Error or warning message.
+    """
 
     dataset_id: str
     version: str | None
-    message: str
+    status: DBStatus
+    message: str | None = None
+
+    def is_ok(self) -> bool:
+        """Return True if the status is OK."""
+        return self.status == DBStatus.OK
+
+
+@unique
+class DBStatus(str, Enum):
+    OK = "OK"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
 
 
 def sanity_check(
     db_root: str | Path,
     *,
     revision: str | None = None,
     include_warning: bool = False,
-) -> DBException | None:
+) -> DBException:
     """Perform sanity check and report exception or warning encountered while loading the dataset.
 
     Args:
@@ -44,14 +64,28 @@ def sanity_check(
             warnings.filterwarnings("ignore")
 
         try:
-            _ = Tier4(data_root=db_root, revision=revision, verbose=False)
-            exception = None
+            t4 = Tier4(data_root=db_root, revision=revision, verbose=False)
+            exception = DBException(
+                dataset_id=t4.dataset_id,
+                version=t4.version,
+                status=DBStatus.OK,
+            )
+        except Warning as w:
+            metadata = load_metadata(db_root)
+
+            exception = DBException(
+                dataset_id=metadata.dataset_id,
+                version=metadata.version,
+                status=DBStatus.WARNING,
+                message=str(w),
+            )
         except Exception as e:
             metadata = load_metadata(db_root)
 
             exception = DBException(
                 dataset_id=metadata.dataset_id,
                 version=metadata.version,
+                status=DBStatus.ERROR,
                 message=str(e),
             )
     return exception