object_color_rules.py

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