catalogue_graph/src/models/indexable_concept.py

@@ -4,7 +4,6 @@
 from pydantic import BaseModel
 
 from models.catalogue_concept import CatalogueConcept
-from models.graph_node import ConceptSource
 
 # Query
 
@@ -32,7 +31,7 @@ class ConceptDisplayIdentifierType(BaseModel):
     type: str = "IdentifierType"
 
     @classmethod
-    def from_source_type(cls, source_type: ConceptSource) -> "ConceptDisplayIdentifierType":
+    def from_source_type(cls, source_type: str) -> "ConceptDisplayIdentifierType":
         if source_type == "label-derived":
             label = "Identifier derived from the label of the referent"
         elif source_type == "nlm-mesh":

catalogue_graph/src/sources/catalogue/concepts_source.py

@@ -6,7 +6,7 @@
 from sources.gzip_source import GZipSource
 
 
-def extract_concepts_from_work(raw_work: dict) -> Generator[(dict, WorkConceptKey)]:
+def extract_concepts_from_work(raw_work: dict) -> Generator[tuple[dict, WorkConceptKey]]:
     for subject in raw_work.get("subjects", []):
         for concept in subject.get("concepts", []):
             yield concept, "subjects"
@@ -24,7 +24,7 @@ class CatalogueConceptsSource(BaseSource):
     def __init__(self, url: str):
         self.url = url
 
-    def stream_raw(self) -> Generator[dict]:
+    def stream_raw(self) -> Generator[tuple[dict, WorkConceptKey]]:
         """Streams raw concept nodes from a work's subjects, genres, and contributors."""
         catalogue_source = GZipSource(self.url)
         for work in catalogue_source.stream_raw():

catalogue_graph/src/transformers/catalogue/concepts_transformer.py

@@ -19,7 +19,7 @@ def __init__(self, url: str):
         )
         self.id_lookup: set = set()
 
-    def transform_node(self, raw_data: (dict, WorkConceptKey)) -> Concept | None:
+    def transform_node(self, raw_data: tuple[dict, WorkConceptKey]) -> Concept | None:
         raw_concept = RawCatalogueConcept(raw_data[0], self.id_label_checker)
 
         if not raw_concept.is_concept:
@@ -37,7 +37,7 @@ def transform_node(self, raw_data: (dict, WorkConceptKey)) -> Concept | None:
             type=raw_concept.type,
         )
 
-    def extract_edges(self, raw_data: (dict, WorkConceptKey)) -> Generator[ConceptHasSourceConcept]:
+    def extract_edges(self, raw_data: tuple[dict, WorkConceptKey]) -> Generator[ConceptHasSourceConcept]:
         raw_concept = RawCatalogueConcept(raw_data[0], self.id_label_checker)
 
         if not raw_concept.is_concept:

catalogue_graph/tests/transformers/test_catalogue_concepts_transformer.py

@@ -25,7 +25,7 @@ def _add_catalogue_request() -> None:
     )
 
 
-def _check_edge(all_edges: list[BaseEdge], from_id: str, to_id: str, expected_edge: BaseEdge):
+def _check_edge(all_edges: list[BaseEdge], from_id: str, to_id: str, expected_edge: BaseEdge) -> None:
     filtered_edges = [edge for edge in all_edges if edge.from_id == from_id and edge.to_id == to_id]
     assert len(filtered_edges) == 1
     assert filtered_edges[0] == expected_edge

catalogue_graph/tests/transformers/test_catalogue_works_transformer.py

@@ -6,7 +6,7 @@
 from transformers.catalogue.works_transformer import CatalogueWorksTransformer
 
 
-def _check_edge(all_edges: list[BaseEdge], from_id: str, to_id: str, expected_edge: BaseEdge):
+def _check_edge(all_edges: list[BaseEdge], from_id: str, to_id: str, expected_edge: BaseEdge) -> None:
     filtered_edges = [edge for edge in all_edges if edge.from_id == from_id and edge.to_id == to_id]
     assert len(filtered_edges) == 1
     assert filtered_edges[0] == expected_edge