Use generic type annotation for nodes in NxOntology (#9)

merges #9

nxontology/examples.py

@@ -1,15 +1,15 @@
 from nxontology import NXOntology
 
 
-def create_metal_nxo() -> NXOntology:
+def create_metal_nxo() -> NXOntology[str]:
     """
     Metals ontology from Fig 1 of "Semantic Similarity Definition" at
     https://doi.org/10.1016/B978-0-12-809633-8.20401-9.
     Also published at https://jbiomedsem.biomedcentral.com/articles/10.1186/2041-1480-2-5/figures/1
     Note edge direction is opposite of the drawing.
     Edges go from general to specific.
     """
-    nxo = NXOntology()
+    nxo: NXOntology[str] = NXOntology()
     nxo.graph.graph["name"] = "Metals"
     nxo.set_graph_attributes(node_label_attribute="{node}")
     edges = [
@@ -27,13 +27,13 @@ def create_metal_nxo() -> NXOntology:
     return nxo
 
 
-def create_disconnected_nxo() -> NXOntology:
+def create_disconnected_nxo() -> NXOntology[str]:
     """
     Fictitious ontology with disjoint / disconnected components.
     Has multiple root nodes. Helpful for testing.
     https://github.com/related-sciences/nxontology/issues/4
     """
-    nxo = NXOntology()
+    nxo: NXOntology[str] = NXOntology()
     nxo.add_node("water")
     edges = [
         ("metal", "precious"),

nxontology/imports.py

@@ -8,13 +8,13 @@
 from nxontology.exceptions import NodeNotFound
 
 
-def pronto_to_nxontology(onto: Prontology) -> NXOntology:
+def pronto_to_nxontology(onto: Prontology) -> NXOntology[str]:
     """
     Create an `NXOntology` from an input `pronto.Ontology`.
     Obsolete terms are omitted as nodes.
     Only is_a / subClassOf relationships are used for edges.
     """
-    nxo = NXOntology()
+    nxo: NXOntology[str] = NXOntology()
     nxo.pronto = onto  # type: ignore [attr-defined]
     for term in onto.terms():
         if term.obsolete:
@@ -41,7 +41,7 @@ def pronto_to_nxontology(onto: Prontology) -> NXOntology:
     return nxo
 
 
-def from_obo_library(slug: str) -> NXOntology:
+def from_obo_library(slug: str) -> NXOntology[str]:
     """
     Read ontology from <http://www.obofoundry.org/>.
     Delegates to [`pronto.Ontology.from_obo_library`](https://pronto.readthedocs.io/en/stable/api/pronto.Ontology.html#pronto.Ontology.from_obo_library).
@@ -52,7 +52,7 @@ def from_obo_library(slug: str) -> NXOntology:
     return nxo
 
 
-def from_file(handle: Union[BinaryIO, str, "PathLike[AnyStr]"]) -> NXOntology:
+def from_file(handle: Union[BinaryIO, str, "PathLike[AnyStr]"]) -> NXOntology[str]:
     """
     Read ontology in OBO, OWL, or JSON (OBO Graphs) format via pronto.
 

nxontology/ontology.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import abc
 import functools
 import itertools
@@ -7,6 +9,7 @@
     Any,
     Callable,
     Dict,
+    Generic,
     Hashable,
     Iterable,
     List,
@@ -33,8 +36,7 @@ def frozen(self) -> bool:
 
 # Type definitions. networkx does not declare types.
 # https://github.com/networkx/networkx/issues/3988#issuecomment-639969263
-Node = Hashable
-Node_Set = Set[Node]
+Node = TypeVar("Node", bound=Hashable)
 T = TypeVar("T")
 T_Freezable = TypeVar("T_Freezable", bound=Freezable)
 
@@ -69,7 +71,7 @@ def wrapped(self: T_Freezable) -> T:
     return wrapped
 
 
-class NXOntology(Freezable):
+class NXOntology(Freezable, Generic[Node]):
     """
     Encapsulate a networkx.DiGraph to represent an ontology.
     Regarding edge directionality, parent terms should point to child term.
@@ -79,7 +81,7 @@ class NXOntology(Freezable):
     def __init__(self, graph: Optional[nx.DiGraph] = None):
         self.graph = nx.DiGraph(graph)
         self.check_is_dag()
-        self._node_info_cache: Dict[Hashable, Node_Info] = {}
+        self._node_info_cache: Dict[Node, Node_Info[Node]] = {}
 
     def check_is_dag(self) -> None:
         if not nx.is_directed_acyclic_graph(self.graph):
@@ -103,7 +105,7 @@ def write_node_link_json(self, path: str) -> None:
             json.dump(obj=nld, fp=write_file, indent=2, ensure_ascii=False)
 
     @classmethod
-    def read_node_link_json(cls, path: str) -> "NXOntology":
+    def read_node_link_json(cls, path: str) -> NXOntology[Node]:
         """
         Retrun a new graph from node-link format as written by `write_node_link_json`.
         """
@@ -141,7 +143,7 @@ def add_edge(self, u_of_edge: Node, v_of_edge: Node, **attr: Any) -> None:
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def roots(self) -> "Node_Set":
+    def roots(self) -> Set[Node]:
         """
         Return all top-level nodes.
         """
@@ -155,7 +157,7 @@ def roots(self) -> "Node_Set":
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def leaves(self) -> "Node_Set":
+    def leaves(self) -> Set[Node]:
         """
         Return all bottom-level nodes.
         """
@@ -185,7 +187,7 @@ def similarity(
         node_0: Node,
         node_1: Node,
         ic_metric: str = "intrinsic_ic_sanchez",
-    ) -> "SimilarityIC":
+    ) -> SimilarityIC[Node]:
         """SimilarityIC instance for the specified nodes"""
         return SimilarityIC(self, node_0, node_1, ic_metric)
 
@@ -219,7 +221,7 @@ def compute_similarities(
             metrics = self.similarity_metrics(node_0, node_1, ic_metric=ic_metric)
             yield metrics
 
-    def node_info(self, node: Node) -> "Node_Info":
+    def node_info(self, node: Node) -> Node_Info[Node]:
         """
         Return Node_Info instance for `node`.
         If frozen, cache node info in `self._node_info_cache`.
@@ -271,7 +273,7 @@ def set_graph_attributes(
             self.graph.graph["node_url_attribute"] = node_url_attribute
 
 
-class Node_Info(Freezable):
+class Node_Info(Freezable, Generic[Node]):
     """
     Compute metrics and values for a node of an NXOntology.
     Includes intrinsic information content (IC) metrics.
@@ -288,7 +290,7 @@ class Node_Info(Freezable):
     Each ic_metric has a scaled version accessible by adding a _scaled suffix.
     """
 
-    def __init__(self, nxo: NXOntology, node: Node):
+    def __init__(self, nxo: NXOntology[Node], node: Node):
         if node not in nxo.graph:
             raise NodeNotFound(f"{node} not in graph.")
         self.nxo = nxo
@@ -337,7 +339,7 @@ def data(self) -> Dict[Any, Any]:
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def ancestors(self) -> Node_Set:
+    def ancestors(self) -> Set[Node]:
         """
         Get ancestors of node in graph, including the node itself.
         Ancestors refers to more general concepts in an ontology,
@@ -350,7 +352,7 @@ def ancestors(self) -> Node_Set:
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def descendants(self) -> Node_Set:
+    def descendants(self) -> Set[Node]:
         """
         Get descendants of node in graph, including the node itself.
         Descendants refers to more specific concepts in an ontology,
@@ -445,7 +447,7 @@ def intrinsic_ic_sanchez_scaled(self) -> float:
         return self.intrinsic_ic_sanchez / math.log(len(self.nxo.leaves) + 1)
 
 
-class Similarity(Freezable):
+class Similarity(Freezable, Generic[Node]):
     """
     Compute intrinsic similarity metrics for a pair of nodes.
     """
@@ -461,7 +463,7 @@ class Similarity(Freezable):
         "batet_log",
     ]
 
-    def __init__(self, nxo: NXOntology, node_0: Node, node_1: Node):
+    def __init__(self, nxo: NXOntology[Node], node_0: Node, node_1: Node):
         self.nxo = nxo
         self.node_0 = node_0
         self.node_1 = node_1
@@ -484,12 +486,12 @@ def node_1_subsumes_0(self) -> bool:
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def common_ancestors(self) -> "Node_Set":
+    def common_ancestors(self) -> Set[Node]:
         return self.info_0.ancestors & self.info_1.ancestors
 
     @property  # type: ignore [misc]
     @cache_on_frozen
-    def union_ancestors(self) -> "Node_Set":
+    def union_ancestors(self) -> Set[Node]:
         return self.info_0.ancestors | self.info_1.ancestors
 
     @property
@@ -531,7 +533,7 @@ def results(self, keys: Optional[List[str]] = None) -> Dict[str, Any]:
         return {key: getattr(self, key) for key in keys}
 
 
-class SimilarityIC(Similarity):
+class SimilarityIC(Similarity[Node]):
     """
     Compute intrinsic similarity metrics for a pair of nodes,
     including Information Content (IC) derived metrics.
@@ -540,7 +542,7 @@ class SimilarityIC(Similarity):
 
     def __init__(
         self,
-        graph: NXOntology,
+        graph: NXOntology[Node],
         node_0: Node,
         node_1: Node,
         ic_metric: str = "intrinsic_ic_sanchez",
@@ -565,7 +567,7 @@ def __init__(
         "jiang_seco",
     ]
 
-    def _get_ic(self, node_info: Node_Info, ic_metric: str) -> float:
+    def _get_ic(self, node_info: Node_Info[Node], ic_metric: str) -> float:
         ic = getattr(node_info, ic_metric)
         assert isinstance(ic, float)
         return ic

nxontology/tests/conftest.py

@@ -5,13 +5,13 @@
 
 
 @pytest.fixture
-def metal_nxo() -> NXOntology:
+def metal_nxo() -> NXOntology[str]:
     """Returns a newly created metal ontology for each test."""
     return create_metal_nxo()
 
 
 @pytest.fixture(scope="module")
-def metal_nxo_frozen() -> NXOntology:
+def metal_nxo_frozen() -> NXOntology[str]:
     """
     Frozen metals ontology,
     scoped such that all tests in this module will receive the same NXOntology instance.