From 948f7c95a41c11316e36b540777e9921e264b9fa Mon Sep 17 00:00:00 2001
From: Alexander Druz <alexandersdruz@gmail.com>
Date: Fri, 10 May 2024 17:07:46 +0200
Subject: [PATCH] Satisfy mypy

---
 pyproject.toml                                 |  1 -
 renumics/spotlight/embeddings/preprocessors.py |  6 ++++--
 renumics/spotlight/io/audio.py                 | 18 ++++++++++--------
 renumics/spotlight/media/mesh.py               |  5 +++--
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ec19c059..c2a306c1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -162,7 +162,6 @@ plugins = "pydantic.mypy"
 [[tool.mypy.overrides]]
 module = [
   "pygltflib",
-  "av",
   "validators",
   "h5py",
   "pandas",
diff --git a/renumics/spotlight/embeddings/preprocessors.py b/renumics/spotlight/embeddings/preprocessors.py
index 8d89a220..00397577 100644
--- a/renumics/spotlight/embeddings/preprocessors.py
+++ b/renumics/spotlight/embeddings/preprocessors.py
@@ -4,9 +4,10 @@
 """
 
 import io
-from typing import List
+from typing import List, cast
 
 import av
+import av.audio
 import numpy as np
 import PIL.Image
 
@@ -36,10 +37,11 @@ def preprocess_audio_batch(
     for raw_data in raw_values:
         with av.open(io.BytesIO(raw_data), "r") as container:
             resampler = av.AudioResampler(
-                format="dbl", layout="mono", rate=sampling_rate
+                format="dbl", layout="mono", rate=sampling_rate  # type: ignore
             )
             data = []
             for frame in container.decode(audio=0):
+                frame = cast(av.audio.AudioFrame, frame)
                 resampled_frames = resampler.resample(frame)
                 for resampled_frame in resampled_frames:
                     frame_array = resampled_frame.to_ndarray()[0]
diff --git a/renumics/spotlight/io/audio.py b/renumics/spotlight/io/audio.py
index 32fe54c9..367c9f07 100644
--- a/renumics/spotlight/io/audio.py
+++ b/renumics/spotlight/io/audio.py
@@ -4,9 +4,10 @@
 
 import io
 import os
-from typing import IO, Dict, Tuple, Union
+from typing import IO, Dict, Tuple, Union, cast
 
 import av
+import av.audio
 import numpy as np
 import requests
 import validators
@@ -82,6 +83,7 @@ def read_audio(file: FileType) -> Tuple[np.ndarray, int]:
 
         data = []
         for frame in container.decode(audio=0):
+            frame = cast(av.audio.AudioFrame, frame)
             frame_array = frame.to_ndarray()
             if len(frame_array) == 1:
                 frame_array = frame_array.reshape((-1, num_channels))
@@ -123,13 +125,13 @@ def write_audio(
     # `AudioFrame.from_ndarray` expects an C-contiguous array as input.
     data = np.ascontiguousarray(data)
     num_channels = len(data)
-    frame = av.audio.AudioFrame.from_ndarray(data, data_format, num_channels)
+    frame = av.audio.AudioFrame.from_ndarray(data, data_format, num_channels)  # type: ignore
     frame.rate = sampling_rate
     with av.open(file, "w", format_) as container:
         stream = container.add_stream(codec, sampling_rate)
-        stream.channels = num_channels
-        container.mux(stream.encode(frame))
-        container.mux(stream.encode(None))
+        stream.channels = num_channels  # type: ignore
+        container.mux(stream.encode(frame))  # type: ignore
+        container.mux(stream.encode(None))  # type: ignore
 
 
 def transcode_audio(
@@ -155,10 +157,10 @@ def transcode_audio(
 
             for frame in input_container.decode(input_stream):
                 frame.pts = None
-                for packet in output_stream.encode(frame):
+                for packet in output_stream.encode(frame):  # type: ignore
                     output_container.mux(packet)
 
-            for packet in output_stream.encode(None):
+            for packet in output_stream.encode(None):  # type: ignore
                 output_container.mux(packet)
 
 
@@ -169,7 +171,7 @@ def get_format_codec(file: FileType) -> Tuple[str, str]:
     file = prepare_input_file(file)
     with av.open(file, "r") as input_container:
         stream = input_container.streams.audio[0]
-        return input_container.format.name, stream.name
+    return input_container.format.name, stream.name  # type: ignore
 
 
 def get_waveform(file: FileType) -> np.ndarray:
diff --git a/renumics/spotlight/media/mesh.py b/renumics/spotlight/media/mesh.py
index 6968b047..c7af1314 100644
--- a/renumics/spotlight/media/mesh.py
+++ b/renumics/spotlight/media/mesh.py
@@ -1,7 +1,7 @@
 import io
 import math
 import os
-from typing import IO, Dict, List, Optional, Tuple, Union
+from typing import IO, Dict, List, Optional, Tuple, Union, cast
 from urllib.parse import urlparse
 
 import numpy as np
@@ -122,7 +122,7 @@ def from_trimesh(cls, mesh: trimesh.Trimesh) -> "Mesh":
         Import a `trimesh.Trimesh` mesh.
         """
         return cls(
-            mesh.vertices, mesh.faces, mesh.vertex_attributes, mesh.face_attributes
+            mesh.vertices, mesh.faces, mesh.vertex_attributes, mesh.face_attributes  # type: ignore
         )
 
     @classmethod
@@ -155,6 +155,7 @@ def from_file(cls, filepath: PathType) -> "Mesh":
             raise exceptions.InvalidFile(
                 f"Mesh {filepath} does not exist or could not be read."
             ) from e
+        mesh = cast(trimesh.Trimesh, mesh)
         return cls.from_trimesh(mesh)
 
     @classmethod