Add strict type checking (#50)

Author: cthoyt

src/semra/api.py

@@ -8,7 +8,7 @@
 import typing as t
 from collections import Counter, defaultdict
 from collections.abc import Iterable
-from typing import cast, overload
+from typing import Literal, TypeVar, cast, overload
 
 import bioregistry
 import networkx as nx
@@ -96,15 +96,26 @@
 #: An index allows for the aggregation of evidences for each core triple
 Index = dict[Triple, list[Evidence]]
 
+X = TypeVar("X")
 
-def _tqdm(mappings, desc: str | None = None, *, progress: bool = True, leave: bool = True):
-    return tqdm(
-        mappings,
-        unit_scale=True,
-        unit="mapping",
-        desc=desc,
-        leave=leave,
-        disable=not progress,
+
+def _tqdm(
+    mappings: Iterable[X],
+    desc: str | None = None,
+    *,
+    progress: bool = True,
+    leave: bool = True,
+) -> Iterable[X]:
+    return cast(
+        Iterable[X],
+        tqdm(
+            mappings,
+            unit_scale=True,
+            unit="mapping",
+            desc=desc,
+            leave=leave,
+            disable=not progress,
+        ),
     )
 
 
@@ -292,7 +303,17 @@ def infer_reversible(mappings: t.Iterable[Mapping], *, progress: bool = True) ->
 # TODO infer negative mappings for exact match from narrow/broad match
 
 
-def flip(mapping: Mapping) -> Mapping | None:
+# docstr-coverage:excused `overload`
+@overload
+def flip(mapping: Mapping, *, strict: Literal[True] = True) -> Mapping: ...
+
+
+# docstr-coverage:excused `overload`
+@overload
+def flip(mapping: Mapping, *, strict: Literal[False] = False) -> Mapping | None: ...
+
+
+def flip(mapping: Mapping, *, strict: bool = False) -> Mapping | None:
     """Flip a mapping, if the relation is configured with an inversion.
 
     :param mapping: An input mapping
@@ -305,14 +326,17 @@ def flip(mapping: Mapping) -> Mapping | None:
         with an inversion (e.g., for practical purposes, regular dbrefs and
         close matches are not configured to invert), then None is returned
     """
-    if (p := FLIP.get(mapping.p)) is None:
+    if (p := FLIP.get(mapping.p)) is not None:
+        return Mapping(
+            s=mapping.o,
+            p=p,
+            o=mapping.s,
+            evidence=[ReasonedEvidence(justification=INVERSION_MAPPING, mappings=[mapping])],
+        )
+    elif strict:
+        raise ValueError
+    else:
         return None
-    return Mapping(
-        s=mapping.o,
-        p=p,
-        o=mapping.s,
-        evidence=[ReasonedEvidence(justification=INVERSION_MAPPING, mappings=[mapping])],
-    )
 
 
 def to_digraph(mappings: t.Iterable[Mapping]) -> nx.DiGraph:
@@ -361,7 +385,7 @@ def _from_digraph_edge(graph: nx.Graph, s: Reference, o: Reference) -> t.Iterabl
 def iter_components(mappings: t.Iterable[Mapping]) -> t.Iterable[set[Reference]]:
     """Iterate over connected components in the multidigraph view over the mappings."""
     graph = to_digraph(mappings)
-    return nx.weakly_connected_components(graph)
+    return cast(t.Iterable[set[Reference]], nx.weakly_connected_components(graph))
 
 
 def to_multidigraph(mappings: t.Iterable[Mapping], *, progress: bool = False) -> nx.MultiDiGraph:
@@ -499,9 +523,9 @@ def infer_chains(
     return [*mappings, *new_mappings]
 
 
-def _path_has_prefix_duplicates(path) -> bool:
+def _path_has_prefix_duplicates(path: Iterable[tuple[Reference, Reference, Reference]]) -> bool:
     """Return if the path has multiple unique."""
-    elements = set()
+    elements: set[Reference] = set()
     for u, v, _ in path:
         elements.add(u)
         elements.add(v)
@@ -899,7 +923,7 @@ def project(
     source_prefix: str,
     target_prefix: str,
     *,
-    return_sus: typing.Literal[True] = True,
+    return_sus: typing.Literal[True] = ...,
     progress: bool = False,
 ) -> tuple[list[Mapping], list[Mapping]]: ...
 
@@ -911,7 +935,7 @@ def project(
     source_prefix: str,
     target_prefix: str,
     *,
-    return_sus: typing.Literal[False] = False,
+    return_sus: typing.Literal[False] = ...,
     progress: bool = False,
 ) -> list[Mapping]: ...
 

src/semra/cli.py

@@ -16,7 +16,7 @@
 
 @click.group()
 @click.version_option()
-def main():
+def main() -> None:
     """CLI for SeMRA."""
 
 

src/semra/client.py

@@ -12,10 +12,10 @@
 import neo4j.graph
 import networkx as nx
 import pydantic
-from neo4j import unit_of_work
+from neo4j import ManagedTransaction, unit_of_work
 
 import semra
-from semra import Evidence, MappingSet, Reference
+from semra import Evidence, MappingSet, Reference, SimpleEvidence
 from semra.rules import (
     RELATIONS,
     SEMRA_EVIDENCE_PREFIX,
@@ -28,6 +28,7 @@
     "Node",
 ]
 
+
 Node: TypeAlias = t.Mapping[str, Any]
 
 TxResult: TypeAlias = list[list[Any]] | None
@@ -62,7 +63,7 @@ def __init__(
         uri: str | None = None,
         user: str | None = None,
         password: str | None = None,
-    ):
+    ) -> None:
         """Initialize the client.
 
         :param uri: The URI of the Neo4j database.
@@ -72,10 +73,12 @@ def __init__(
         uri = uri or os.environ.get("NEO4J_URL") or "bolt://0.0.0.0:7687"
         user = user or os.environ.get("NEO4J_USER")
         password = password or os.environ.get("NEO4J_PASSWORD")
-
-        self.driver = neo4j.GraphDatabase.driver(
-            uri=uri, auth=(user, password), max_connection_lifetime=180
-        )
+        auth: tuple[str, str] | None
+        if user is not None and password is not None:
+            auth = user, password
+        else:
+            auth = None
+        self.driver = neo4j.GraphDatabase.driver(uri=uri, auth=auth, max_connection_lifetime=180)
 
         self._all_relations = {curie for (curie,) in self.read_query(RELATIONS_CYPHER)}
         self._rel_q = "|".join(
@@ -84,12 +87,12 @@ def __init__(
             if reference.curie in self._all_relations
         )
 
-    def __del__(self):
+    def __del__(self) -> None:
         """Ensure driver is shut down when client is destroyed."""
         if self.driver is not None:
             self.driver.close()
 
-    def read_query(self, query: str, **query_params) -> list[list]:
+    def read_query(self, query: str, **query_params: Any) -> list[list[Any]]:
         """Run a read-only query.
 
         :param query: The cypher query to run
@@ -98,11 +101,11 @@ def read_query(self, query: str, **query_params) -> list[list]:
         :returns: The result of the query
         """
         with self.driver.session() as session:
-            values = session.execute_read(_do_cypher_tx, query, **query_params)
+            values = session.execute_read(_do_cypher_tx, query, **query_params)  # type:ignore
 
         return values
 
-    def write_query(self, query: str, **query_params):
+    def write_query(self, query: str, **query_params: Any) -> None:
         """Run a write query.
 
         :param query: The cypher query to run
@@ -111,7 +114,7 @@ def write_query(self, query: str, **query_params):
         :returns: The result of the write query
         """
         with self.driver.session() as session:
-            return session.write_transaction(_do_cypher_tx, query, **query_params)
+            session.write_transaction(_do_cypher_tx, query, **query_params)  # type:ignore
 
     def create_single_property_node_index(
         self, index_name: str, label: str, property_name: str, *, exist_ok: bool = False
@@ -137,7 +140,7 @@ def _get_node_by_curie(self, curie: ReferenceHint, node_type: str | None = None)
             curie = curie.curie
         query = "MATCH (n%s {curie: $curie}) RETURN n" % (":" + node_type if node_type else "")
         res = self.read_query(query, curie=curie)
-        return res[0][0]
+        return cast(Node, res[0][0])
 
     def get_mapping(self, curie: ReferenceHint) -> semra.Mapping:
         """Get a mapping.
@@ -217,7 +220,7 @@ def get_evidence(self, curie: ReferenceHint) -> Evidence:
         curie = _safe_curie(curie, SEMRA_EVIDENCE_PREFIX)
         query = "MATCH (n:evidence {curie: $curie}) RETURN n"
         res = self.read_query(query, curie=curie)
-        return res[0][0]
+        return SimpleEvidence.model_validate(res[0][0])  # FIXME test this?
 
     def summarize_predicates(self) -> t.Counter[str]:
         """Get a counter of predicates."""
@@ -301,7 +304,7 @@ def get_exact_matches(
             RETURN b.curie, b.name
         """
         return {
-            cast(Reference, Reference.from_curie(n_curie)): name
+            Reference.from_curie(n_curie, name=name): name
             for n_curie, name in self.read_query(query, curie=curie)
         }
 
@@ -359,8 +362,8 @@ def get_connected_component_graph(self, curie: ReferenceHint) -> nx.MultiDiGraph
         for path in paths:
             for relationship in path.relationships:
                 g.add_edge(
-                    path.start_node["curie"],  # type: ignore
-                    path.end_node["curie"],  # type: ignore
+                    path.start_node["curie"],
+                    path.end_node["curie"],
                     key=relationship.id,
                     type=relationship.type,
                 )
@@ -375,9 +378,11 @@ def get_concept_name(self, curie: ReferenceHint) -> str | None:
         except Exception:
             return None
         else:
-            return name
+            return cast(str, name)
 
-    def sample_mappings_from_set(self, curie: ReferenceHint, n: int = 10) -> list:
+    def sample_mappings_from_set(
+        self, curie: ReferenceHint, n: int = 10
+    ) -> list[tuple[str, str, str, str, str, str]]:
         """Get n mappings from a given set (by CURIE)."""
         if isinstance(curie, Reference):
             curie = curie.curie
@@ -392,13 +397,13 @@ def sample_mappings_from_set(self, curie: ReferenceHint, n: int = 10) -> list:
         RETURN n.curie, n.predicate, s.curie, s.name, t.curie, t.name
         LIMIT {n}
         """
-        return list(self.read_query(query, curie=curie))
+        return list(self.read_query(query, curie=curie))  # type:ignore
 
 
 # Follows example here:
 # https://neo4j.com/docs/python-manual/current/session-api/#python-driver-simple-transaction-fn
 # and from the docstring of neo4j.Session.read_transaction
 @unit_of_work()
-def _do_cypher_tx(tx, query, **query_params) -> list[list]:
+def _do_cypher_tx(tx: ManagedTransaction, query: str, **query_params: Any) -> list[list[Any]]:
     result = tx.run(query, parameters=query_params)
     return [record.values() for record in result]

src/semra/gilda_utils.py

@@ -6,11 +6,11 @@
 import logging
 import typing as t
 from collections import defaultdict
+from typing import cast
 
 import bioregistry
 import gilda
-from gilda import Term
-from gilda.term import filter_out_duplicates
+import gilda.term
 from tabulate import tabulate
 from tqdm.auto import tqdm
 from tqdm.contrib.concurrent import process_map
@@ -49,7 +49,7 @@
 REVERSE_GILDA_MAP = {v: k for k, v in GILDA_TO_BIOREGISTRY.items()}
 
 
-def update_terms(terms: list[Term], mappings: list[Mapping]) -> list[Term]:
+def update_terms(terms: list[gilda.Term], mappings: list[Mapping]) -> list[gilda.Term]:
     """Use a priority mapping to re-write terms with priority groundings.
 
     :param terms: A list of Gilda term objects
@@ -98,24 +98,29 @@ def update_terms(terms: list[Term], mappings: list[Mapping]) -> list[Term]:
 
     # Unwind the terms index
     new_terms = list(itt.chain.from_iterable(terms_index.values()))
-    return filter_out_duplicates(new_terms)
+    return cast(list[gilda.Term], gilda.term.filter_out_duplicates(new_terms))
 
 
-def standardize_terms(terms: t.Iterable[Term], *, multiprocessing: bool = True) -> list[Term]:
+def standardize_terms(
+    terms: t.Iterable[gilda.Term], *, multiprocessing: bool = True
+) -> list[gilda.Term]:
     """Standardize a list of terms."""
     if not multiprocessing:
         return [standardize_term(t) for t in terms]
-    return process_map(
-        standardize_term,
-        terms,
-        unit="term",
-        unit_scale=True,
-        desc="standardizing",
-        chunksize=40_000,
+    return cast(
+        list[gilda.Term],
+        process_map(
+            standardize_term,
+            terms,
+            unit="term",
+            unit_scale=True,
+            desc="standardizing",
+            chunksize=40_000,
+        ),
     )
 
 
-def standardize_term(term: Term) -> Term:
+def standardize_term(term: gilda.Term) -> gilda.Term:
     """Standardize a term's prefix and identifier to the Bioregistry standard."""
     prefix = bioregistry.normalize_prefix(term.db)
     if prefix is None:
@@ -132,17 +137,17 @@ def standardize_term(term: Term) -> Term:
 
 
 def make_new_term(
-    term: Term,
+    term: gilda.Term,
     target_db: str,
     target_id: str,
     target_name: str | None = None,
-) -> Term:
+) -> gilda.Term:
     """Make a new gilda term object by replacing the database, identifier, and name."""
     if target_name is None:
         from indra.ontology.bio import bio_ontology
 
         target_name = bio_ontology.get_name(target_db, target_id)
-    return Term(
+    return gilda.Term(
         norm_text=term.norm_text,
         text=term.text,
         db=target_db,