transformers: de-duplicate fixup code

1 files changed, 25 insertions, 18 deletions

diff --git a/fietsboek/transformers/__init__.py b/fietsboek/transformers/__init__.py
index f98a545..f6318d7 100644
--- a/fietsboek/transformers/__init__.py
+++ b/fietsboek/transformers/__init__.py
@@ -10,7 +10,7 @@ function to load and apply transformers.
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Callable, Iterable, Mapping
 from itertools import islice
 from typing import Literal, NamedTuple, TypeVar
 
@@ -155,7 +155,7 @@ class FixNullElevation(Transformer):
     def parameters(self, value):
         pass
 
-    def execute(self, gpx):
+    def execute(self, gpx: GPX):
         def all_points():
             return (
                 point
@@ -172,30 +172,37 @@ class FixNullElevation(Transformer):
                 for point in reversed(segment.points)
             )
 
-        max_slope = 1.0
+        # First, from the front
+        self.fixup(all_points)
+        # Then, from the back
+        self.fixup(rev_points)
 
-        # First, from the front, find the first point with non-zero elevation (or low enough slope)
-        bad_until = 0
-        final_elevation = 0
-        for i, (point, next_point) in enumerate(zip(all_points(), islice(all_points(), 1, None))):
-            if point.elevation != 0.0 and self.slope(point, next_point) < max_slope:
-                bad_until = i
-                final_elevation = point.elevation
-                break
+    @classmethod
+    def fixup(cls, points: Callable[[], Iterable[GPXTrackPoint]]):
+        """Fixes the given GPX points.
 
-        for point in islice(all_points(), None, bad_until):
-            point.elevation = final_elevation
+        This iterates over the points and checks for the first point that has a
+        non-zero elevation, and a slope that doesn't exceed 100%. All previous
+        points will have their elevation adjusted to match this first "good
+        point".
+
+        :param points: A function that generates the iterable of points.
+        """
+        max_slope = 1.0
 
-        # Second, from the back
         bad_until = 0
-        final_elevation = 0
-        for i, (point, prev_point) in enumerate(zip(rev_points(), islice(rev_points(), 1, None))):
-            if point.elevation != 0.0 and self.slope(point, prev_point) < max_slope:
+        final_elevation = 0.0
+        for i, (point, next_point) in enumerate(zip(points(), islice(points(), 1, None))):
+            if (
+                point.elevation is not None
+                and point.elevation != 0.0
+                and cls.slope(point, next_point) < max_slope
+            ):
                 bad_until = i
                 final_elevation = point.elevation
                 break
 
-        for point in islice(rev_points(), None, bad_until):
+        for point in islice(points(), None, bad_until):
             point.elevation = final_elevation
 
     @staticmethod