object_color_rules.py

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