Various typing fixes (#320)

fabnemEPFL · JCHAVEROT · web-flow · commit 570b013cf5b1 · 2026-06-12T14:17:28.000+02:00
Co-authored-by: Jérémy Chaverot &lt;jeremy.chaverot@epfl.ch&gt;
diff --git a/.github/pyright/pyrightconfig.colvision.json b/.github/pyright/pyrightconfig.colvision.json
@@ -0,0 +1,16 @@
+{
+  "include": ["../../src/mmore/colvision", "../../src/mmore/tui"],
+  "exclude": ["../../build", "../../dist", "**/__pycache__", "*.egg-info"],
+  "executionEnvironments": [
+    {
+      "root": "../../src/mmore/colvision",
+      "extraPaths": ["../../src"]
+    },
+    {
+      "root": "../../src/mmore/tui",
+      "extraPaths": ["../../src"]
+    }
+  ],
+  "pythonVersion": "3.11",
+  "typeCheckingMode": "standard"
+}
diff --git a/.github/workflows/pyright.yml b/.github/workflows/pyright.yml
@@ -8,39 +8,29 @@ on:
 
 jobs:
   pyright:
+    name: "Pyright"
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v6
 
-      - name: Set up Python 3.12
+      - name: Set up Python 3.11
         uses: actions/setup-python@v6
         with:
-          python-version: "3.12"
+          python-version: "3.11"
 
-      - name: Install uv and create venv
-        run: |
-          pipx install uv
-          uv venv .venv
+      - name: Install uv
+        run: pipx install uv
 
-      - name: Install dependencies
-        run: |
-          source .venv/bin/activate
-          uv pip install -e ".[process,index,rag,api,cpu,dev,websearch]"
+      - name: Install base dependencies
+        run: uv venv .venv && source .venv/bin/activate && uv pip install -e ".[process,index,rag,api,tui,cpu,dev,websearch]"
 
-      - name: Run Pyright
+      - name: Run Pyright (base)
         continue-on-error: true
-        run: |
-          source .venv/bin/activate
-          pyright
+        run: source .venv/bin/activate && pyright --project pyrightconfig.json
 
-      - name: Install dependencies - colvision
-        run: |
-          uv venv .venv-colvision
-          source .venv-colvision/bin/activate
-          uv pip install -e ".[colvision,cpu,dev]"
+      - name: Install colvision dependencies
+        run: uv venv .venv-colvision && source .venv-colvision/bin/activate && uv pip install -e ".[colvision,tui,cpu,dev]"
 
-      - name: Run Pyright - colvision
+      - name: Run Pyright (colvision)
         continue-on-error: true
-        run: |
-          source .venv-colvision/bin/activate
-          pyright src/mmore/colvision
+        run: source .venv-colvision/bin/activate && pyright --project .github/pyright/pyrightconfig.colvision.json
diff --git a/README.md b/README.md
@@ -160,7 +160,7 @@ MultimodalSample.to_jsonl(out_file, result_pdf)
 
 To launch the MMORE pipeline, follow the specialised instructions in the docs.
 
-![The MMORE pipelines architecture](https://github.com/user-attachments/assets/0cd61466-1680-43ed-9d55-7bd483a04a09)
+![The MMORE pipelines architecture](docs/source/doc_images/pipeline_mmore+.png)
 
 
 1. **:page_facing_up: Input Documents**
diff --git a/docs/source/advanced_usage/profiler.md b/docs/source/advanced_usage/profiler.md
@@ -119,7 +119,7 @@ If you only need to measure execution time (wall clock), use `time_function` or
 ```python
 from mmore.profiler import time_function, time_context
 
-@time_function
+@time_function()
 def quick_check():
     pass
 
diff --git a/docs/source/doc_images/pipeline_mmore+.png b/docs/source/doc_images/pipeline_mmore+.png
diff --git a/pyrightconfig.json b/pyrightconfig.json
@@ -0,0 +1,22 @@
+{
+  "include": ["src/mmore", "tests"],
+  "exclude": [
+    "build",
+    "dist",
+    "**/__pycache__",
+    "*.egg-info",
+    "src/mmore/colvision"
+  ],
+  "executionEnvironments": [
+    {
+      "root": "src/mmore",
+      "extraPaths": ["src"]
+    },
+    {
+      "root": "tests",
+      "extraPaths": ["src"]
+    }
+  ],
+    "pythonVersion": "3.11",
+    "typeCheckingMode": "standard"
+}
diff --git a/src/mmore/colvision/model_utils.py b/src/mmore/colvision/model_utils.py
@@ -57,7 +57,12 @@ def resolve_model_classes(model_name: str) -> Tuple[Type, Type]:
     Raises:
         ValueError: if the model name does not match any known pattern.
     """
-    import colpali_engine.models as models_module
+    try:
+        import colpali_engine.models as models_module
+    except ImportError as e:
+        raise ImportError(
+            "colpali_engine is required for ColVision. Install with: pip install mmore[colvision]"
+        ) from e
 
     name_lower = model_name.lower()
 
diff --git a/src/mmore/colvision/retriever.py b/src/mmore/colvision/retriever.py
@@ -59,7 +59,7 @@ def embed_queries(texts: List[str], model, processor) -> List[np.ndarray]:
         with torch.no_grad():
             batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
             emb = model(**batch_query)
-            vectors.extend(list(torch.unbind(emb.to("cpu"))))
+            vectors.extend(list(emb.to("cpu").unbind()))
     return [v.float().numpy() for v in vectors]
 
 
diff --git a/src/mmore/process/post_processor/__init__.py b/src/mmore/process/post_processor/__init__.py
@@ -33,13 +33,15 @@ def load_postprocessor(config: BasePostProcessorConfig) -> BasePostProcessor:
         from .translator import TranslatorConfig, TranslatorPostProcessor
 
         config_translator = load_config(config.args, TranslatorConfig)
-        return TranslatorPostProcessor.from_config(config_translator)
+        return cast(
+            BasePostProcessor, TranslatorPostProcessor.from_config(config_translator)
+        )
 
     elif config.type == "metafuse":
         from .metafuse import MetaDataInfusor, MetaDataInfusorConfig
 
         config_metafuse = load_config(config.args, MetaDataInfusorConfig)
-        return MetaDataInfusor.from_config(config_metafuse)
+        return cast(BasePostProcessor, MetaDataInfusor.from_config(config_metafuse))
 
     else:
         raise ValueError(f"Unrecognized postprocessor type: {config.type}")
diff --git a/src/mmore/process/processors/base.py b/src/mmore/process/processors/base.py
@@ -216,12 +216,13 @@ def __getstate__(self):
             del state["_pool"]
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: Dict[str, Any]):
         """
         Called when the object is unpickled (received by the worker).
         We restore the state and set _pool to None (workers don't need the pool manager).
         """
-        self.__dict__.update(state)
+        for key, value in state.items():
+            setattr(self, key, value)
         # Initialize _pool as None in the worker process
         self._pool = None
         # Workers should never own the pool
diff --git a/src/mmore/process/processors/media_processor.py b/src/mmore/process/processors/media_processor.py
@@ -3,10 +3,12 @@
 import tempfile
 from typing import List
 
+import numpy as np
 import torch
 from moviepy.audio.io.AudioFileClip import AudioFileClip
 from moviepy.video.io.VideoFileClip import VideoFileClip
 from PIL import Image
+from torch._C import device as torch_device
 from transformers.pipelines import pipeline as pipeline_t
 
 from ...type import DocumentMetadata, FileDescriptor, MultimodalSample
@@ -19,10 +21,10 @@ class MediaProcessor(Processor):
     @staticmethod
     def _get_available_devices():
         if torch.cuda.is_available():
-            return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
+            return [torch_device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
         if torch.backends.mps.is_available():
-            return [torch.device("mps")]
-        return [torch.device("cpu")]
+            return [torch_device("mps")]
+        return [torch_device("cpu")]
 
     devices = _get_available_devices()
     pipelines = []
@@ -158,7 +160,7 @@ def _extract_video_frames(file_path: str) -> List[Image.Image]:
                     num_thumbnails = max(1, int(duration / sample_rate))
                     for i in range(num_thumbnails):
                         t = min(i * sample_rate, duration - 0.1)
-                        frame = clip.get_frame(t)
+                        frame = np.asarray(clip.get_frame(t))
                         image = Image.fromarray(frame).convert("RGB")
                         images.append(image)
                 logger.info(f"Extracted {len(images)} images from {file_path}.")
diff --git a/src/mmore/process/processors/pdf_processor.py b/src/mmore/process/processors/pdf_processor.py
@@ -231,7 +231,7 @@ def _extract_images(pdf_doc, xref) -> Optional[Image.Image]:
                 if image_bytes is None:
                     logging.error(f"No image data found for xref {xref}")
 
-                return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+                return Image.open(io.BytesIO(cast(bytes, image_bytes))).convert("RGB")
 
             except KeyError as e:
                 logging.error(f"KeyError while extracting image: {e}")
@@ -249,7 +249,7 @@ def _extract_images(pdf_doc, xref) -> Optional[Image.Image]:
                 )
                 return None
 
-        for page_num, page in enumerate(pdf_doc):
+        for page_num, page in enumerate(pdf_doc):  # pyright: ignore[reportArgumentType]
             text = clean_text(page.get_text())  # type: ignore[attr-defined]
 
             if text.strip():
diff --git a/src/mmore/profiler.py b/src/mmore/profiler.py
@@ -10,10 +10,12 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Callable, Optional, TypeVar
 
 logger = logging.getLogger(__name__)
 
+R = TypeVar("R")
+
 
 @dataclass
 class ProfilingConfig:
@@ -204,20 +206,19 @@ def time_context(name: str, log: bool = True):
             logger.info(f"⏱️  {name} took {elapsed:.4f} seconds")
 
 
-def time_function(func: Optional[Callable] = None, log: bool = True):
+def time_function(log: bool = True) -> Callable[[Callable[..., R]], Callable[..., R]]:
     """Decorator to time a function execution.
 
     Args:
-        func: Function to decorate (if used as decorator without parentheses)
         log: Whether to log the timing result
 
     Returns:
         Decorated function
     """
 
-    def decorator(f: Callable) -> Callable:
+    def decorator(f: Callable[..., R]) -> Callable[..., R]:
         @functools.wraps(f)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: object, **kwargs: object) -> R:
             start_time = time.time()
             try:
                 result = f(*args, **kwargs)
@@ -229,10 +230,7 @@ def wrapper(*args, **kwargs):
 
         return wrapper
 
-    if func is None:
-        return decorator
-    else:
-        return decorator(func)
+    return decorator
 
 
 class Profiler:
diff --git a/src/mmore/rag/model/dense/multimodal.py b/src/mmore/rag/model/dense/multimodal.py
@@ -14,8 +14,9 @@
 class MultimodalEmbeddings(Embeddings):
     def __init__(self, model_name: str):
         super().__init__()
+        float16: torch.dtype = torch.float16
         self.model = AutoModelForImageTextToText.from_pretrained(
-            model_name, torch_dtype=torch.float16, device_map="auto"
+            model_name, torch_dtype=float16, device_map="auto"
         )
         self.processor = AutoProcessor.from_pretrained(model_name)
         self.device = self.model.device
diff --git a/src/mmore/run_ragcli.py b/src/mmore/run_ragcli.py
@@ -14,10 +14,9 @@
 )
 
 from mmore.profiler import enable_profiling_from_env, profile_function
-
-from .rag.pipeline import RAGPipeline
-from .run_rag import RAGInferenceConfig
-from .utils import load_config
+from mmore.rag.pipeline import RAGPipeline
+from mmore.run_rag import RAGInferenceConfig
+from mmore.utils import load_config
 
 
 class RagCLI:
diff --git a/src/mmore/run_websearch.py b/src/mmore/run_websearch.py
@@ -16,13 +16,12 @@
 from pydantic import BaseModel, Field
 
 from mmore.profiler import enable_profiling_from_env, profile_function
-
-from .rag.pipeline import RAGPipeline
-from .run_rag import APIConfig, LocalConfig, RAGInferenceConfig
-from .utils import load_config
-from .websearchRAG.config import WebsearchConfig
-from .websearchRAG.logging_config import logger
-from .websearchRAG.pipeline import WebsearchPipeline
+from mmore.rag.pipeline import RAGPipeline
+from mmore.run_rag import APIConfig, LocalConfig, RAGInferenceConfig
+from mmore.utils import load_config
+from mmore.websearchRAG.config import WebsearchConfig
+from mmore.websearchRAG.logging_config import logger
+from mmore.websearchRAG.pipeline import WebsearchPipeline
 
 load_dotenv()
 
diff --git a/src/mmore/tui/config_builder.py b/src/mmore/tui/config_builder.py
@@ -802,7 +802,7 @@ def pick_or_build_config(
                     )
                 )
                 path = picked
-                if choice == "edit":
+                if choice == "edit" and path is not None:
                     _edit_config(path)
 
         if choice == "manual":
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -4,7 +4,7 @@
 import pytest
 from langchain_milvus.utils.sparse import BaseSparseEmbedding
 
-from mmore.type import MultimodalSample
+from mmore.type import DocumentMetadata, MultimodalSample
 
 
 class FakeSparseEmbedding(BaseSparseEmbedding):
@@ -23,21 +23,21 @@ def embed_documents(self, texts: List[str]) -> List[Dict[int, float]]:
         document_id="doc-1",
         text="Paris is the capital of France.",
         modalities=[],
-        metadata={},
+        metadata=DocumentMetadata(),
     ),
     MultimodalSample(
         id="doc-2",
         document_id="doc-2",
         text="The Eiffel Tower stands 330 metres tall.",
         modalities=[],
-        metadata={"author": "Alice"},
+        metadata=DocumentMetadata(extra={"author": "Alice"}),
     ),
     MultimodalSample(
         id="doc-3",
         document_id="doc-3",
         text="Milvus is an open-source vector database.",
         modalities=[],
-        metadata={},
+        metadata=DocumentMetadata(),
     ),
 ]
 
diff --git a/tests/test_indexer.py b/tests/test_indexer.py
@@ -11,7 +11,7 @@
 from mmore.index.indexer import Indexer
 from mmore.rag.model import DenseModelConfig, SparseModelConfig
 from mmore.run_index import index
-from mmore.type import MultimodalSample
+from mmore.type import DocumentMetadata, MultimodalSample
 
 
 @pytest.fixture
@@ -98,7 +98,7 @@ def test_indexer_idempotent_collection(mock_sparse, tmp_path):
             document_id="doc-4",
             text="A fourth document.",
             modalities=[],
-            metadata={},
+            metadata=DocumentMetadata(),
         )
     ]
 
@@ -153,7 +153,7 @@ def test_indexer_error_on_missing_collection(mock_sparse, tmp_path):
     )
 
     doc = MultimodalSample(
-        id="x", document_id="x", text="test", modalities=[], metadata={}
+        id="x", document_id="x", text="test", modalities=[], metadata=DocumentMetadata()
     )
     with pytest.raises(Exception):
         # Bypasses index_documents (which creates the collection) and inserts directly
diff --git a/tests/test_profiler.py b/tests/test_profiler.py
