Refactor out the merging logic

[climbing.kerkeslager.com] / src / core / utils.py

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
+
+