From d3b930155973e685993064031ecfc05dff3380cb Mon Sep 17 00:00:00 2001 From: Tamas Nepusz Date: Tue, 9 Apr 2024 15:59:52 +0200 Subject: [PATCH] chore: typing improvements and fixes --- .../sbstudio/plugin/model/light_effects.py | 34 ++++--- typings/bpy/__init__.pyi | 3 +- typings/bpy/path.pyi | 1 + typings/bpy/types.pyi | 88 ++++++++++++++++++- 4 files changed, 111 insertions(+), 15 deletions(-) diff --git a/src/modules/sbstudio/plugin/model/light_effects.py b/src/modules/sbstudio/plugin/model/light_effects.py index 68685ae..df95432 100644 --- a/src/modules/sbstudio/plugin/model/light_effects.py +++ b/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 diff --git a/typings/bpy/__init__.pyi b/typings/bpy/__init__.pyi index db7f19d..86fb171 100644 --- a/typings/bpy/__init__.pyi +++ b/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 diff --git a/typings/bpy/path.pyi b/typings/bpy/path.pyi index a275691..b2807bd 100644 --- a/typings/bpy/path.pyi +++ b/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: ... diff --git a/typings/bpy/types.pyi b/typings/bpy/types.pyi index 79245c4..9dde30e 100644 --- a/typings/bpy/types.pyi +++ b/typings/bpy/types.pyi @@ -1,4 +1,15 @@ -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 @@ -6,9 +17,10 @@ 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,6 +29,19 @@ 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): @@ -24,15 +49,46 @@ class ID(bpy_struct): 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]