object_color_rules.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 bl_info = {
   6     "name": "Object Color Rules",
   7     "author": "Campbell Barton",
   8     "version": (0, 0, 2),
   9     "blender": (2, 80, 0),
  10     "location": "Properties > Object Buttons",
  11     "description": "Rules for assigning object color (for object & wireframe colors).",
  12     "doc_url": "{BLENDER_MANUAL_URL}/addons/object/color_rules.html",
  13     "category": "Object",
  14 }
  15
  16
  17 def test_name(rule, needle, haystack, cache):
  18     if rule.use_match_regex:
  19         if not cache:
  20             import re
  21             re_needle = re.compile(needle)
  22             cache[:] = [re_needle]
  23         else:
  24             re_needle = cache[0]
  25         return (re_needle.match(haystack) is not None)
  26     else:
  27         return (needle in haystack)
  28
  29
  30 class rule_test:
  31     __slots__ = ()
  32
  33     def __new__(cls, *args, **kwargs):
  34         raise RuntimeError("%s should not be instantiated" % cls)
  35
  36     @staticmethod
  37     def NAME(obj, rule, cache):
  38         match_name = rule.match_name
  39         return test_name(rule, match_name, obj.name, cache)
  40
  41     def DATA(obj, rule, cache):
  42         match_name = rule.match_name
  43         obj_data = obj.data
  44         if obj_data is not None:
  45             return test_name(rule, match_name, obj_data.name, cache)
  46         else:
  47             return False
  48
  49     @staticmethod
  50     def COLLECTION(obj, rule, cache):
  51         if not cache:
  52             match_name = rule.match_name
  53             objects = {o for g in bpy.data.collections if test_name(rule, match_name, g.name, cache) for o in g.objects}
  54             cache["objects"] = objects
  55         else:
  56             objects = cache["objects"]
  57
  58         return obj in objects
  59
  60     @staticmethod
  61     def MATERIAL(obj, rule, cache):
  62         match_name = rule.match_name
  63         materials = getattr(obj.data, "materials", None)
  64
  65         return ((materials is not None) and
  66                 (any((test_name(rule, match_name, m.name) for m in materials if m is not None))))
  67
  68     @staticmethod
  69     def TYPE(obj, rule, cache):
  70         return (obj.type == rule.match_object_type)
  71
  72     @staticmethod
  73     def EXPR(obj, rule, cache):
  74         if not cache:
  75             match_expr = rule.match_expr
  76             expr = compile(match_expr, rule.name, 'eval')
  77
  78             namespace = {}
  79             namespace.update(__import__("math").__dict__)
  80
  81             cache["expr"] = expr
  82             cache["namespace"] = namespace
  83         else:
  84             expr = cache["expr"]
  85             namespace = cache["namespace"]
  86
  87         try:
  88             return bool(eval(expr, {}, {"self": obj}))
  89         except:
  90             import traceback
  91             traceback.print_exc()
  92             return False
  93
  94
  95 class rule_draw:
  96     __slots__ = ()
  97
  98     def __new__(cls, *args, **kwargs):
  99         raise RuntimeError("%s should not be instantiated" % cls)
 100
 101     @staticmethod
 102     def _generic_match_name(layout, rule):
 103         layout.label(text="Match Name:")
 104         row = layout.row(align=True)
 105         row.prop(rule, "match_name", text="")
 106         row.prop(rule, "use_match_regex", text="", icon='SORTALPHA')
 107
 108     @staticmethod
 109     def NAME(layout, rule):
 110         rule_draw._generic_match_name(layout, rule)
 111
 112     @staticmethod
 113     def DATA(layout, rule):
 114         rule_draw._generic_match_name(layout, rule)
 115
 116     @staticmethod
 117     def COLLECTION(layout, rule):
 118         rule_draw._generic_match_name(layout, rule)
 119
 120     @staticmethod
 121     def MATERIAL(layout, rule):
 122         rule_draw._generic_match_name(layout, rule)
 123
 124     @staticmethod
 125     def TYPE(layout, rule):
 126         row = layout.row()
 127         row.prop(rule, "match_object_type")
 128
 129     @staticmethod
 130     def EXPR(layout, rule):
 131         col = layout.column()
 132         col.label(text="Scripted Expression:")
 133         col.prop(rule, "match_expr", text="")
 134
 135
 136 def object_colors_calc(rules, objects):
 137     from mathutils import Color
 138
 139     rules_cb = [getattr(rule_test, rule.type) for rule in rules]
 140     rules_blend = [(1.0 - rule.factor, rule.factor) for rule in rules]
 141     rules_color = [Color(rule.color) for rule in rules]
 142     rules_cache = [{} for i in range(len(rules))]
 143     rules_inv = [rule.use_invert for rule in rules]
 144     changed_count = 0
 145
 146     for obj in objects:
 147         is_set = False
 148         obj_color = Color(obj.color[0:3])
 149
 150         for (rule, test_cb, color, blend, cache, use_invert) \
 151              in zip(rules, rules_cb, rules_color, rules_blend, rules_cache, rules_inv):
 152
 153             if test_cb(obj, rule, cache) is not use_invert:
 154                 if is_set is False:
 155                     obj_color = color
 156                 else:
 157                     # prevent mixing colors losing saturation
 158                     obj_color_s = obj_color.s
 159                     obj_color = (obj_color * blend[0]) + (color * blend[1])
 160                     obj_color.s = (obj_color_s * blend[0]) + (color.s * blend[1])
 161
 162                 is_set = True
 163
 164         if is_set:
 165             obj.color[0:3] = obj_color
 166             changed_count += 1
 167     return changed_count
 168
 169
 170 def object_colors_select(rule, objects):
 171     cache = {}
 172
 173     rule_type = rule.type
 174     test_cb = getattr(rule_test, rule_type)
 175
 176     for obj in objects:
 177         obj.select_set(test_cb(obj, rule, cache))
 178
 179
 180 def object_colors_rule_validate(rule, report):
 181     rule_type = rule.type
 182
 183     if rule_type in {'NAME', 'DATA', 'COLLECTION', 'MATERIAL'}:
 184         if rule.use_match_regex:
 185             import re
 186             try:
 187                 re.compile(rule.match_name)
 188             except Exception as e:
 189                 report({'ERROR'}, "Rule %r: %s" % (rule.name, str(e)))
 190                 return False
 191
 192     elif rule_type == 'EXPR':
 193         try:
 194             compile(rule.match_expr, rule.name, 'eval')
 195         except Exception as e:
 196             report({'ERROR'}, "Rule %r: %s" % (rule.name, str(e)))
 197             return False
 198
 199     return True
 200
 201
 202
 203 import bpy
 204 from bpy.types import (
 205     Operator,
 206     Panel,
 207     UIList,
 208 )
 209 from bpy.props import (
 210     StringProperty,
 211     BoolProperty,
 212     IntProperty,
 213     FloatProperty,
 214     EnumProperty,
 215     CollectionProperty,
 216     BoolVectorProperty,
 217     FloatVectorProperty,
 218 )
 219
 220
 221 class OBJECT_PT_color_rules(Panel):
 222     bl_label = "Color Rules"
 223     bl_space_type = 'PROPERTIES'
 224     bl_region_type = 'WINDOW'
 225     bl_context = "object"
 226     bl_options = {'DEFAULT_CLOSED'}
 227
 228     def draw(self, context):
 229         layout = self.layout
 230
 231         scene = context.scene
 232
 233         # Rig type list
 234         row = layout.row()
 235         row.template_list(
 236             "OBJECT_UL_color_rule", "color_rules",
 237             scene, "color_rules",
 238             scene, "color_rules_active_index",
 239         )
 240
 241         col = row.column()
 242         colsub = col.column(align=True)
 243         colsub.operator("object.color_rules_add", icon='ADD', text="")
 244         colsub.operator("object.color_rules_remove", icon='REMOVE', text="")
 245
 246         colsub = col.column(align=True)
 247         colsub.operator("object.color_rules_move", text="", icon='TRIA_UP').direction = -1
 248         colsub.operator("object.color_rules_move", text="", icon='TRIA_DOWN').direction = 1
 249
 250         colsub = col.column(align=True)
 251         colsub.operator("object.color_rules_select", text="", icon='RESTRICT_SELECT_OFF')
 252
 253         if scene.color_rules:
 254             index = scene.color_rules_active_index
 255             rule = scene.color_rules[index]
 256
 257             box = layout.box()
 258             row = box.row(align=True)
 259             row.prop(rule, "name", text="")
 260             row.prop(rule, "type", text="")
 261             row.prop(rule, "use_invert", text="", icon='ARROW_LEFTRIGHT')
 262
 263             draw_cb = getattr(rule_draw, rule.type)
 264             draw_cb(box, rule)
 265
 266             row = layout.split(factor=0.75, align=True)
 267             props = row.operator("object.color_rules_assign", text="Assign Selected")
 268             props.use_selection = True
 269             props = row.operator("object.color_rules_assign", text="All")
 270             props.use_selection = False
 271
 272
 273 class OBJECT_UL_color_rule(UIList):
 274     def draw_item(self, context, layout, data, rule, icon, active_data, active_propname, index):
 275         # assert(isinstance(rule, bpy.types.ShapeKey))
 276         # scene = active_data
 277         split = layout.split(factor=0.5)
 278         row = split.split(align=False)
 279         row.label(text="%s (%s)" % (rule.name, rule.type.lower()))
 280         split = split.split(factor=0.7)
 281         split.prop(rule, "factor", text="", emboss=False)
 282         split.prop(rule, "color", text="")
 283
 284
 285 class OBJECT_OT_color_rules_assign(Operator):
 286     """Assign colors to objects based on user rules"""
 287     bl_idname = "object.color_rules_assign"
 288     bl_label = "Assign Colors"
 289     bl_options = {'UNDO'}
 290
 291     use_selection: BoolProperty(
 292         name="Selected",
 293         description="Apply to selected (otherwise all objects in the scene)",
 294         default=True,
 295     )
 296
 297     def execute(self, context):
 298         scene = context.scene
 299
 300         if self.use_selection:
 301             objects = context.selected_editable_objects
 302         else:
 303             objects = scene.objects
 304
 305         rules = scene.color_rules[:]
 306         for rule in rules:
 307             if not object_colors_rule_validate(rule, self.report):
 308                 return {'CANCELLED'}
 309
 310         changed_count = object_colors_calc(rules, objects)
 311         self.report({'INFO'}, "Set colors for {} of {} objects".format(changed_count, len(objects)))
 312         return {'FINISHED'}
 313
 314
 315 class OBJECT_OT_color_rules_select(Operator):
 316     """Select objects matching the current rule"""
 317     bl_idname = "object.color_rules_select"
 318     bl_label = "Select Rule"
 319     bl_options = {'UNDO'}
 320
 321     def execute(self, context):
 322         scene = context.scene
 323         rule = scene.color_rules[scene.color_rules_active_index]
 324
 325         if not object_colors_rule_validate(rule, self.report):
 326             return {'CANCELLED'}
 327
 328         objects = context.visible_objects
 329         object_colors_select(rule, objects)
 330         return {'FINISHED'}
 331
 332
 333 class OBJECT_OT_color_rules_add(Operator):
 334     bl_idname = "object.color_rules_add"
 335     bl_label = "Add Color Layer"
 336     bl_options = {'UNDO'}
 337
 338     def execute(self, context):
 339         scene = context.scene
 340         rules = scene.color_rules
 341         rule = rules.add()
 342         rule.name = "Rule.%.3d" % len(rules)
 343         scene.color_rules_active_index = len(rules) - 1
 344         return {'FINISHED'}
 345
 346
 347 class OBJECT_OT_color_rules_remove(Operator):
 348     bl_idname = "object.color_rules_remove"
 349     bl_label = "Remove Color Layer"
 350     bl_options = {'UNDO'}
 351
 352     def execute(self, context):
 353         scene = context.scene
 354         rules = scene.color_rules
 355         rules.remove(scene.color_rules_active_index)
 356         if scene.color_rules_active_index > len(rules) - 1:
 357             scene.color_rules_active_index = len(rules) - 1
 358         return {'FINISHED'}
 359
 360
 361 class OBJECT_OT_color_rules_move(Operator):
 362     bl_idname = "object.color_rules_move"
 363     bl_label = "Remove Color Layer"
 364     bl_options = {'UNDO'}
 365     direction: IntProperty()
 366
 367     def execute(self, context):
 368         scene = context.scene
 369         rules = scene.color_rules
 370         index = scene.color_rules_active_index
 371         index_new = index + self.direction
 372         if index_new < len(rules) and index_new >= 0:
 373             rules.move(index, index_new)
 374             scene.color_rules_active_index = index_new
 375             return {'FINISHED'}
 376         else:
 377             return {'CANCELLED'}
 378
 379
 380 class ColorRule(bpy.types.PropertyGroup):
 381     name: StringProperty(
 382         name="Rule Name",
 383     )
 384     color: FloatVectorProperty(
 385         name="Color",
 386         description="Color to assign",
 387         subtype='COLOR', size=3, min=0, max=1, precision=3, step=0.1,
 388         default=(0.5, 0.5, 0.5),
 389     )
 390     factor: FloatProperty(
 391         name="Opacity",
 392         description="Color to assign",
 393         min=0, max=1, precision=1, step=0.1,
 394         default=1.0,
 395     )
 396     type: EnumProperty(
 397         name="Rule Type",
 398         items=(
 399             ('NAME', "Name", "Object name contains this text (or matches regex)"),
 400             ('DATA', "Data Name", "Object data name contains this text (or matches regex)"),
 401             ('COLLECTION', "Collection Name", "Object in collection that contains this text (or matches regex)"),
 402             ('MATERIAL', "Material Name", "Object uses a material name that contains this text (or matches regex)"),
 403             ('TYPE', "Type", "Object type"),
 404             ('EXPR', "Expression", (
 405                 "Scripted expression (using 'self' for the object) eg:\n"
 406                 "  self.type == 'MESH' and len(self.data.vertices) > 20"
 407             )
 408             ),
 409         ),
 410     )
 411
 412     use_invert: BoolProperty(
 413         name="Invert",
 414         description="Match when the rule isn't met",
 415     )
 416
 417     # ------------------
 418     # Matching Variables
 419
 420     # shared by all name matching
 421     match_name: StringProperty(
 422         name="Match Name",
 423     )
 424     use_match_regex: BoolProperty(
 425         name="Regex",
 426         description="Use regular expressions for pattern matching",
 427     )
 428     # type == 'TYPE'
 429     match_object_type: EnumProperty(
 430         name="Object Type",
 431         items=([(i.identifier, i.name, "")
 432                 for i in bpy.types.Object.bl_rna.properties['type'].enum_items]
 433                )
 434     )
 435     # type == 'EXPR'
 436     match_expr: StringProperty(
 437         name="Expression",
 438         description="Python expression, where 'self' is the object variable"
 439     )
 440
 441
 442 classes = (
 443     OBJECT_PT_color_rules,
 444     OBJECT_OT_color_rules_add,
 445     OBJECT_OT_color_rules_remove,
 446     OBJECT_OT_color_rules_move,
 447     OBJECT_OT_color_rules_assign,
 448     OBJECT_OT_color_rules_select,
 449     OBJECT_UL_color_rule,
 450     ColorRule,
 451 )
 452
 453
 454 def register():
 455     for cls in classes:
 456         bpy.utils.register_class(cls)
 457
 458     bpy.types.Scene.color_rules = CollectionProperty(type=ColorRule)
 459     bpy.types.Scene.color_rules_active_index = IntProperty()
 460
 461
 462 def unregister():
 463     for cls in classes:
 464         bpy.utils.unregister_class(cls)
 465
 466     del bpy.types.Scene.color_rules