chore: typing improvements and fixes

src/modules/sbstudio/plugin/model/light_effects.py

@@ -1,11 +1,11 @@
 import types
 import bpy
-import bmesh
 
 from functools import partial
 from operator import itemgetter
 from typing import cast, Callable, Iterable, List, Optional, Sequence, Tuple
 
+from bpy.path import abspath
 from bpy.props import (
     BoolProperty,
     CollectionProperty,
@@ -140,7 +140,7 @@ def get_color_function_names(self, context: Context) -> List[Tuple[str, str, str
     names: List[str]
 
     if self.path:
-        absolute_path = bpy.path.abspath(self.path)
+        absolute_path = abspath(self.path)
         module = load_module(absolute_path)
         names = [
             name
@@ -374,13 +374,11 @@ def apply_on_colors(
                 randomization is turned on
         """
 
-        time_fraction = (frame - self.frame_start) / self.duration
-
         def get_output_based_on_output_type(
             output_type: str,
             mapping_mode: str,
             output_function,
-        ) -> Tuple[Optional[List[float]], Optional[float]]:
+        ) -> Tuple[Optional[List[Optional[float]]], Optional[float]]:
             """Get the float output(s) for color ramp or image indexing based on the output type.
 
             Args:
@@ -390,7 +388,7 @@ def get_output_based_on_output_type(
             Returns:
                 individual and common outputs
             """
-            outputs: Optional[List[float]] = None
+            outputs: Optional[List[Optional[float]]] = None
             common_output: Optional[float] = None
             order: Optional[List[int]] = None
 
@@ -437,7 +435,7 @@ def get_output_based_on_output_type(
                         # axes
                         sort_key = lambda index: query_axes(positions[index])
 
-                outputs = [1.0] * num_positions
+                outputs = [1.0] * num_positions  # type: ignore
                 order = list(range(num_positions))
                 if num_positions > 1:
                     if proportional and sort_key is not None:
@@ -459,6 +457,7 @@ def get_output_based_on_output_type(
                         if sort_key is not None:
                             order.sort(key=sort_key)
 
+                        assert outputs is not None
                         for u, v in enumerate(order):
                             outputs[v] = u / (num_positions - 1)
 
@@ -496,9 +495,9 @@ def get_output_based_on_output_type(
                         outputs = [None if x is None else x / np_m1 for x in mapping]
                 else:
                     # if there is no mapping at all, we do not change color of drones
-                    outputs = [None] * num_positions
+                    outputs = [None] * num_positions  # type: ignore
             elif output_type == "CUSTOM":
-                absolute_path = bpy.path.abspath(output_function.path)
+                absolute_path = abspath(output_function.path)
                 module = load_module(absolute_path)
                 if self.output_function.name:
                     fn = getattr(module, self.output_function.name)
@@ -507,7 +506,9 @@ def get_output_based_on_output_type(
                             frame=frame,
                             time_fraction=time_fraction,
                             drone_index=index,
-                            formation_index=mapping[index],
+                            formation_index=(
+                                mapping[index] if mapping is not None else None
+                            ),
                             position=positions[index],
                             drone_count=num_positions,
                         )
@@ -525,6 +526,7 @@ def get_output_based_on_output_type(
         if not self.enabled or not self.contains_frame(frame):
             return
 
+        time_fraction = (frame - self.frame_start) / self.duration
         num_positions = len(positions)
 
         color_ramp = self.color_ramp
@@ -559,6 +561,8 @@ def get_output_based_on_output_type(
                 if outputs_x[index] is None:
                     continue
                 output_x = outputs_x[index]
+            assert isinstance(output_x, float)
+
             if color_image is not None:
                 if common_output_y is not None:
                     output_y = common_output_y
@@ -569,6 +573,7 @@ def get_output_based_on_output_type(
                     if outputs_y[index] is None:
                         continue
                     output_y = outputs_y[index]
+                assert isinstance(output_y, float)
 
             # Randomize the output value if needed
             if self.randomness != 0:
@@ -590,7 +595,9 @@ def get_output_based_on_output_type(
                         frame=frame,
                         time_fraction=time_fraction,
                         drone_index=index,
-                        formation_index=mapping[index],
+                        formation_index=(
+                            mapping[index] if mapping is not None else None
+                        ),
                         position=position,
                         drone_count=num_positions,
                     )
@@ -655,7 +662,7 @@ def color_image(self, image):
     def color_function_ref(self) -> Optional[Callable]:
         if self.type != "FUNCTION" or not self.color_function:
             return None
-        absolute_path = bpy.path.abspath(self.color_function.path)
+        absolute_path = abspath(self.color_function.path)
         module = load_module(absolute_path)
         return getattr(module, self.color_function.name, None)
 
@@ -717,6 +724,7 @@ def update_from(self, other: "LightEffect") -> None:
         self.output_function_y.copy_to(other.output_function_y)
 
         if self.color_ramp is not None:
+            assert other.color_ramp is not None  # because we copied the type
             update_color_ramp_from(self.color_ramp, other.color_ramp)
 
     def _evaluate_influence_at(
@@ -881,6 +889,8 @@ def append_new_entry(
                 sensible default
             select: whether to select the newly added entry after it was created
         """
+        assert context is not None
+
         scene = context.scene
 
         fps = scene.render.fps

typings/bpy/__init__.pyi

@@ -1,6 +1,6 @@
 from typing import Tuple
 
-from .types import Context
+from .types import BlendData, Context
 
 class _App:
     version: Tuple[int, int, int]
@@ -9,3 +9,4 @@ class _App:
 
 app: _App
 context: Context
+data: BlendData

typings/bpy/path.pyi

@@ -1,4 +1,5 @@
 from pathlib import Path
 from typing import Union
 
+def abspath(path: Union[str, Path]) -> str: ...
 def basename(path: Union[str, Path]) -> str: ...

typings/bpy/types.pyi

@@ -1,14 +1,26 @@
-from typing import Generic, Optional, TypeVar, Union, overload
+from __future__ import annotations
+
+from typing import (
+    Generic,
+    Literal,
+    MutableSequence,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from mathutils import Matrix
 from sbstudio.plugin.model import DroneShowAddonProperties
 
 T = TypeVar("T")
 U = TypeVar("U")
 
+RGBAColor = MutableSequence[float]
 Vector3 = tuple[float, float, float]
 
-class bpy_prop_collection(Generic[T]):
+class bpy_prop_collection(Generic[T], Sequence[T]):
     def find(self, key: str) -> int: ...
     @overload
     def get(self, key: str) -> Optional[T]: ...
@@ -17,22 +29,66 @@ class bpy_prop_collection(Generic[T]):
 
 class bpy_struct: ...
 
+class ColorRampElement(bpy_struct):
+    alpha: float
+    color: RGBAColor
+    position: float
+
+class ColorRamp(bpy_struct):
+    color_mode: Literal["RGB", "HSV", "HSL"]
+    elements: bpy_prop_collection[ColorRampElement]
+    hue_interpolation: Literal["NEAR", "FAR", "CW", "CCW"]
+    interpolation: Literal["EASE", "CARDINAL", "LINEAR", "B_SPLINE", "CONSTANT"]
+
+    def evaluate(self, position: float) -> RGBAColor: ...
+
 Self = TypeVar("Self", bound="ID")
 
 class ID(bpy_struct):
     name: str
 
     def copy(self: Self) -> Self: ...
 
+class PropertyGroup(bpy_struct):
+    name: str
+
+class RenderSettings(bpy_struct):
+    fps: int
+    fps_base: float
+
+class Image(ID):
+    depth: int
+    size: tuple[int, int]
+
+class Material(ID): ...
+
 class Mesh(ID):
     def transform(self, matrix: Matrix, shape_keys: bool = False) -> None: ...
 
+class Texture(ID):
+    color_ramp: ColorRamp
+    use_color_ramp: bool
+
+class ImageTexture(Texture):
+    image: Image
+
 class Depsgraph(bpy_struct):
     objects: bpy_prop_collection[Object]
     scene: Scene
     scene_eval: Scene
 
 class Scene:
+    frame_current: int
+    frame_current_final: float
+    frame_end: int
+    frame_float: float
+    frame_preview_end: int
+    frame_preview_start: int
+    frame_start: int
+    frame_step: int
+    frame_subframe: float
+
+    render: RenderSettings
     skybrush: DroneShowAddonProperties
 
 class Context(bpy_struct):
@@ -49,3 +105,31 @@ class Object(ID):
     matrix_local: Matrix
     matrix_parent_inverse: Matrix
     matrix_world: Matrix
+
+class BlendDataImage(bpy_prop_collection[Image]):
+    def new(
+        self,
+        name: str,
+        width: int,
+        height: int,
+        alpha=False,
+        float_buffer=False,
+        stereo3d=False,
+        is_data=False,
+        tiled=False,
+    ) -> Image: ...
+
+class BlendDataTextures(bpy_prop_collection[Texture]):
+    @overload
+    def new(self, name: str, type: Literal["IMAGE"]) -> ImageTexture: ...
+    @overload
+    def new(self, name: str, type: str) -> Texture: ...
+
+class BlendData(bpy_struct):
+    images: BlendDataImage
+    materials: bpy_prop_collection[Material]
+    meshes: bpy_prop_collection[Mesh]
+    objects: bpy_prop_collection[Object]
+    scenes: bpy_prop_collection[Scene]
+    textures: BlendDataTextures
+    version: tuple[int, int, int]