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