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 loosing 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
 240     def draw(self, context):
 241         layout = self.layout
 242
 243         scene = context.scene
 244
 245         # Rig type list
 246         row = layout.row()
 247         row.template_list(
 248             "OBJECT_UL_color_rule", "color_rules",
 249             scene, "color_rules",
 250             scene, "color_rules_active_index",
 251         )
 252
 253         col = row.column()
 254         colsub = col.column(align=True)
 255         colsub.operator("object.color_rules_add", icon='ADD', text="")
 256         colsub.operator("object.color_rules_remove", icon='REMOVE', text="")
 257
 258         colsub = col.column(align=True)
 259         colsub.operator("object.color_rules_move", text="", icon='TRIA_UP').direction = -1
 260         colsub.operator("object.color_rules_move", text="", icon='TRIA_DOWN').direction = 1
 261
 262         colsub = col.column(align=True)
 263         colsub.operator("object.color_rules_select", text="", icon='RESTRICT_SELECT_OFF')
 264
 265         if scene.color_rules:
 266             index = scene.color_rules_active_index
 267             rule = scene.color_rules[index]
 268
 269             box = layout.box()
 270             row = box.row(align=True)
 271             row.prop(rule, "name", text="")
 272             row.prop(rule, "type", text="")
 273             row.prop(rule, "use_invert", text="", icon='ARROW_LEFTRIGHT')
 274
 275             draw_cb = getattr(rule_draw, rule.type)
 276             draw_cb(box, rule)
 277
 278             row = layout.split(factor=0.75, align=True)
 279             props = row.operator("object.color_rules_assign", text="Assign Selected")
 280             props.use_selection = True
 281             props = row.operator("object.color_rules_assign", text="All")
 282             props.use_selection = False
 283
 284
 285 class OBJECT_UL_color_rule(UIList):
 286     def draw_item(self, context, layout, data, rule, icon, active_data, active_propname, index):
 287         # assert(isinstance(rule, bpy.types.ShapeKey))
 288         # scene = active_data
 289         split = layout.split(factor=0.5)
 290         row = split.split(align=False)
 291         row.label(text="%s (%s)" % (rule.name, rule.type.lower()))
 292         split = split.split(factor=0.7)
 293         split.prop(rule, "factor", text="", emboss=False)
 294         split.prop(rule, "color", text="")
 295
 296
 297 class OBJECT_OT_color_rules_assign(Operator):
 298     """Assign colors to objects based on user rules"""
 299     bl_idname = "object.color_rules_assign"
 300     bl_label = "Assign Colors"
 301     bl_options = {'UNDO'}
 302
 303     use_selection: BoolProperty(
 304         name="Selected",
 305         description="Apply to selected (otherwise all objects in the scene)",
 306         default=True,
 307     )
 308
 309     def execute(self, context):
 310         scene = context.scene
 311
 312         if self.use_selection:
 313             objects = context.selected_editable_objects
 314         else:
 315             objects = scene.objects
 316
 317         rules = scene.color_rules[:]
 318         for rule in rules:
 319             if not object_colors_rule_validate(rule, self.report):
 320                 return {'CANCELLED'}
 321
 322         changed_count = object_colors_calc(rules, objects)
 323         self.report({'INFO'}, "Set colors for {} of {} objects".format(changed_count, len(objects)))
 324         return {'FINISHED'}
 325
 326
 327 class OBJECT_OT_color_rules_select(Operator):
 328     """Select objects matching the current rule"""
 329     bl_idname = "object.color_rules_select"
 330     bl_label = "Select Rule"
 331     bl_options = {'UNDO'}
 332
 333     def execute(self, context):
 334         scene = context.scene
 335         rule = scene.color_rules[scene.color_rules_active_index]
 336
 337         if not object_colors_rule_validate(rule, self.report):
 338             return {'CANCELLED'}
 339
 340         objects = context.visible_objects
 341         object_colors_select(rule, objects)
 342         return {'FINISHED'}
 343
 344
 345 class OBJECT_OT_color_rules_add(Operator):
 346     bl_idname = "object.color_rules_add"
 347     bl_label = "Add Color Layer"
 348     bl_options = {'UNDO'}
 349
 350     def execute(self, context):
 351         scene = context.scene
 352         rules = scene.color_rules
 353         rule = rules.add()
 354         rule.name = "Rule.%.3d" % len(rules)
 355         scene.color_rules_active_index = len(rules) - 1
 356         return {'FINISHED'}
 357
 358
 359 class OBJECT_OT_color_rules_remove(Operator):
 360     bl_idname = "object.color_rules_remove"
 361     bl_label = "Remove Color Layer"
 362     bl_options = {'UNDO'}
 363
 364     def execute(self, context):
 365         scene = context.scene
 366         rules = scene.color_rules
 367         rules.remove(scene.color_rules_active_index)
 368         if scene.color_rules_active_index > len(rules) - 1:
 369             scene.color_rules_active_index = len(rules) - 1
 370         return {'FINISHED'}
 371
 372
 373 class OBJECT_OT_color_rules_move(Operator):
 374     bl_idname = "object.color_rules_move"
 375     bl_label = "Remove Color Layer"
 376     bl_options = {'UNDO'}
 377     direction: IntProperty()
 378
 379     def execute(self, context):
 380         scene = context.scene
 381         rules = scene.color_rules
 382         index = scene.color_rules_active_index
 383         index_new = index + self.direction
 384         if index_new < len(rules) and index_new >= 0:
 385             rules.move(index, index_new)
 386             scene.color_rules_active_index = index_new
 387             return {'FINISHED'}
 388         else:
 389             return {'CANCELLED'}
 390
 391
 392 class ColorRule(bpy.types.PropertyGroup):
 393     name: StringProperty(
 394         name="Rule Name",
 395     )
 396     color: FloatVectorProperty(
 397         name="Color",
 398         description="Color to assign",
 399         subtype='COLOR', size=3, min=0, max=1, precision=3, step=0.1,
 400         default=(0.5, 0.5, 0.5),
 401     )
 402     factor: FloatProperty(
 403         name="Opacity",
 404         description="Color to assign",
 405         min=0, max=1, precision=1, step=0.1,
 406         default=1.0,
 407     )
 408     type: EnumProperty(
 409         name="Rule Type",
 410         items=(
 411             ('NAME', "Name", "Object name contains this text (or matches regex)"),
 412             ('DATA', "Data Name", "Object data name contains this text (or matches regex)"),
 413             ('COLLECTION', "Collection Name", "Object in collection that contains this text (or matches regex)"),
 414             ('MATERIAL', "Material Name", "Object uses a material name that contains this text (or matches regex)"),
 415             ('TYPE', "Type", "Object type"),
 416             ('EXPR', "Expression", (
 417                 "Scripted expression (using 'self' for the object) eg:\n"
 418                 "  self.type == 'MESH' and len(self.data.vertices) > 20"
 419             )
 420             ),
 421         ),
 422     )
 423
 424     use_invert: BoolProperty(
 425         name="Invert",
 426         description="Match when the rule isn't met",
 427     )
 428
 429     # ------------------
 430     # Matching Variables
 431
 432     # shared by all name matching
 433     match_name: StringProperty(
 434         name="Match Name",
 435     )
 436     use_match_regex: BoolProperty(
 437         name="Regex",
 438         description="Use regular expressions for pattern matching",
 439     )
 440     # type == 'TYPE'
 441     match_object_type: EnumProperty(
 442         name="Object Type",
 443         items=([(i.identifier, i.name, "")
 444                 for i in bpy.types.Object.bl_rna.properties['type'].enum_items]
 445                )
 446     )
 447     # type == 'EXPR'
 448     match_expr: StringProperty(
 449         name="Expression",
 450         description="Python expression, where 'self' is the object variable"
 451     )
 452
 453
 454 classes = (
 455     OBJECT_PT_color_rules,
 456     OBJECT_OT_color_rules_add,
 457     OBJECT_OT_color_rules_remove,
 458     OBJECT_OT_color_rules_move,
 459     OBJECT_OT_color_rules_assign,
 460     OBJECT_OT_color_rules_select,
 461     OBJECT_UL_color_rule,
 462     ColorRule,
 463 )
 464
 465
 466 def register():
 467     for cls in classes:
 468         bpy.utils.register_class(cls)
 469
 470     bpy.types.Scene.color_rules = CollectionProperty(type=ColorRule)
 471     bpy.types.Scene.color_rules_active_index = IntProperty()
 472
 473
 474 def unregister():
 475     for cls in classes:
 476         bpy.utils.unregister_class(cls)
 477
 478     del bpy.types.Scene.color_rules