Refactor out the merging logic

diff --git a/src/climbing/models.py b/src/climbing/models.py
index 60aa717..83d437f 100644 (file)
--- a/src/climbing/models.py
+++ b/src/climbing/models.py
@@ -1,6 +1,8 @@
 from django.contrib.auth.models import User
 from django.db import models
 
+from core import utils
+
 class Area(models.Model):
     name = models.CharField(max_length=64)
     notes = models.TextField(blank=True, null=True)
@@ -10,46 +12,11 @@ class Area(models.Model):
 
     @property
     def sub_areas(self):
-        crags = iter(self.crags.order_by('name'))
-        clusters = iter(self.clusters.order_by('name'))
-
-        try:
-            crag = next(crags)
-        except StopIteration:
-            crag = None
-
-        try:
-            cluster = next(clusters)
-        except StopIteration:
-            cluster = None
-
-        while crag and cluster:
-            if crag.name < cluster.name:
-                yield crag
-                try:
-                    crag = next(crags)
-                except:
-                    crag = None
-            else:
-                yield cluster
-                try:
-                    cluster = next(clusters)
-                except:
-                    cluster = None
-
-        while crag:
-            yield crag
-            try:
-                crag = next(crags)
-            except:
-                crag = None
-
-        while cluster:
-            yield cluster
-            try:
-                cluster = next(clusters)
-            except:
-                cluster = None
+        return utils.merge(
+            self.crags.order_by('name'),
+            self.clusters.order_by('name'),
+            'name',
+        )
 
 class Crag(models.Model):
     area = models.ForeignKey(

diff --git a/src/core/utils.py b/src/core/utils.py
new file mode 100644 (file)
index 0000000..679de84
--- /dev/null
+++ b/src/core/utils.py
@@ -0,0 +1,54 @@
+def merge(xs:iter, ys:iter, x_key, y_key=None):
+    if y_key is None:
+        y_key = x_key
+
+    if isinstance(x_key, str):
+        x_attr = x_key
+        x_key = lambda x: getattr(x, x_attr)
+
+    if isinstance(y_key, str):
+        y_attr = y_key
+        y_key = lambda y: getattr(y, y_attr)
+
+    xs = iter(xs)
+    ys = iter(ys)
+
+    try:
+        x = next(xs)
+    except StopIteration:
+        x = None
+
+    try:
+        y = next(ys)
+    except StopIteration:
+        y = None
+
+    while x and y:
+        if x_key(x) < y_key(y):
+            yield x
+            try:
+                x = next(xs)
+            except StopIteration:
+                x = None
+        else:
+            yield y
+            try:
+                y = next(ys)
+            except StopIteration:
+                y = None
+
+    while x:
+        yield x
+        try:
+            x = next(xs)
+        except StopIteration:
+            x = None
+
+    while y:
+        yield y
+        try:
+            y = next(ys)
+        except StopIteration:
+            y = None
+
+