Improve typing in algorithms module

pyt-team · Mar 5, 2024 · 172a237 · 172a237
1 parent 28bc02b
commit 172a237
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 38 deletions.
diff --git a/toponetx/algorithms/components.py b/toponetx/algorithms/components.py
@@ -1,12 +1,12 @@
 """Module to compute connected components on topological domains."""
 from collections.abc import Generator, Hashable
+from typing import TypeVar
 
 import networkx as nx
 
 from toponetx.classes.cell_complex import CellComplex
 from toponetx.classes.colored_hypergraph import ColoredHyperGraph
 from toponetx.classes.combinatorial_complex import CombinatorialComplex
-from toponetx.classes.complex import Complex
 
 __all__ = [
     "s_connected_components",
@@ -15,22 +15,27 @@
     "connected_component_subcomplexes",
 ]
 
+# In this module, only cell complexes, combinatorial complexes and colored
+# hypergraphs are supported. We bound the type variable to these types.
+ComplexType = CellComplex | CombinatorialComplex | ColoredHyperGraph
+ComplexTypeVar = TypeVar("ComplexTypeVar", bound=ComplexType)
+
 
 def s_connected_components(
-    domain: Complex, s: int = 1, cells: bool = True, return_singletons: bool = False
+    domain: ComplexType, s: int, cells: bool = True, return_singletons: bool = False
 ) -> Generator[set[Hashable] | set[tuple[Hashable, ...]], None, None]:
     """Return generator for the s-connected components.
 
     Parameters
     ----------
-    domain : Complex
-        Supported complexes are cell/combintorial and hypegraphs.
-    s : int, optional
+    domain : CellComplex or CombinatorialComplex or ColoredHyperGraph
+        The domain on which to compute the s-connected components.
+    s : int
         The number of intersections between pairwise consecutive cells.
-    cells : bool, optional
+    cells : bool, default=True
         If True will return cell components, if False will return node components.
-    return_singletons : bool, optional
-        When True, returns singletons connected components.
+    return_singletons : bool, default=False
+        When True, returns singleton connected components.
 
     Notes
     -----
@@ -60,21 +65,17 @@ def s_connected_components(
     >>> CC.add_cell([2, 3, 4], rank=2)
     >>> CC.add_cell([5, 6, 7], rank=2)
     >>> list(s_connected_components(CC, s=1, cells=False))
-    >>> #  [{2, 3, 4}, {5, 6, 7}]
+    [{2, 3, 4}, {5, 6, 7}]
     >>> list(s_connected_components(CC, s=1, cells=True))
-    >>> # [{(2, 3), (2, 3, 4), (2, 4), (3, 4)},
-    >>> # {(5, 6), (5, 6, 7), (5, 7), (6, 7)}]
+    [{(2, 3), (2, 3, 4), (2, 4), (3, 4)}, {(5, 6), (5, 6, 7), (5, 7), (6, 7)}]
     >>> CHG = CC.to_colored_hypergraph()
     >>> list(s_connected_components(CHG, s=1, cells=False))
     >>> CC.add_cell([4, 5], rank=1)
     >>> list(s_connected_components(CC, s=1, cells=False))
-    >>> # [{2, 3, 4, 5, 6, 7}]
+    [{2, 3, 4, 5, 6, 7}]
     >>> CCC = CC.to_combinatorial_complex()
     >>> list(s_connected_components(CCC, s=1, cells=False))
     """
-    if not isinstance(domain, CellComplex | ColoredHyperGraph | CombinatorialComplex):
-        raise TypeError(f"Input complex {domain} is not supported.")
-
     if cells:
         cell_dict, A = domain.all_cell_to_node_coadjacency_matrix(s=s, index=True)
         cell_dict = {v: k for k, v in cell_dict.items()}
@@ -107,23 +108,26 @@ def s_connected_components(
 
 
 def s_component_subcomplexes(
-    domain: Complex, s: int = 1, cells: bool = True, return_singletons: bool = False
-) -> Generator[Complex, None, None]:
+    domain: ComplexTypeVar,
+    s: int = 1,
+    cells: bool = True,
+    return_singletons: bool = False,
+) -> Generator[ComplexTypeVar, None, None]:
     """Return a generator for the induced subcomplexes of s_connected components.
 
     Removes singletons unless return_singletons is set to True.
 
     Parameters
     ----------
-    domain : Complex
-        Supported complexes are cell/combintorial and hypegraphs.
-    s : int, optional
+    domain : CellComplex or CombinatorialComplex or ColoredHyperGraph
+        The domain for which to compute the the s-connected subcomplexes.
+    s : int, default=1
         The number of intersections between pairwise consecutive cells.
-    cells : bool, optional
+    cells : bool, default=True
         Determines if cell or node components are desired. Returns
         subcomplexes equal to the cell complex restricted to each set of nodes(cells) in the
         s-connected components or s-cell-connected components.
-    return_singletons : bool, optional
+    return_singletons : bool, default=False
         When True, returns singletons connected components.
 
     Yields
@@ -155,7 +159,7 @@ def s_component_subcomplexes(
 
 
 def connected_components(
-    domain: Complex, cells: bool = False, return_singletons: bool = True
+    domain: ComplexType, cells: bool = False, return_singletons: bool = True
 ) -> Generator[set[Hashable] | set[tuple[Hashable, ...]], None, None]:
     """Compute s-connected components with s=1.
 
