Optimize polygon retrievement in database

1 files changed, 45 insertions, 0 deletions

diff --git a/chimere/models.py b/chimere/models.py
index a68945a..2c27644 100644
--- a/chimere/models.py
+++ b/chimere/models.py
@@ -904,6 +904,51 @@ class Polygon(GeographicItem):
                                      "inner_color": inner_color}}
         return json.dumps(attributes)
 
+    @classmethod
+    def getGeoJSONs(self, queryset, color="#000",
+                    inner_color='rgba(180, 180, 180, 0.3)',
+                    limit_to_categories=[]):
+        vals, default_color, default_inner_color = [], color, inner_color
+        q = queryset.select_related('categories').extra(
+            select={'json': 'ST_AsGeoJSON(polygon)'}).values(
+                'json', 'name', 'pk', 'inner_color', 'color', 'categories__pk')
+        added = []
+        current_categories = {}
+        for polygon in q.all():
+            if polygon['pk'] in added:
+                continue
+            if limit_to_categories and \
+                    polygon["categories__pk"] not in limit_to_categories:
+                continue
+            color = default_color
+            if polygon["color"]:
+                color = polygon['color']
+            elif polygon["categories__pk"]:
+                if polygon["categories__pk"] not in current_categories:
+                    cat = SubCategory.objects.get(pk=polygon["categories__pk"])
+                    # [index, color list]
+                    current_categories[polygon["categories__pk"]] = \
+                        [0, list(Color.objects.filter(
+                            color_theme=cat.color_theme))]
+                idx, colors = current_categories[polygon["categories__pk"]]
+                # category have a color theme
+                if colors:
+                    color = colors[idx % len(colors)].code
+                    # index += 1
+                    current_categories[polygon["categories__pk"]][0] += 1
+
+            inner_color = default_inner_color if not polygon["inner_color"] \
+                else polygon["inner_color"]
+            vals.append({
+                "type": "Feature",
+                "geometry": json.loads(polygon['json']),
+                "properties": {"pk": polygon['pk'], "name": polygon['name'],
+                               'key': "polygon-{}".format(polygon['pk']),
+                               'color': color,
+                               'inner_color': inner_color}})
+            added.append(polygon['pk'])
+        return vals
+
 
 class MultimediaType(models.Model):
     MEDIA_TYPES = (('A', _(u"Audio")),