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