source: products/quintagroup.analytics/branches/treemap/quintagroup/analytics/browser/treemap.py @ 3372

Last change on this file since 3372 was 3372, checked in by potar, 12 years ago

Added new branch

File size: 17.1 KB
Line 
1""" Module dedicated to create treemap """
2
3import colorsys
4from StringIO import StringIO
5
6from PIL import Image
7from PIL import ImageDraw
8from PIL import ImageFont
9
10from Products.Five.browser import BrowserView
11from views import SizeByPath
12from const import FONT_PATH, IMAGE_MASK, USER_FONT, WIDTH, HEIGHT, \
13                  TEXT_SATURATION, TEXT_LIGHT, MAX_PERCENT, MAX_ANGLE, \
14                  DIV_RGB, DIV_SIDE, DIVERSION_FATHER, DIVERSION_CHILD, \
15                  DIVERSION_TEXT_CHILD_V, DIVERSION_TEXT_CHILD_H, \
16                  FONT_SIZE_FATHER, FONT_SIZE_CHILDREN_MIN, FRAME_COLOR, \
17                  USER_COLOR, COLOR_FIELD, DIVERSION_TEXT_BORDER, \
18                  DIVERSION_CHILD, DIVERSION_VERTICAL, FRAME_DIVERSION
19
20class TreemapBTree:
21    """ Class dedicated to create treemap binary tree """
22    def __init__(self, left=None, right=None, data=None):
23        self.left = left
24        self.right = right
25        self.data = data
26       
27    def getWeight(self, treemap):
28        """ Method gets size of treemap objects """
29        return float(sum(x.size for x in treemap))
30   
31    def getHalfSize(self, items):
32        """  Method gets index of treemap objects (half size) """
33        half_size = self.getWeight(items) / 2
34        pre_size = next_size = 0.0
35        for i,treemap in enumerate(items):
36            next_size = pre_size + treemap.size
37            if abs(half_size - next_size) > abs(half_size - pre_size):
38                return i
39            pre_size = next_size
40
41       
42    def addNodes(self, treemap_data):
43        """  Recursive method adds treemap objects into TreemapBTree """
44        def leaf(self,treemap_data):
45            """ Method sets leaf into treemap tree """
46            if len(treemap_data) == 0:
47                return True
48           
49            elif isinstance(treemap_data, Treemap) and treemap_data.title == 'Root':
50                self.addNodes(treemap_data[0])
51                return False 
52
53            elif len(treemap_data) == 1:
54                if len(treemap_data[0]) >= 1:
55                    self.data = treemap_data[0]
56                    self.addNodes(treemap_data[0])
57                elif type(treemap_data[0]) is Treemap:
58                    self.data = treemap_data[0]
59                return True
60
61        if leaf(self, treemap_data):
62            return
63
64        half = self.getHalfSize(treemap_data)
65        treemap1 = treemap_data[0:half]
66        treemap2 = treemap_data[half:]
67        weight1 = self.getWeight(treemap1)
68        weight2 = self.getWeight(treemap2)
69       
70        self.left = TreemapBTree(data=Treemap(size=weight1))
71        self.right = TreemapBTree(data=Treemap(size=weight2))
72       
73        self.left.addNodes(treemap1)
74        self.right.addNodes(treemap2)
75   
76class TreemapControl(BrowserView):
77    """ Class dedicated to control Treemap objects """
78    def __init__(self, context, request, treemap, data):
79        super(BrowserView, self).__init__(context, request)
80
81        # eg: self.data = SizeByPath(context, request)
82        self.data = data
83        self.treemap = treemap
84        self.binary_tree = TreemapBTree(data=self.treemap) 
85
86    def setFieldColor(self, data):
87        """
88        >>> x = TreemapControl()
89        >>> x.setFieldColor([{'id': 'admin'}, {'id': 'admin2'}])
90        [{'color': '#fde5be', 'id': 'admin'}, {'color': '#4af6a0', 'id': 'admin2'}]
91        """
92        colors_list = COLOR_FIELD[0:len(data)]
93        for i,x in enumerate(data):
94            x['color'] = colors_list[i]
95        return data
96       
97    def sortItems(self, data):
98        """ Method sorts fields for treemap
99        >>> TreemapControl().sortItems([{'id': 'admin', 'size': 10}, {'id': 'admin2', 'size': 20}])
100        [{'id': 'admin2', 'size': 20}, {'id': 'admin', 'size': 10}]
101        """
102        getWeight = lambda x: x.get('size')
103        return sorted(data, key=getWeight, reverse=True)
104
105    def createRectangles(self, data):
106        """ Method create treemap rectangles """
107        for field_data in self.sortItems(data):
108             self.treemap.addItem(Treemap(size=field_data.get('size'),
109                                          color=field_data.get('color'),
110                                          title=field_data.get('id')))
111    def setTreemapData(self):
112        """ Method sets treemap attributes """
113        # self.data.getTreemapInfo() -> [{'id': 'admin', size: '10', 'type': 'Image'},...]
114        data = self.data.getTreemapInfo()
115        self.setFieldColor(data)
116        self.createRectangles(data)
117
118class Treemap(list):
119    """ Class dedicated to create rectangle """
120    def __init__(self, x=None, y=None, w=None, h=None, color=None, title="", size=0):
121        self.x = x
122        self.y = y
123        self.w = w    # width
124        self.h = h    # hight
125        self.color = color
126        self.size = float( size )
127        self.title = title
128
129    def coordinates_get(self):
130        """ Method gets rectangle coordinates """
131        return (self.x, self.y, self.w, self.h)
132    def coordinates_set(self, value):
133        """ Method sets rectangle coordinates """
134        (self.x, self.y, self.w, self.h) = value
135    coordinates = property(fget=coordinates_get, fset=coordinates_set)
136
137    def addItem(self, item):
138        """ Method adds treemap object into root treemap """
139        if float(item.size) != 0.0:
140            self.append(item)
141       
142    def addItems(self, items):
143        """ Method adds treemap objects into root treemap """
144        for item in items:
145            self.addItem(item)
146
147    def setCoordinates(self, treemap):
148        """  Recursive method adds coordinates into rectangles """
149        if not(treemap.left and treemap.right):
150            return
151
152        if treemap.data.w > treemap.data.h:
153            treemap.left.data.coordinates = (treemap.data.x, treemap.data.y,
154                                             treemap.data.w * (treemap.left.data.size/
155                                            (treemap.left.data.size + treemap.right.data.size)), 
156                                             treemap.data.h)
157            treemap.right.data.coordinates = (treemap.data.x + treemap.left.data.w,
158                                              treemap.data.y, 
159                                              treemap.data.w - treemap.left.data.w,
160                                              treemap.data.h)
161        else:
162            treemap.left.data.coordinates = (treemap.data.x, treemap.data.y, 
163                                             treemap.data.w,
164                                             treemap.data.h * (treemap.left.data.size/
165                                             (treemap.left.data.size + treemap.right.data.size)))
166            treemap.right.data.coordinates = (treemap.data.x, treemap.data.y + treemap.left.data.h, 
167                                              treemap.data.w,
168                                              treemap.data.h - treemap.left.data.h)
169
170        self.setCoordinates(treemap.left)
171        self.setCoordinates(treemap.right)
172
173
174         
175
176class TreemapHtmlControl(TreemapControl):
177    """ Class dedicated to control Treemap objects """
178    def __init__(self, context, request):
179        data = SizeByPath(context, request)
180        super(TreemapHtmlControl, self).__init__(context, request, 
181                                                 TreemapHtml(), 
182                                                 data)
183
184    def getTreemap(self):
185        """ Method generates treemap (using html) """
186        self.setTreemapData()
187        self.binary_tree.addNodes(self.treemap)
188        self.treemap.setTreemapData(self.binary_tree)
189        return self.treemap.createTreemap(self.binary_tree)
190
191class TreemapImageControl(TreemapControl):
192    """ Class dedicated to control Treemap objects """
193    def __init__(self, context, request):
194        super(TreemapImageControl, self).__init__(context, request, 
195                                             TreemapImage(),
196                                             SizeByPath(context, request))
197    def getTreemap(self):
198        """ Method generates treemap view (PIL)"""
199        self.setTreemapData()
200        self.binary_tree.addNodes(self.treemap)
201        self.treemap.setCoordinates(self.binary_tree)
202
203        treemap = ImageDraw.Draw(self.treemap.image)
204        self.treemap.drawTreemap(self.binary_tree, treemap)
205        self.treemap.image = self.treemap.drawContainer()
206        self.treemap.image = self.treemap.createFrame()
207        self.treemap.image.show()
208
209        output_file = StringIO()
210        self.treemap.image.save(output_file,"PNG")
211        try:
212            return output_file.getvalue()
213        finally:
214            output_file.close()
215
216
217class TreemapHtml(Treemap):
218    """ Class dedicated to create Treemap objects """
219    def __init__(self, x = 0, y = 0, w = 100, h = 100, title = "Root", 
220                                               position=""):
221        Treemap.__init__(self, x, y, w, h, title = title)
222#        self.__position = position
223
224#    def position_get(self):
225#        """ Method gets rectangle position """
226#        return self.__position
227#    def position_set(self, value):
228#        """ Method sets rectangle position """
229#        if value in ['horizontal', 'vertical']:
230#            self.__position = value
231#    position = property(fget=position_get, fset=position_set)
232
233    def getClasses(self, treemap):
234        """ Method provides classes for html tag """
235        if not(treemap.left and treemap.right):
236            return "" 
237
238        elif treemap.data.w >= treemap.data.h:
239            return 'horizontal prop%d' % int(round(treemap.left.data.w, -1))
240        else:
241            return 'vertical prop%d' % int(round(treemap.left.data.h, -1))
242
243           
244    def _writeHtml(self, treemap):
245        """ Method generates html tag (<div>) """
246        if not(treemap.left and treemap.right):
247            return "" 
248   
249        return  '<div class="%s" style="background-color:%s">' % \
250                              (self.getClasses(treemap.left), treemap.left.data.color) + \
251                              self._writeHtml(treemap.left) + '</div>' + \
252                '<div class="%s" style="background-color:%s">'% \
253                              (self.getClasses(treemap.right), treemap.right.data.color) + \
254                               self._writeHtml(treemap.right) +'</div>'
255
256    def createTreemap(self, treemap):
257        """ Method generates html tag (treemap) """
258        return  '<div class="%s">' % \
259                       self.getClasses(treemap) + \
260                       self._writeHtml(treemap) + \
261                '</div>'
262               
263       
264    def setTreemapData(self, treemap):
265        """  Recursive method adds coordinates into rectangles """
266        if not(treemap.left and treemap.right):
267            return
268        if treemap.data.w >= treemap.data.h:
269            treemap.left.data.coordinates = (None, None,
270                                             100 * (treemap.left.data.size/
271                                            (treemap.left.data.size + treemap.right.data.size)), 
272                                             100)
273            treemap.right.data.coordinates = (None,
274                                              None, 
275                                              treemap.data.w - treemap.left.data.w,
276                                              100)
277        else:
278            treemap.left.data.coordinates = (None, None, 
279                                             100,
280                                             100 * (treemap.left.data.size/
281                                             (treemap.left.data.size + treemap.right.data.size)))
282            treemap.right.data.coordinates = (None, None, 
283                                              100,
284                                              treemap.data.h - treemap.left.data.h)
285
286        self.setTreemapData(treemap.left)
287        self.setTreemapData(treemap.right)
288
289class TreemapImage(Treemap):
290    """ Class dedicated to create treemap image """
291    def __init__(self, x = 0, y = 0, w = WIDTH, h = HEIGHT, title = "Root"):
292        Treemap.__init__(self, x, y, w, h, title = title)
293        self.image = Image.new("RGB", (w, h), "white")
294        self.size = (w, h)
295
296    def drawContainer(self):
297        """ Method draws treemap (name of item in the top) """
298        font = self.createFont(FONT_SIZE_FATHER, USER_FONT)
299        treemap = Image.new( "RGB", self.size, "white")
300        rectangle = ImageDraw.Draw(treemap)
301        for main in self:
302             rectangle.rectangle((main.x, main.y, main.x + main.w, 
303                                        main.h + main.y),
304                                        fill=main.color)
305             if font.getsize(main.title)[0] < main.w:
306                 rectangle.text((main.x, main.y), text = main.title, 
307                                       font = font)
308             region = self.image.crop((int(round(main.x)) + DIVERSION_CHILD, 
309                                               int(round(main.y)) + DIVERSION_CHILD,
310                                               int(round(main.x + main.w) - DIVERSION_CHILD), 
311                                               int(round(main.h + main.y)) - DIVERSION_CHILD))
312             try:
313                 region = region.resize((int(round(main.w - DIVERSION_FATHER * 2)),
314                                        int(round(main.h - DIVERSION_FATHER * 2 - DIVERSION_VERTICAL))))
315             except MemoryError:
316                 continue
317             treemap.paste(region,(int(round( main.x + DIVERSION_FATHER)),
318                           int(round( main.y + DIVERSION_FATHER + DIVERSION_VERTICAL))))
319        return treemap
320
321    def rgbHexToDecimal(self, hex_str):
322        """
323        >>> TreemapImage().rgbHexToDecimal('#fde5be')
324        (253, 229, 190)
325        """
326        hex_str = hex_str[1:]
327        R = int(hex_str[0:2], 16)
328        G = int(hex_str[2:4], 16)
329        B = int(hex_str[4:], 16)
330        return R,G,B
331
332    def getHlsPIL(self, r, g, b):
333        """
334        >>> TreemapImage().getHlsPIL(10,10,20)
335        'hsl(10,10%,20%)'
336        """
337        return "hsl(%d,%d%%,%d%%)"%(r, g, b)
338
339    def getFontSize(self, field):
340        """ Method gets font size for rectangles """
341#        font_size = field.w/DIV_SID if field.h < field.w else field.h / DIV_SIDE
342
343        if field.h < field.w:
344            font_size =  field.w / DIV_SIDE
345        else:
346            font_size =  field.h / DIV_SIDE
347        if font_size < FONT_SIZE_CHILDREN_MIN:
348            return FONT_SIZE_CHILDREN_MIN
349        else:
350            return font_size
351
352    def createFont(self, font_size, font=''):
353        """ Method creates font for treemap """
354        return ImageFont.truetype( FONT_PATH + '/vera/Vera%s.ttf'%font, font_size)
355
356    def setMask(self, field):
357        """ Method sets mask on rectangle """
358        try:
359            image_mask = IMAGE_MASK.resize((field.w - DIVERSION_CHILD,
360                                            field.h - DIVERSION_CHILD))
361        except MemoryError:
362            return
363        self.image.paste("white", (field.x + DIVERSION_CHILD,
364                                 field.y + DIVERSION_CHILD),
365                                 mask = image_mask)
366
367    def getPilColor(self, color):
368        """
369        >>> TreemapImage().getPilColor('#fde5be')
370        'hsl(37,144%,56%)'
371        """
372        r, g, b = self.rgbHexToDecimal(color)
373        h, l, s = colorsys.rgb_to_hls(r / DIV_RGB, g / DIV_RGB, b / DIV_RGB)
374        l -= TEXT_LIGHT
375        s += TEXT_SATURATION
376        return self.getHlsPIL(int(h * MAX_ANGLE), int(s * MAX_PERCENT), 
377                              int(l * MAX_PERCENT))
378
379    def writeText(self, treemap, image):
380        """ Method dedicated to write text in treemap rectangles """
381        field_name = treemap.data.title
382        font = self.createFont(self.getFontSize(treemap.data))
383   
384        # font_size = (126, 47); font_size[0] - width, font_size[1] - height
385        font_size = font.getsize(field_name)
386        if treemap.data.w - DIVERSION_TEXT_CHILD_H + DIVERSION_TEXT_BORDER > font_size[0] and \
387              font_size[1] < treemap.data.h - DIVERSION_TEXT_CHILD_V + DIVERSION_TEXT_BORDER:
388            image.text((treemap.data.x + DIVERSION_TEXT_CHILD_H ,
389                        treemap.data.y + DIVERSION_TEXT_CHILD_V),
390                        text = field_name, font = font, 
391                        fill = self.getPilColor(treemap.data.color))
392       
393    def drawTreemap(self, treemap, image):
394        """ Method generates treemap image """
395        if treemap.data.color:
396            image.rectangle((treemap.data.x + DIVERSION_CHILD, treemap.data.y + DIVERSION_CHILD, 
397                             treemap.data.x + treemap.data.w - DIVERSION_CHILD, 
398                             treemap.data.h + treemap.data.y - DIVERSION_CHILD),
399                             fill= treemap.data.color )
400   
401            self.setMask(treemap.data)
402            self.writeText(treemap, image)
403
404        if not(treemap.left and treemap.right):
405            return 
406       
407        self.drawTreemap(treemap.left, image)
408        self.drawTreemap(treemap.right, image)
409       
410    def createFrame(self):
411        """ Method creates frame for treemap image """
412        treemap = Image.new("RGB", 
413                           (WIDTH + FRAME_DIVERSION,HEIGHT + FRAME_DIVERSION),
414                            FRAME_COLOR)
415        treemap.paste(self.image,(FRAME_DIVERSION/2, FRAME_DIVERSION/2))
416        return treemap
Note: See TracBrowser for help on using the repository browser.