refactor: replace custom missing provider implementation with utility function

Author: CaliLuke

graphium_core/cross_encoder/rerankers/__init__.py

@@ -1,33 +1,23 @@
 from __future__ import annotations
 
+from graphium_core.utils import missing_provider
+
 from .generic import GenericLiteLLMReranker
 from .openai import OpenAIReranker
 
 __all__ = ['GenericLiteLLMReranker', 'OpenAIReranker']
 
 
-def _missing_provider(name: str, dependency: str) -> type:
-    class _MissingProvider:  # pragma: no cover - simple error shim
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                f'{name} requires the optional dependency "{dependency}". '
-                f'Install it via `uv sync --extra {dependency}` or the matching extras group.'
-            )
-
-    _MissingProvider.__name__ = f'Missing{name}'
-    return _MissingProvider
-
-
 try:  # pragma: no cover - optional dependency
     from .gemini import GeminiReranker  # type: ignore[unused-import]
 except ImportError:  # pragma: no cover
-    GeminiReranker = _missing_provider('GeminiReranker', 'google-genai')  # type: ignore[assignment]
+    GeminiReranker = missing_provider('GeminiReranker', 'google-genai')  # type: ignore[assignment]
 else:  # pragma: no cover
     __all__.append('GeminiReranker')
 
 try:  # pragma: no cover - optional dependency
     from .bge import BGEReranker  # type: ignore[unused-import]
 except ImportError:  # pragma: no cover
-    BGEReranker = _missing_provider('BGEReranker', 'sentence-transformers')  # type: ignore[assignment]
+    BGEReranker = missing_provider('BGEReranker', 'sentence-transformers')  # type: ignore[assignment]
 else:  # pragma: no cover
     __all__.append('BGEReranker')

graphium_core/embedder/providers/__init__.py

@@ -1,42 +1,32 @@
 from __future__ import annotations
 
+from graphium_core.utils import missing_provider
+
 from .openai import OpenAIEmbedder, OpenAIEmbedderConfig
 
 __all__ = ['OpenAIEmbedder', 'OpenAIEmbedderConfig']
 
 
-def _missing_provider(name: str, dependency: str) -> type:
-    class _MissingProvider:  # pragma: no cover - shim for optional extras
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                f'{name} requires the optional dependency "{dependency}". '
-                f'Install it via `uv sync --extra {dependency}` or the corresponding extras group.'
-            )
-
-    _MissingProvider.__name__ = f'Missing{name}'
-    return _MissingProvider
-
-
 try:  # pragma: no cover - optional dependency
     from .gemini import GeminiEmbedder, GeminiEmbedderConfig  # type: ignore[unused-import]
 except ImportError:  # pragma: no cover
-    GeminiEmbedder = _missing_provider('GeminiEmbedder', 'google-genai')  # type: ignore[assignment]
-    GeminiEmbedderConfig = _missing_provider('GeminiEmbedderConfig', 'google-genai')  # type: ignore[assignment]
+    GeminiEmbedder = missing_provider('GeminiEmbedder', 'google-genai')  # type: ignore[assignment]
+    GeminiEmbedderConfig = missing_provider('GeminiEmbedderConfig', 'google-genai')  # type: ignore[assignment]
 else:  # pragma: no cover
     __all__.extend(['GeminiEmbedder', 'GeminiEmbedderConfig'])
 
 try:  # pragma: no cover - optional dependency
     from .voyage import VoyageAIEmbedder, VoyageAIEmbedderConfig  # type: ignore[unused-import]
 except ImportError:  # pragma: no cover
-    VoyageAIEmbedder = _missing_provider('VoyageAIEmbedder', 'voyageai')  # type: ignore[assignment]
-    VoyageAIEmbedderConfig = _missing_provider('VoyageAIEmbedderConfig', 'voyageai')  # type: ignore[assignment]
+    VoyageAIEmbedder = missing_provider('VoyageAIEmbedder', 'voyageai')  # type: ignore[assignment]
+    VoyageAIEmbedderConfig = missing_provider('VoyageAIEmbedderConfig', 'voyageai')  # type: ignore[assignment]
 else:  # pragma: no cover
     __all__.extend(['VoyageAIEmbedder', 'VoyageAIEmbedderConfig'])
 
 try:  # pragma: no cover - optional dependency
     from .embeddinggemma import EmbeddingGemmaConfig, EmbeddingGemmaEmbedder  # type: ignore[unused-import]
 except ImportError:  # pragma: no cover
