Sliceable and ordered reordering

arkindex/images/management/commands/reorder_polygons.py

 from datetime import timedelta
+from math import ceil
 from timeit import default_timer
 from django.core.management.base import BaseCommand
 from arkindex.project.polygon import Polygon
@@ -17,11 +18,27 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         super().add_arguments(parser)
+        parser.add_argument(
+            '--offset',
+            type=int,
+            default=0,
+        )
+        parser.add_argument(
+            '--limit',
+            type=int,
+            default=0,
+        )
         parser.add_argument(
             '--batch-size',
+            help='Batch size used to retrieve zones.',
             type=int,
             default=1000,
         )
+        parser.add_argument(
+            '--update-batch-size',
+            help='Batch size used to update zones. Updates are run with each retrieval batch.',
+            type=int,
+        )
         parser.add_argument(
             '--dry-run',
             action='store_true',
@@ -34,12 +51,23 @@ class Command(BaseCommand):
             default=True,
         )
 
-    def handle(self, *args, batch_size=1000, remove=True, **kwargs):
-        completed, failed, removed, total = 0, 0, 0, Zone.objects.count()
-        logger.info('{} zones to reorder ({} batches)'.format(total, int(total / batch_size) + 1))
+    def handle(self, *args, offset=0, limit=0, batch_size=1000, update_batch_size=None, remove=True, **kwargs):
+        assert batch_size > 0, 'Batch size must be positive.'
+        assert update_batch_size > 0, 'Update batch size must be positive.'
+
+        assert offset >= 0, 'Offset cannot be negative.'
+        queryset = Zone.objects.order_by('id')
+        if limit:
+            assert limit > 0, 'Limit must be positive.'
+            queryset = queryset[offset:offset+limit]
+        else:
+            queryset = queryset[offset:]
+
+        completed, reordered, failed, removed, total = 0, 0, 0, 0, queryset.count()
+        logger.info('{} zones to reorder ({} batches)'.format(total, ceil(total / batch_size)))
         start_time = default_timer()
         for i in range(0, total, batch_size):
-            batch = Zone.objects.all()[i:i+batch_size]
+            batch = queryset[i:min(i+batch_size, total)]
             reordered_zones = []
 
             for zone in batch:
@@ -64,11 +92,12 @@ class Command(BaseCommand):
                 if zone.polygon != old_polygon:
                     reordered_zones.append(zone)
 
-            Zone.objects.bulk_update(reordered_zones, ['polygon'])
+            Zone.objects.bulk_update(reordered_zones, ['polygon'], batch_size=update_batch_size)
+            reordered += len(reordered_zones)
             completed += batch.count()
 
             percent = (completed / total) * 100
             elapsed = timedelta(seconds=default_timer() - start_time)
-            logger.info('{:6.2f}% completed ({:7d} / {}) in {}, {} failed, {} removed'.format(
-                percent, completed, total, elapsed, failed, removed
+            logger.info('{:6.2f}% completed ({:7d} / {}) in {}, {} reordered, {} failed, {} removed'.format(
+                percent, completed, total, elapsed, reordered, failed, removed
             ))