Fix mypy type errors and improve type annotations

- Fix method signature mismatches between base and derived classes
- Update load() and save() methods to return bool instead of dict
- Fix return types for methods that should return List but were returning Set
- Change dictionary type annotations from Dict[str, Any] to Dict[Any, Any] to support various key types
- Fix draw_hypergraph() function to accept BaseHypergraphDB instead of HypergraphDB
- Add proper type annotations and fix Liskov Substitution Principle violations
- Update HypergraphViewer constructor to accept BaseHypergraphDB
- Add missing imports for BaseHypergraphDB in draw.py
- Format code with black to maintain consistency

Author: yifanfeng97

hyperdb/base.py

@@ -12,13 +12,16 @@ class BaseHypergraphDB:
 
     storage_file: Union[str, Path] = field(default="my_hypergraph.hgdb", compare=False)
 
-    def save(self, file_path: Union[str, Path]):
+    def save(self, file_path: Union[str, Path]) -> bool:
         r"""
         Save the hypergraph to a file.
 
         Args:
             ``file_path`` (``Union[str, Path]``): The file path to save the
                 hypergraph.
+
+        Returns:
+            ``bool``: True if successful, False otherwise.
         """
         raise NotImplementedError
 
@@ -32,13 +35,15 @@ def save_as(self, format: str, file_path: Union[str, Path]):
         """
         raise NotImplementedError
 
-    @staticmethod
-    def load(self, file_path: Union[str, Path]):
+    def load(self, file_path: Union[str, Path]) -> bool:
         r"""
         Load the hypergraph from a file.
 
         Args:
             ``file_path`` (``Union[str, Path]``): The file path to load the hypergraph from.
+
+        Returns:
+            ``bool``: True if successful, False otherwise.
         """
         raise NotImplementedError
 
@@ -153,21 +158,23 @@ def remove_e(self, e_tuple: Tuple):
         """
         raise NotImplementedError
 
-    def update_v(self, v_id: Any):
+    def update_v(self, v_id: Any, v_data: Dict[Any, Any]):
         r"""
         Update the vertex data.
 
         Args:
             ``v_id`` (``Any``): The vertex id.
+            ``v_data`` (``Dict[Any, Any]``): The vertex data.
         """
         raise NotImplementedError
 
-    def update_e(self, e_tuple: Tuple):
+    def update_e(self, e_tuple: Union[List[Any], Set[Any], Tuple[Any, ...]], e_data: Dict[Any, Any]):
         r"""
         Update the hyperedge data.
 
         Args:
-            ``e_tuple`` (``Tuple``): The hyperedge tuple: (v1_name, v2_name, ..., vn_name).
+            ``e_tuple`` (``Union[List[Any], Set[Any], Tuple[Any, ...]]``): The hyperedge tuple.
+            ``e_data`` (``Dict[Any, Any]``): The hyperedge data.
         """
         raise NotImplementedError
 

hyperdb/draw.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict
 from urllib.parse import parse_qs, urlparse
 
+from .base import BaseHypergraphDB
 from .hypergraph import HypergraphDB
 
 
@@ -224,7 +225,7 @@ def _get_html_template(self) -> str:
 class HypergraphViewer:
     """Hypergraph visualization tool"""
 
-    def __init__(self, hypergraph_db: HypergraphDB, port: int = 8080):
+    def __init__(self, hypergraph_db: BaseHypergraphDB, port: int = 8080):
         self.hypergraph_db = hypergraph_db
         self.port = port
 
@@ -262,7 +263,9 @@ def stop_server(self):
             self.httpd.server_close()
 
 
-def draw_hypergraph(hypergraph_db: HypergraphDB, port: int = 8080, open_browser: bool = True, blocking: bool = True):
+def draw_hypergraph(
+    hypergraph_db: BaseHypergraphDB, port: int = 8080, open_browser: bool = True, blocking: bool = True
+):
     """
     Main function to draw hypergraph
 

hyperdb/hypergraph.py

@@ -16,9 +16,9 @@ class HypergraphDB(BaseHypergraphDB):
     Hypergraph database.
     """
 
