getGeoJSONs refactoring - add slice ability

1 files changed, 152 insertions, 142 deletions

diff --git a/chimere/models.py b/chimere/models.py
index 1e3a57b..46a778c 100644
--- a/chimere/models.py
+++ b/chimere/models.py
@@ -451,9 +451,9 @@ IMPORTER_CHOICES_DICT = dict(IMPORTER_CHOICES)
 
 
 class Importer(models.Model):
-    '''
+    """
     Data importer for a specific subcategory
-    '''
+    """
     importer_type = models.CharField(_("Importer type"), max_length=4,
                                      choices=IMPORTER_CHOICES)
     filtr = models.TextField(_("Filter"), blank=True, null=True)
@@ -535,6 +535,7 @@ class ImporterKeyCategories(models.Model):
     class Meta:
         verbose_name = _("Importer - Key categories")
 
+
 OVERLAY_CHOICES = (
     ('JSON', 'GeoJSON'),
 )
@@ -608,6 +609,9 @@ class GeographicItem(models.Model):
     search_vector = SearchVectorField(_("Search vector"), blank=True, null=True,
                                       help_text=_("Auto filled at save"))
 
+    default_values = {}
+    json_color_keys = []
+
     class Meta:
         abstract = True
 
@@ -842,17 +846,76 @@ class GeographicItem(models.Model):
             (float((max_weight - min_weight)) or 1), 5)
 
     @classmethod
-    def getGeoJSONs(self, queryset, limit_to_categories=[]):
+    def _json_get_icon_and_colors(cls, item, default_values,
+                                  current_categories):
+        """
+        Utility function for getGeoJSONs specific for each geometry type.
+        Update colors and icons.
+
+        :param item: queried item
+        :param default_values: current values
+        :param current_categories: dict containing categories and how many items
+        these category have been reached (for color rotation)
+        :return: tuple of new properties and updated current categories dict
+        """
         raise NotImplementedError()
 
     @classmethod
-    def search(cls, query, area=None, get_json=False):
+    def getGeoJSONs(cls, queryset, limit_to_categories=None, slice=None):
+        default_values = copy.copy(cls.default_values)
+        q_values = ['json', 'name', 'pk', 'categories__pk'] + \
+            cls.json_color_keys
+        q = queryset.select_related('categories').extra(
+            select={'json': 'ST_AsGeoJSON({})'.format(cls.geom_attr)}
+        ).values(*q_values)
+
+        vals, added, current_categories = [], [], {}
+        start, end = 0, None
+        if slice:
+            start, end = slice
+        for item in q.all():
+            if end and len(vals) > end:
+                break
+            if item['pk'] in added:
+                continue
+            if limit_to_categories and \
+                            item["categories__pk"] not in limit_to_categories:
+                continue
+            added.append(item['pk'])
+            if start:
+                start -= 1
+                continue
+            properties = {"pk": item['pk'], "name": item['name'],
+                          'key': "{}-{}".format(cls.geom_attr, item['pk'])}
+            icon_and_colors, current_categories = \
+                cls._json_get_icon_and_colors(item, default_values,
+                                              current_categories)
+            if not icon_and_colors:
+                # something went wrong when fetching icon / colors
+                continue
+            properties.update(icon_and_colors)
+            try:
+                # bad geometry are not displayed
+                geom = json.loads(item['json'])
+            except json.decoder.JSONDecodeError:
+                continue
+
+            vals.append({
+                "type": "Feature",
+                "geometry": geom,
+                "properties": properties
+            })
+        return vals
+
+    @classmethod
+    def search(cls, query, area=None, get_json=False, slice=None):
         """
         Search items using search_vector
 
         :param query: terms to search
         :param area: area limitation (default None)
         :param get_json: output as json if True else return a queryset
+        :param slice: startting and ending bounds
         :return: items matching the search query and the area
         """
         subcats = SubCategory.getAvailable(instance=True, area=area)
@@ -861,8 +924,11 @@ class GeographicItem(models.Model):
                                    query,
                                    config=settings.CHIMERE_SEARCH_LANGUAGE))
         if get_json:
-            return cls.getGeoJSONs(q)
-        return q
+            return cls.getGeoJSONs(q, slice=slice)
+        if slice:
+            return q.all()[slice[0]:slice[1]]
+        else:
+            return q.all()
 
     def update_search_vector(self, save=True):
         """