-    EmbeddingGemmaEmbedder = _missing_provider('EmbeddingGemmaEmbedder', 'sentence-transformers')  # type: ignore[assignment]
-    EmbeddingGemmaConfig = _missing_provider('EmbeddingGemmaConfig', 'sentence-transformers')  # type: ignore[assignment]
+    EmbeddingGemmaEmbedder = missing_provider('EmbeddingGemmaEmbedder', 'sentence-transformers')  # type: ignore[assignment]
+    EmbeddingGemmaConfig = missing_provider('EmbeddingGemmaConfig', 'sentence-transformers')  # type: ignore[assignment]
 else:  # pragma: no cover
     __all__.extend(['EmbeddingGemmaEmbedder', 'EmbeddingGemmaConfig'])

graphium_core/search/edges.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from collections.abc import Coroutine
 from typing import Any
 
@@ -24,7 +23,7 @@
     node_distance_reranker,
     rrf,
 )
-from graphium_core.utils.async_utils import semaphore_gather
+from graphium_core.search.shared import gather_search_results
 
 
 def _collect_seed_bfs_nodes(search_results: list[list[EntityEdge]]) -> list[str]:
@@ -98,36 +97,6 @@ def _build_edge_search_tasks(
     return tasks, requires_seeded_bfs
 
 
-async def _gather_edge_search_results(
-    driver: GraphDriver,
-    tasks: list[Coroutine[Any, Any, list[EntityEdge]]],
-    requires_seeded_bfs: bool,
-    config: EdgeSearchConfig,
-    search_filter: SearchFilters,
-    group_ids: list[str] | None,
-    limit: int,
-) -> list[list[EntityEdge]]:
-    search_results = list(await semaphore_gather(*tasks)) if tasks else []
-
-    if not requires_seeded_bfs:
-        return search_results
-
-    seed_node_uuids = _collect_seed_bfs_nodes(search_results)
-    if not seed_node_uuids:
-        return search_results
-
-    bfs_results = await edge_bfs_search(
-        driver,
-        seed_node_uuids,
-        config.bfs_max_depth,
-        search_filter,
-        group_ids,
-        2 * limit,
-    )
-    search_results.append(bfs_results)
-    return search_results
-
-
 def _rerank_edges_with_rrf(
     search_result_uuids: list[list[str]],
     edge_uuid_map: dict[str, EntityEdge],
@@ -311,14 +280,18 @@ async def edge_search(
         limit,
     )
 
-    search_results = await _gather_edge_search_results(
-        driver,
+    search_results = await gather_search_results(
         search_tasks,
         requires_seeded_bfs,
-        config,
-        search_filter,
-        group_ids,
-        limit,
+        _collect_seed_bfs_nodes,
+        lambda seed_node_uuids: edge_bfs_search(
+            driver,
+            seed_node_uuids,
+            config.bfs_max_depth,
+            search_filter,
+            group_ids,
+            2 * limit,
+        ),
     )
 
     reranked_edges, edge_scores = await _rerank_edges(
@@ -336,4 +309,4 @@ async def edge_search(
     return reranked_edges[:limit], edge_scores[:limit]
 
 
-__all__ = ['edge_search', '_build_edge_search_tasks', '_gather_edge_search_results']
+__all__ = ['edge_search', '_build_edge_search_tasks']

graphium_core/search/nodes.py

@@ -32,7 +32,7 @@
     node_similarity_search,
     rrf,
 )
-from graphium_core.utils.async_utils import semaphore_gather
+from graphium_core.search.shared import gather_search_results
 
 
 async def node_search(
@@ -61,14 +61,18 @@ async def node_search(
         limit,
     )
 
-    search_results = await _gather_node_search_results(
-        driver,
+    search_results = await gather_search_results(
         search_tasks,
         requires_seeded_bfs,
-        config,
-        search_filter,
-        group_ids,
-        limit,
+        _collect_seed_nodes,
+        lambda seed_nodes: node_bfs_search(
+            driver,
+            seed_nodes,
+            search_filter,
+            config.bfs_max_depth,
+            group_ids,
+            2 * limit,
+        ),
     )
 
     reranked_nodes, node_scores = await _rerank_nodes(
@@ -195,36 +199,6 @@ def _build_node_search_tasks(
     return tasks, requires_seeded_bfs
 
 
-async def _gather_node_search_results(
-    driver: GraphDriver,
-    tasks: list[Coroutine[Any, Any, list[EntityNode]]],
-    requires_seeded_bfs: bool,
-    config: NodeSearchConfig,
-    search_filter: SearchFilters,
-    group_ids: list[str] | None,
-    limit: int,
-) -> list[list[EntityNode]]:
-    search_results = list(await semaphore_gather(*tasks)) if tasks else []
-
-    if not requires_seeded_bfs:
-        return search_results
-
-    seed_nodes = _collect_seed_nodes(search_results)
-    if not seed_nodes:
-        return search_results
-
-    bfs_results = await node_bfs_search(
-        driver,
-        seed_nodes,
-        search_filter,
-        config.bfs_max_depth,
-        group_ids,
-        2 * limit,
-    )
-    search_results.append(bfs_results)
-    return search_results
-
-
 def _collect_seed_nodes(search_results: list[list[EntityNode]]) -> list[str]:
     seen: set[str] = set()
     seeds: list[str] = []

graphium_core/search/shared.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Coroutine, Sequence
+from typing import Any, TypeVar
+
+from graphium_core.utils.async_utils import semaphore_gather
+
+T = TypeVar('T')
+
+
+async def gather_search_results(
+    tasks: Sequence[Coroutine[Any, Any, list[T]]],
+    requires_seeded_bfs: bool,
+    collect_seed_ids: Callable[[list[list[T]]], list[str]],
+    bfs_fetch: Callable[[list[str]], Awaitable[list[T]]],
+) -> list[list[T]]:
+    """Execute search tasks and optionally augment results with a BFS fallback."""
+    search_results = list(await semaphore_gather(*tasks)) if tasks else []
+
+    if not requires_seeded_bfs:
+        return search_results
+
+    seed_ids = collect_seed_ids(search_results)
+    if not seed_ids:
+        return search_results
+
+    bfs_results = await bfs_fetch(seed_ids)
+    search_results.append(bfs_results)
+    return search_results
+
+
+__all__ = ['gather_search_results']

graphium_core/utils/__init__.py

@@ -0,0 +1,3 @@
+from .optional_dependencies import missing_provider
+
+__all__ = ['missing_provider']

graphium_core/utils/optional_dependencies.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import Final
+
+_DEFAULT_GUIDANCE: Final[str] = 'or the matching extras group.'
+
+
+def missing_provider(
+    name: str,
+    dependency: str,
+    *,
+    extras_group: str | None = None,
+    guidance: str | None = None,
+) -> type:
+    """Return a shim type that raises a helpful ImportError when instantiated.
+
+    Both embedder and reranker packages use the same pattern for optional providers.
+    This helper keeps the messaging consistent while avoiding copy/paste implementations.
+    """
+
+    extras_label = extras_group or dependency
+    message = (
+        f'{name} requires the optional dependency "{dependency}". '
+        f'Install it via `uv sync --extra {extras_label}` {guidance or _DEFAULT_GUIDANCE}'
+    )
+
+    class _MissingProvider:  # pragma: no cover - shim for optional extras
+        def __init__(self, *args, **kwargs):
+            raise ImportError(message)
+
+    _MissingProvider.__name__ = f'Missing{name}'
+    return _MissingProvider
+
+
+__all__ = ['missing_provider']

pyproject.toml

@@ -70,6 +70,7 @@ dev = [
     "pytest-xdist>=3.6.1",
     "ruff>=0.7.1",
     "opentelemetry-sdk>=1.20.0",
+    "pylint>=4.0.1",
 ]
 
 [build-system]