@@ -165,14 +169,14 @@ def connected_components(
     ----------
     domain : Complex
         Supported complexes are cell/combintorial and hypegraphs.
-    cells : bool, optional
+    cells : bool, default=False
         If True will return cell components, if False will return node components.
-    return_singletons : bool, optional
+    return_singletons : bool, default=True
         When True, returns singletons connected components.
 
     Yields
     ------
-    set[Hashable] | set[tuple[Hashable, ...]], None, None
+    set[Hashable] | set[tuple[Hashable, ...]]
         Yields subcomplexes generated by the cells (or nodes) in the
         cell(node) components of complex.
 
@@ -190,20 +194,20 @@ def connected_components(
     >>> CC.add_cell([4, 5], rank=1)
     >>> list(CC.connected_components(CC, cells=False))
     """
-    yield from s_connected_components(domain, s=1, cells=cells, return_singletons=True)
+    yield from s_connected_components(domain, s=1, cells=cells, return_singletons=return_singletons)
 
 
 def connected_component_subcomplexes(
-    domain: Complex, return_singletons: bool = True
-) -> Generator[Complex, None, None]:
+    domain: ComplexTypeVar, return_singletons: bool = True
+) -> Generator[ComplexTypeVar, None, None]:
     """Compute connected component subcomplexes with s=1.
 
     Same as :meth:`s_component_subcomplexes` with s=1.
 
     Parameters
     ----------
-    domain : Complex
-        Supported complexes are cell/combintorial and hypegraphs.
+    domain : CellComplex or CombinaorialComplex or ColoredHyperGraph
+        The domain for which to compute the the connected subcomplexes.
     return_singletons : bool, optional
         When True, returns singletons connected components.
 

diff --git a/toponetx/algorithms/distance.py b/toponetx/algorithms/distance.py
@@ -9,19 +9,24 @@
 from toponetx.classes.cell_complex import CellComplex
 from toponetx.classes.colored_hypergraph import ColoredHyperGraph
 from toponetx.classes.combinatorial_complex import CombinatorialComplex
-from toponetx.classes.complex import Complex
 from toponetx.classes.hyperedge import HyperEdge
 
 __all__ = ["distance", "cell_distance"]
 
+# In this module, only cell complexes, combinatorial complexes and colored
+# hypergraphs are supported.
+ComplexType = CellComplex | CombinatorialComplex | ColoredHyperGraph
 
-def distance(domain: Complex, source: Hashable, target: Hashable, s: int = 1) -> int:
+
+def distance(
+    domain: ComplexType, source: Hashable, target: Hashable, s: int = 1
+) -> int:
     """Return shortest s-walk distance between two nodes in the cell complex.
 
     Parameters
     ----------
-    domain : Complex
-        Supported complexes are cell/combintorial and hypegraphs.
+    domain : CellComplex or CombinaorialComplex or ColoredHyperGraph
+        The domain on which to compute the s-walk distance between source and target.
     source : Hashable
         A node in the input complex.
     target : Hashable
@@ -81,7 +86,7 @@ def distance(domain: Complex, source: Hashable, target: Hashable, s: int = 1) ->
 
 
 def cell_distance(
-    domain: Complex,
+    domain: ComplexType,
     source: Iterable | HyperEdge | Cell,
     target: Iterable | HyperEdge | Cell,
     s: int = 1,
@@ -90,8 +95,8 @@ def cell_distance(
 
     Parameters
     ----------
-    domain : Complex
-        Supported complexes are cell/combintorial and hypegraphs.
+    domain : CellComplex or CombinatorialComplex or ColoredHyperGraph
+        The domain on which to compute the s-walk distance between source and target cells.
     source : Iterable or HyperEdge or Cell
         An Iterable representing a cell in the input complex cell complex.
     target : Iterable or HyperEdge or Cell