-    _v_data: Dict[str, Any] = field(default_factory=dict)
+    _v_data: Dict[Any, Any] = field(default_factory=dict)
     _e_data: Dict[Tuple, Any] = field(default_factory=dict)
-    _v_inci: Dict[str, Set[Tuple]] = field(default_factory=lambda: defaultdict(set))
+    _v_inci: Dict[Any, Set[Tuple]] = field(default_factory=lambda: defaultdict(set))
 
     def __post_init__(self):
         assert isinstance(self.storage_file, (str, Path))
@@ -27,7 +27,7 @@ def __post_init__(self):
         if self.storage_file.exists():
             self.load(self.storage_file)
 
-    def load(self, storage_file: Path) -> dict:
+    def load(self, storage_file: Union[str, Path]) -> bool:
         r"""
         Load the hypergraph database from the storage file.
         """
@@ -41,7 +41,7 @@ def load(self, storage_file: Path) -> dict:
         except Exception:
             return False
 
-    def save(self, storage_file: Path) -> dict:
+    def save(self, storage_file: Union[str, Path]) -> bool:
         r"""
         Save the hypergraph database to the storage file.
         """
@@ -114,14 +114,14 @@ def all_v(self) -> List[str]:
         r"""
         Return a list of all vertices in the hypergraph.
         """
-        return set(self._v_data.keys())
+        return list(self._v_data.keys())
 
     @cached_property
     def all_e(self) -> List[Tuple]:
         r"""
         Return a list of all hyperedges in the hypergraph.
         """
-        return set(self._e_data.keys())
+        return list(self._e_data.keys())
 
     @cached_property
     def num_v(self) -> int:
@@ -307,7 +307,7 @@ def nbr_e_of_v(self, v_id: Any) -> list:
         """
         assert isinstance(v_id, Hashable), "The vertex id must be hashable."
         assert v_id in self._v_data, f"The vertex {v_id} does not exist in the hypergraph."
-        return set(self._v_inci[v_id])
+        return list(self._v_inci[v_id])
 
     def nbr_v_of_e(self, e_tuple: Union[List, Set, Tuple]) -> list:
         r"""
@@ -319,7 +319,7 @@ def nbr_v_of_e(self, e_tuple: Union[List, Set, Tuple]) -> list:
         assert isinstance(e_tuple, (list, set, tuple)), "The hyperedge must be a list, set, or tuple of vertex ids."
         e_tuple = self.encode_e(e_tuple)
         assert e_tuple in self._e_data, f"The hyperedge {e_tuple} does not exist in the hypergraph."
-        return set(e_tuple)
+        return list(e_tuple)
 
     def nbr_v(self, v_id: Any, exclude_self=True) -> list:
         r"""
@@ -330,9 +330,9 @@ def nbr_v(self, v_id: Any, exclude_self=True) -> list:
         """
         assert isinstance(v_id, Hashable), "The vertex id must be hashable."
         assert v_id in self._v_data, f"The vertex {v_id} does not exist in the hypergraph."
-        nbrs = set()
+        nbrs: set = set()
         for e_tuple in self._v_inci[v_id]:
             nbrs.update(e_tuple)
         if exclude_self:
             nbrs.remove(v_id)
-        return set(nbrs)
+        return list(nbrs)

tests/test_hypergraph.py

@@ -85,7 +85,7 @@ def test_remove_v(hg):
     hg.remove_v(6)
     assert hg.has_v(6) is False
     assert hg.has_e((1, 5, 6)) is False
-    assert hg.has_e((5, 1)) == {"relation": "study"}
+    assert hg.e((1, 5)) == {"relation": "study"}
 
 
 def test_remove_e(hg):