diff --git a/src/modules/sbstudio/plugin/objects.py b/src/modules/sbstudio/plugin/objects.py index 4529e30..0f02e36 100644 --- a/src/modules/sbstudio/plugin/objects.py +++ b/src/modules/sbstudio/plugin/objects.py @@ -1,8 +1,8 @@ import bpy -from bpy.types import Collection, Context, MeshVertex, Object, Scene, VertexGroup +from bpy.types import Collection, Context, Mesh, MeshVertex, Object, Scene, VertexGroup from mathutils import Vector -from typing import Any, Iterable, List, Optional, Union, Tuple +from typing import Any, Iterable, List, Optional, Union, Tuple, cast from sbstudio.model.types import Coordinate3D @@ -80,9 +80,10 @@ def get_vertices_of_object_in_vertex_group( object: the object to query name: the name of the vertex group """ - result = [] + result: list[MeshVertex] = [] mesh = object.data if object else None if mesh is not None: + mesh = cast(Mesh, mesh) index = group.index for vertex in mesh.vertices: if any(g.group == index for g in vertex.groups): diff --git a/src/modules/sbstudio/plugin/utils/transition.py b/src/modules/sbstudio/plugin/utils/transition.py index 36a096f..f05571b 100644 --- a/src/modules/sbstudio/plugin/utils/transition.py +++ b/src/modules/sbstudio/plugin/utils/transition.py @@ -1,5 +1,15 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, cast + +from bpy.types import Constraint, CopyLocationConstraint, Object + from .identifiers import create_internal_id, is_internal_id +if TYPE_CHECKING: + from sbstudio.plugin.model import StoryboardEntry + + __all__ = ( "create_transition_constraint_between", "find_transition_constraint_between", @@ -9,14 +19,16 @@ ) -def get_id_for_formation_constraint(storyboard_entry): +def get_id_for_formation_constraint(storyboard_entry: StoryboardEntry): """Returns a unique identifier for the given storyboard entry.""" # Make sure to update is_transition_constraint() as well if you change the # format of the ID return create_internal_id(f"Entry {storyboard_entry.id}") -def create_transition_constraint_between(drone, storyboard_entry): +def create_transition_constraint_between( + drone: Object, storyboard_entry: StoryboardEntry +) -> CopyLocationConstraint: """Creates a transition constraint between the given drone and an arbitrary entry in the storyboard. @@ -30,10 +42,12 @@ def create_transition_constraint_between(drone, storyboard_entry): constraint.name = get_id_for_formation_constraint(storyboard_entry) constraint.influence = 0 - return constraint + return cast(CopyLocationConstraint, constraint) -def find_transition_constraint_between(drone, storyboard_entry): +def find_transition_constraint_between( + drone: Object, storyboard_entry: StoryboardEntry +) -> Optional[CopyLocationConstraint]: """Finds the Blender "copy location" constraint object that exists between the given drone and any of the points of the formation of the given storyboard entry; the purpose of this constraint is to keep the drone at @@ -47,12 +61,12 @@ def find_transition_constraint_between(drone, storyboard_entry): for constraint in drone.constraints: if constraint.type == "COPY_LOCATION" and constraint.name == expected_id: - return constraint + return cast(CopyLocationConstraint, constraint) return None -def is_transition_constraint(constraint): +def is_transition_constraint(constraint: Constraint) -> bool: """Returns whether the given constraint object is a transition constraint, judging from its name and type. """ @@ -64,7 +78,9 @@ def is_transition_constraint(constraint): ) -def set_constraint_name_from_storyboard_entry(constraint, storyboard_entry): +def set_constraint_name_from_storyboard_entry( + constraint: Constraint, storyboard_entry: StoryboardEntry +) -> None: """Updates the name of the given constraint such that it refers to the given storyboard entry. """ diff --git a/typings/bpy/types.pyi b/typings/bpy/types.pyi index c3c4e21..b8f76bc 100644 --- a/typings/bpy/types.pyi +++ b/typings/bpy/types.pyi @@ -25,6 +25,8 @@ class bpy_prop_collection(Sequence[T]): def get(self, key: str) -> Optional[T]: ... @overload def get(self, key: str, default: U) -> Union[T, U]: ... + def __getitem__(self, key: Union[int, str]) -> T: ... + def __len__(self) -> int: ... class bpy_struct: ... @@ -41,6 +43,44 @@ class ColorRamp(bpy_struct): def evaluate(self, position: float) -> RGBAColor: ... +class Constraint(bpy_struct): + name: str + influence: float + type: Literal[ + "CAMERA_SOLVER", + "FOLLOW_TRACK", + "OBJECT_SOLVER", + "COPY_LOCATION", + "COPY_ROTATION", + "COPY_SCALE", + "COPY_TRANSFORMS", + "LIMIT_DISTANCE", + "LIMIT_LOCATION", + "LIMIT_ROTATION", + "LIMIT_SCALE", + "MAINTAIN_VOLUME", + "TRANSFORM", + "TRANSFORM_CACHE", + "TRACK_TO", + "DAMPED_TRACK", + "IK", + "LOCKED_TRACK", + "SPLINE_IK", + "STRETCH_TO", + "TRACK_TO", + "ACTION", + "ARMATURE", + "CHILD_OF", + "FLOOR", + "FOLLOW_PATH", + "PIVOT", + "SHRINKWRAP", + ] + +class CopyLocationConstraint(Constraint): + target: Object + subtarget: str + Self = TypeVar("Self", bound="ID") class ID(bpy_struct): @@ -49,7 +89,10 @@ class ID(bpy_struct): def copy(self: Self) -> Self: ... -class MeshVertex(bpy_struct): ... +class MeshVertex(bpy_struct): + groups: bpy_prop_collection[VertexGroupElement] + index: int + select: bool class PropertyGroup(bpy_struct): name: str @@ -58,6 +101,15 @@ class RenderSettings(bpy_struct): fps: int fps_base: float +class VertexGroup(bpy_struct): + index: int + name: str + +class VertexGroupElement(bpy_struct): + group: int + weight: float + +class VertexGroups(bpy_prop_collection[VertexGroup]): ... class ViewLayer(bpy_struct): ... class Collection(ID): @@ -119,6 +171,8 @@ class Context(bpy_struct): class Object(ID): active_material: Material + constraints: ObjectConstraints + vertex_groups: VertexGroups data: ID location: Vector3 @@ -146,6 +200,9 @@ class CollectionObjects(bpy_prop_collection[Object]): def link(self, object: Object) -> None: ... def unlink(self, object: Object) -> None: ... +class ObjectConstraints(bpy_prop_collection[Constraint]): + def new(self, type: str) -> Constraint: ... + class BlendDataCollections(bpy_prop_collection[Collection]): def new(self, name: str) -> Collection: ... def remove(