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