diff options
Diffstat (limited to 'rigify/utils/widgets.py')
| -rw-r--r-- | rigify/utils/widgets.py | 162 |
1 files changed, 98 insertions, 64 deletions
diff --git a/rigify/utils/widgets.py b/rigify/utils/widgets.py index 5a16065b..46300197 100644 --- a/rigify/utils/widgets.py +++ b/rigify/utils/widgets.py @@ -5,21 +5,25 @@ import math import inspect import functools +from typing import Optional, Callable +from bpy.types import Mesh, Object, UILayout from mathutils import Matrix, Vector, Euler from itertools import count from .errors import MetarigError from .collections import ensure_collection +from .misc import ArmatureObject, MeshObject, AnyVector, verify_mesh_obj from .naming import change_name_side, get_name_side, Side WGT_PREFIX = "WGT-" # Prefix for widget objects -#============================================= -# Widget creation -#============================================= +############################################## +# Widget creation +############################################## -def obj_to_bone(obj, rig, bone_name, bone_transform_name=None): +def obj_to_bone(obj: Object, rig: ArmatureObject, bone_name: str, + bone_transform_name: Optional[str] = None): """ Places an object at the location/rotation/scale of the given bone. """ if bpy.context.mode == 'EDIT_ARMATURE': @@ -45,8 +49,13 @@ def obj_to_bone(obj, rig, bone_name, bone_transform_name=None): obj.matrix_basis = rig.matrix_world @ bone.bone.matrix_local @ shape_mat -def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, widget_force_new=False, subsurf=0): - """ Creates an empty widget object for a bone, and returns the object. +def create_widget(rig: ArmatureObject, bone_name: str, + bone_transform_name: Optional[str] = None, *, + widget_name: Optional[str] = None, + widget_force_new=False, subsurf=0) -> Optional[MeshObject]: + """ + Creates an empty widget object for a bone, and returns the object. + If the object already existed, returns None. """ assert rig.mode != 'EDIT' @@ -61,16 +70,17 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, if generator: collection = generator.widget_collection else: + # noinspection SpellCheckingInspection collection = ensure_collection(bpy.context, 'WGTS_' + rig.name, hidden=True) use_mirror = generator and generator.use_mirror_widgets - - if use_mirror: - bone_mid_name = change_name_side(bone_name, Side.MIDDLE) + bone_mid_name = change_name_side(bone_name, Side.MIDDLE) if use_mirror else bone_name obj_name = widget_name or WGT_PREFIX + rig.name + '_' + bone_name reuse_mesh = None + obj: Optional[MeshObject] + # Check if it already exists in the scene if not widget_force_new: obj = None @@ -87,7 +97,9 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, # Search the scene by name obj = scene.objects.get(obj_name) if obj and obj.library: - local_objs = [obj for obj in scene.objects if obj.name == obj_name and not obj.library] + # Second brute force try if the first result is linked + local_objs = [obj for obj in scene.objects + if obj.name == obj_name and not obj.library] obj = local_objs[0] if local_objs else None if obj: @@ -100,7 +112,7 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, collection.objects.link(obj) # Flip scale for originally mirrored widgets - if obj.scale.x < 0 and bone.custom_shape_scale_xyz.x > 0: + if obj.scale.x < 0 < bone.custom_shape_scale_xyz.x: bone.custom_shape_scale_xyz.x *= -1 # Move object to bone position, in case it changed @@ -132,7 +144,7 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, mesh = bpy.data.meshes.new(obj_name) # Create the object - obj = bpy.data.objects.new(obj_name, mesh) + obj = verify_mesh_obj(bpy.data.objects.new(obj_name, mesh)) collection.objects.link(obj) # Add the subdivision surface modifier @@ -158,9 +170,9 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None, return obj -#============================================= +############################################## # Widget choice dropdown -#============================================= +############################################## _registered_widgets = {} @@ -170,7 +182,7 @@ def _get_valid_args(callback, skip): return set(spec.args[skip:] + spec.kwonlyargs) -def register_widget(name, callback, **default_args): +def register_widget(name: str, callback, **default_args): unwrapped = inspect.unwrap(callback) if unwrapped != callback: valid_args = _get_valid_args(unwrapped, 1) @@ -180,10 +192,11 @@ def register_widget(name, callback, **default_args): _registered_widgets[name] = (callback, valid_args, default_args) -def layout_widget_dropdown(layout, props, prop_name, **kwargs): - "Create a UI dropdown to select a widget from the known list." +def layout_widget_dropdown(layout: UILayout, props, prop_name: str, **kwargs): + """Create a UI dropdown to select a widget from the known list.""" id_store = bpy.context.window_manager + # noinspection PyUnresolvedReferences rigify_widgets = id_store.rigify_widgets rigify_widgets.clear() @@ -195,7 +208,7 @@ def layout_widget_dropdown(layout, props, prop_name, **kwargs): layout.prop_search(props, prop_name, id_store, "rigify_widgets", **kwargs) -def create_registered_widget(obj, bone_name, widget_id, **kwargs): +def create_registered_widget(obj: ArmatureObject, bone_name: str, widget_id: str, **kwargs): try: callback, valid_args, default_args = _registered_widgets[widget_id] except KeyError: @@ -210,26 +223,27 @@ def create_registered_widget(obj, bone_name, widget_id, **kwargs): if 'size' in valid_args and not kwargs.get('size'): kwargs['size'] = kwargs['radius'] * 2 - args = { **default_args, **kwargs } + args = {**default_args, **kwargs} - return callback(obj, bone_name, **{ k:v for k,v in args.items() if k in valid_args}) + return callback(obj, bone_name, **{k: v for k, v in args.items() if k in valid_args}) -#============================================= +############################################## # Widget geometry -#============================================= +############################################## class GeometryData: + verts: list[AnyVector] + edges: list[tuple[int, int]] + faces: list[tuple[int, ...]] + def __init__(self): self.verts = [] self.edges = [] self.faces = [] -def widget_generator(generate_func=None, *, register=None, subsurf=0): - if generate_func is None: - return functools.partial(widget_generator, register=register, subsurf=subsurf) - +def widget_generator(generate_func=None, *, register=None, subsurf=0) -> Callable: """ Decorator that encapsulates a call to create_widget, and only requires the actual function to fill the provided vertex and edge lists. @@ -237,15 +251,21 @@ def widget_generator(generate_func=None, *, register=None, subsurf=0): Accepts parameters of create_widget, plus any keyword arguments the wrapped function has. """ + if generate_func is None: + return functools.partial(widget_generator, register=register, subsurf=subsurf) + @functools.wraps(generate_func) - def wrapper(rig, bone_name, bone_transform_name=None, widget_name=None, widget_force_new=False, **kwargs): - obj = create_widget(rig, bone_name, bone_transform_name, widget_name=widget_name, widget_force_new=widget_force_new, subsurf=subsurf) + def wrapper(rig: ArmatureObject, bone_name: str, bone_transform_name=None, + widget_name=None, widget_force_new=False, **kwargs): + obj = create_widget(rig, bone_name, bone_transform_name, + widget_name=widget_name, widget_force_new=widget_force_new, + subsurf=subsurf) if obj is not None: geom = GeometryData() generate_func(geom, **kwargs) - mesh = obj.data + mesh: Mesh = obj.data mesh.from_pydata(geom.verts, geom.edges, geom.faces) mesh.update() @@ -259,7 +279,9 @@ def widget_generator(generate_func=None, *, register=None, subsurf=0): return wrapper -def generate_lines_geometry(geom, points, *, matrix=None, closed_loop=False): +def generate_lines_geometry(geom: GeometryData, + points: list[AnyVector], *, + matrix: Optional[Matrix] = None, closed_loop=False): """ Generates a polyline using given points, optionally closing the loop. """ @@ -282,13 +304,15 @@ def generate_lines_geometry(geom, points, *, matrix=None, closed_loop=False): geom.edges.append((len(geom.verts) - 1, base)) -def generate_circle_geometry(geom, center, radius, *, matrix=None, angle_range=None, - steps=24, radius_x=None, depth_x=0): +def generate_circle_geometry(geom: GeometryData, center: AnyVector, radius: float, *, + matrix: Optional[Matrix] = None, + angle_range: Optional[tuple[float, float]] = None, + steps=24, radius_x: Optional[float] = None, depth_x=0): """ Generates a circle, adding vertices and edges to the lists. center, radius: parameters of the circle matrix: transformation matrix (by default the circle is in the XY plane) - angle_range: pair of angles to generate an arc of the circle + angle_range: a pair of angles to generate an arc of the circle steps: number of edges to cover the whole circle (reduced for arcs) """ assert steps >= 3 @@ -319,7 +343,9 @@ def generate_circle_geometry(geom, center, radius, *, matrix=None, angle_range=N generate_lines_geometry(geom, points, matrix=matrix, closed_loop=not angle_range) -def generate_circle_hull_geometry(geom, points, radius, gap, *, matrix=None, steps=24): +def generate_circle_hull_geometry(geom: GeometryData, points: list[AnyVector], + radius: float, gap: float, *, + matrix: Optional[Matrix] = None, steps=24): """ Given a list of 2D points forming a convex hull, generate a contour around it, with each point being circumscribed with a circle arc of given radius, @@ -337,28 +363,28 @@ def generate_circle_hull_geometry(geom, points, radius, gap, *, matrix=None, ste base = len(geom.verts) points_ex = [points[-1], *points, points[0]] - agap = math.asin(gap / radius) + angle_gap = math.asin(gap / radius) - for i, pprev, pcur, pnext in zip(count(0), points_ex[0:], points_ex[1:], points_ex[2:]): - vprev = pprev - pcur - vnext = pnext - pcur + for i, pt_prev, pt_cur, pt_next in zip(count(0), points_ex[0:], points_ex[1:], points_ex[2:]): + vec_prev = pt_prev - pt_cur + vec_next = pt_next - pt_cur # Compute bearings to adjacent points - aprev = math.atan2(vprev.y, vprev.x) - anext = math.atan2(vnext.y, vnext.x) - if anext <= aprev: - anext += math.pi * 2 + angle_prev = math.atan2(vec_prev.y, vec_prev.x) + angle_next = math.atan2(vec_next.y, vec_next.x) + if angle_next <= angle_prev: + angle_next += math.pi * 2 # Adjust gap for circles that are too close - aprev += max(agap, math.acos(min(1, vprev.length/radius/2))) - anext -= max(agap, math.acos(min(1, vnext.length/radius/2))) + angle_prev += max(angle_gap, math.acos(min(1, vec_prev.length/radius/2))) + angle_next -= max(angle_gap, math.acos(min(1, vec_next.length/radius/2))) - if anext > aprev: + if angle_next > angle_prev: if len(geom.verts) > base: geom.edges.append((len(geom.verts)-1, len(geom.verts))) generate_circle_geometry( - geom, pcur, radius, angle_range=(aprev, anext), + geom, pt_cur, radius, angle_range=(angle_prev, angle_next), matrix=matrix, steps=steps ) @@ -366,7 +392,7 @@ def generate_circle_hull_geometry(geom, points, radius, gap, *, matrix=None, ste geom.edges.append((len(geom.verts)-1, base)) -def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0): +def create_circle_polygon(number_verts: int, axis: str, radius=1.0, head_tail=0.0): """ Creates a basic circle around of an axis selected. number_verts: number of vertices of the polygon axis: axis normal to the circle @@ -380,7 +406,7 @@ def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0): assert(axis in 'XYZ') - while i < (number_verts): + while i < number_verts: a = math.cos(i * angle) b = math.sin(i * angle) @@ -392,7 +418,7 @@ def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0): verts.append((a * radius, b * radius, head_tail)) if i < (number_verts - 1): - edges.append((i , i + 1)) + edges.append((i, i + 1)) i += 1 @@ -401,12 +427,13 @@ def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0): return verts, edges -#============================================= +############################################## # Widget transformation -#============================================= +############################################## -def adjust_widget_axis(obj, axis='y', offset=0.0): +def adjust_widget_axis(obj: Object, axis='y', offset=0.0): mesh = obj.data + assert isinstance(mesh, Mesh) if axis[0] == '-': s = -1.0 @@ -431,31 +458,35 @@ def adjust_widget_axis(obj, axis='y', offset=0.0): vert.co = matrix @ vert.co -def adjust_widget_transform_mesh(obj, matrix, local=None): +def adjust_widget_transform_mesh(obj: Optional[Object], matrix: Matrix, + local: bool | None = None): """Adjust the generated widget by applying a correction matrix to the mesh. If local is false, the matrix is in world space. If local is True, it's in the local space of the widget. If local is a bone, it's in the local space of the bone. """ if obj: + mesh = obj.data + assert isinstance(mesh, Mesh) + if local is not True: if local: assert isinstance(local, bpy.types.PoseBone) - bonemat = local.id_data.matrix_world @ local.bone.matrix_local - matrix = bonemat @ matrix @ bonemat.inverted() + bone_mat = local.id_data.matrix_world @ local.bone.matrix_local + matrix = bone_mat @ matrix @ bone_mat.inverted() - obmat = obj.matrix_basis - matrix = obmat.inverted() @ matrix @ obmat + obj_mat = obj.matrix_basis + matrix = obj_mat.inverted() @ matrix @ obj_mat - obj.data.transform(matrix) + mesh.transform(matrix) -def write_widget(obj, name='thing', use_size=True): +def write_widget(obj: Object, name='thing', use_size=True): """ Write a mesh object as a python script for widget use. """ script = "" script += "@widget_generator\n" - script += "def create_"+name+"_widget(geom"; + script += "def create_"+name+"_widget(geom" if use_size: script += ", *, size=1.0" script += "):\n" @@ -464,23 +495,26 @@ def write_widget(obj, name='thing', use_size=True): szs = "*size" if use_size else "" width = 2 if use_size else 3 + mesh = obj.data + assert isinstance(mesh, Mesh) + script += " geom.verts = [" - for i, v in enumerate(obj.data.vertices): + for i, v in enumerate(mesh.vertices): script += "({:g}{}, {:g}{}, {:g}{}),".format(v.co[0], szs, v.co[1], szs, v.co[2], szs) script += "\n " if i % width == (width - 1) else " " script += "]\n" # Edges script += " geom.edges = [" - for i, e in enumerate(obj.data.edges): + for i, e in enumerate(mesh.edges): script += "(" + str(e.vertices[0]) + ", " + str(e.vertices[1]) + ")," script += "\n " if i % 10 == 9 else " " script += "]\n" # Faces - if obj.data.polygons: + if mesh.polygons: script += " geom.faces = [" - for i, f in enumerate(obj.data.polygons): + for i, f in enumerate(mesh.polygons): script += "(" + ", ".join(str(v) for v in f.vertices) + ")," script += "\n " if i % 10 == 9 else " " script += "]\n" |