@@ -1036,53 +1102,32 @@ class Marker(GeographicItem):
         return json.dumps(jsons)
 
     @classmethod
-    def getGeoJSONs(self, queryset, limit_to_categories=[]):
-        vals = []
-        q = queryset.select_related('categories').extra(
-            select={'json': 'ST_AsGeoJSON(point)'}).values(
-                'json', 'name', 'pk', 'categories__pk')
-        added, cats = [], {}
-        for item in q.all():
-            if item['pk'] in added:
-                continue
-            if limit_to_categories and \
-                    item["categories__pk"] not in limit_to_categories:
-                continue
-            if item['categories__pk'] not in cats:
-                try:
-                    cat = SubCategory.objects.get(
-                        available=True, pk=item['categories__pk'])
-                except SubCategory.DoesNotExist:
-                    continue
-                cats[item['categories__pk']] = {
-                    'icon_path': str(cat.icon.image),
-                    'icon_hover_path': str(cat.hover_icon.image)
-                    if cat.hover_icon else '',
-                    'icon_offset_x': cat.icon.offset_x,
-                    'icon_offset_y': cat.icon.offset_y,
-                    'icon_popup_offset_x': cat.icon.popup_offset_x,
-                    'icon_popup_offset_y': cat.icon.popup_offset_y,
-                    'category_name': cat.name}
-                try:
-                    cats[item['categories__pk']].update(
-                        {'icon_width': cat.icon.image.width,
-                         'icon_height': cat.icon.image.height,
-                         })
-                except IOError:
-                    pass
+    def _json_get_icon_and_colors(cls, item, default_values,
+                                  current_categories):
+        if item['categories__pk'] not in current_categories:
             try:
-                geom = json.loads(item['json'])
-            except json.decoder.JSONDecodeError:
-                continue
-            dct = {
-                "type": "Feature",
-                "geometry": geom,
-                "properties": {"pk": item['pk'], "name": item['name'],
-                               'key': "marker-{}".format(item['pk'])}}
-            dct['properties'].update(cats[item['categories__pk']])
-            vals.append(dct)
-            added.append(item['pk'])
-        return vals
+                cat = SubCategory.objects.get(
+                    available=True, pk=item['categories__pk'])
+            except SubCategory.DoesNotExist:
+                return None, current_categories
+            current_categories[item['categories__pk']] = {
+                'icon_path': str(cat.icon.image),
+                'icon_hover_path': str(cat.hover_icon.image)
+                if cat.hover_icon else '',
+                'icon_offset_x': cat.icon.offset_x,
+                'icon_offset_y': cat.icon.offset_y,
+                'icon_popup_offset_x': cat.icon.popup_offset_x,
+                'icon_popup_offset_y': cat.icon.popup_offset_y,
+                'category_name': cat.name}
+            try:
+                current_categories[item['categories__pk']].update(
+                    {'icon_width': cat.icon.image.width,
+                     'icon_height': cat.icon.image.height,
+                     })
+            except IOError:
+                # image is not available
+                pass
+        return current_categories[item['categories__pk']], current_categories
 
     @property
     def default_category(self):
@@ -1196,6 +1241,9 @@ class Polygon(GeographicItem):
         help_text=_("HTML code/name"), blank=True, null=True)
     objects = models.GeoManager()
     geom_attr = 'polygon'
+    json_color_keys = ['color', 'inner_color']
+    default_values = {'color': "#000",
+                      'inner_color': 'rgba(180, 180, 180, 0.3)'}
 
     class Meta:
         ordering = ('status', 'name')
@@ -1219,53 +1267,33 @@ class Polygon(GeographicItem):
         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:
-                    c = colors[idx % len(colors)]
-                    color = c.code
-                    if c.inner_code:
-                        inner_color = c.inner_code
-                    # index += 1
-                    current_categories[polygon["categories__pk"]][0] += 1
-            if polygon["inner_color"]:
-                inner_color = polygon["inner_color"]
-            elif not inner_color:
-                inner_color = default_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
+    def _json_get_icon_and_colors(cls, item, default_values,
+                                  current_categories):
+        inner_color = None
+        color = default_values['color']
+        if item["color"]:
+            color = item['color']
+        elif item["categories__pk"]:
+            if item["categories__pk"] not in current_categories:
+                cat = SubCategory.objects.get(pk=item["categories__pk"])
+                # [index, color list]
+                current_categories[item["categories__pk"]] = \
+                    [0, list(Color.objects.filter(
+                        color_theme=cat.color_theme))]
+            idx, colors = current_categories[item["categories__pk"]]
+            # category have a color theme
+            if colors:
+                c = colors[idx % len(colors)]
+                color = c.code
+                if c.inner_code:
+                    inner_color = c.inner_code
+                # index += 1
+                current_categories[item["categories__pk"]][0] += 1
+        if item["inner_color"]:
+            inner_color = item["inner_color"]
+        elif not inner_color:
+            inner_color = default_values['inner_color']
+        return {"color": color, "inner_color": inner_color}, current_categories
 
     def get_full_dict(self):
         dct = super(Polygon, self).get_full_dict()
@@ -1290,8 +1318,6 @@ class AggregatedPolygon(models.Model):
         db_table = 'chimere_aggregated_polygons'
 
     def getGeoJSON(self, color="", inner_color=''):
-        '''Return a GeoJSON string
-        '''
         # get colors
         if not color or not inner_color:
             q = Color.objects.filter(
@@ -1660,14 +1686,17 @@ class Route(GeographicItem):
         blank=True, null=True)
     objects = models.GeoManager()
     geom_attr = 'route'
+    json_color_keys = ['color']
+    default_values = {'color': '#000'}
 
     class Meta:
         ordering = ('status', 'name')
         verbose_name = _("Route")
 
     def getGeoJSON(self, color="#000"):
-        '''Return a GeoJSON string
-        '''
+        """
+        Return a GeoJSON string
+        """
         try:
             geom = json.loads(self.route.geojson)
         except json.JSONDecodeError:
@@ -1686,45 +1715,26 @@ class Route(GeographicItem):
         return dct
 
     @classmethod
-    def getGeoJSONs(self, queryset, color="#000",
-                    limit_to_categories=[]):
-        vals, default_color = [], color
-        q = queryset.select_related('categories').extra(
-            select={'json': 'ST_AsGeoJSON(route)'}).values(
-                'json', 'name', 'pk', 'color', 'categories__pk')
-        added = []
-        current_categories = {}
-        for item in q.all():
-            if item['pk'] in added:
-                continue
-            if limit_to_categories and \
-                    item["categories__pk"] not in limit_to_categories:
-                continue
-            color = default_color
-            if item["color"]:
-                color = item['color']
-            elif item["categories__pk"]:
-                if item["categories__pk"] not in current_categories:
-                    cat = SubCategory.objects.get(pk=item["categories__pk"])
-                    # [index, color list]
-                    current_categories[item["categories__pk"]] = \
-                        [0, list(Color.objects.filter(
-                            color_theme=cat.color_theme))]
-                idx, colors = current_categories[item["categories__pk"]]
-                # category have a color theme
-                if colors:
-                    c = colors[idx % len(colors)]
-                    color = c.code
-                    # index += 1
-                    current_categories[item["categories__pk"]][0] += 1
-            vals.append({
-                "type": "Feature",
-                "geometry": json.loads(item['json']),
-                "properties": {"pk": item['pk'], "name": item['name'],
-                               'key': "route-{}".format(item['pk']),
-                               'color': color}})
-            added.append(item['pk'])
-        return vals
+    def _json_get_icon_and_colors(cls, item, default_values,
+                                  current_categories):
+        color = default_values['color']
+        if item["color"]:
+            color = item['color']
+        elif item["categories__pk"]:
+            if item["categories__pk"] not in current_categories:
+                cat = SubCategory.objects.get(pk=item["categories__pk"])
+                # [index, color list]
+                current_categories[item["categories__pk"]] = \
+                    [0, list(Color.objects.filter(
+                        color_theme=cat.color_theme))]
+            idx, colors = current_categories[item["categories__pk"]]
+            # category have a color theme
+            if colors:
+                c = colors[idx % len(colors)]
+                color = c.code
+                # index += 1
+                current_categories[item["categories__pk"]][0] += 1
+        return {"color": color}, current_categories
 
 pre_save_route_values = {}