Improve command to merge paths

GeotrekCE · Jul 25, 2023 · afb1575 · afb1575
1 parent faa2071
commit afb1575
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 31 deletions.
diff --git a/geotrek/core/management/commands/merge_segmented_paths.py b/geotrek/core/management/commands/merge_segmented_paths.py
@@ -33,7 +33,7 @@ def extract_neighbourgs_graph(self, number_of_neighbourgs, extremities=[]):
 
     def try_merge(self, a, b):
         if {a, b} in self.discarded:
-            print(f"├ Already discarded {a} and {b}")
+            self.stdout.write(f"├ Already discarded {a} and {b}")
             return False
         success = 2
         try:
@@ -42,24 +42,20 @@ def try_merge(self, a, b):
             with transaction.atomic():
                 success = patha.merge_path(pathb)
         except Exception:
-            print(f"├ Cannot merge {a} and {b}")
+            self.stdout.write(f"├ Cannot merge {a} and {b}")
             self.discarded.append({a, b})
             return False
-        if success == 2:
-            print(f"├ Cannot merge {a} and {b}")
-            self.discarded.append({a, b})
-            return False
-        elif success == 0:
-            print(f"├ No matching points to merge paths {a} and {b} found")
+        if success == 2 or success == 0:
+            self.stdout.write(f"├ Cannot merge {a} and {b}")
             self.discarded.append({a, b})
             return False
         else:
-            print(f"├ Merged {b} into {a}")
+            self.stdout.write(f"├ Merged {b} into {a}")
             sleep(self.sleeptime)
             return True
 
     def merge_paths_with_one_neighbour(self):
-        print("┌ STEP 1")
+        self.stdout.write("┌ STEP 1")
         neighbours_graph = self.extract_neighbourgs_graph(1)
         successes = 0
         fails = 0
@@ -75,7 +71,7 @@ def merge_paths_with_one_neighbour(self):
         return successes
 
     def merge_paths_with_two_neighbours(self):
-        print("┌ STEP 2")
+        self.stdout.write("┌ STEP 2")
         successes = 0
         neighbours_graph = self.extract_neighbourgs_graph(2)
         mergeables = list(neighbours_graph.keys())
@@ -94,7 +90,7 @@ def merge_paths_with_two_neighbours(self):
         return successes
 
     def merge_paths_with_three_neighbours(self):
-        print("┌ STEP 3")
+        self.stdout.write("┌ STEP 3")
         successes = 0
         neighbours_graph = self.extract_neighbourgs_graph(3)
         mergeables = list(neighbours_graph.keys())
@@ -120,21 +116,21 @@ def handle(self, *args, **options):
         self.discarded = []
         paths_before = Path.include_invisible.count()
 
-        print("\n")
-        print(datetime.now())
+        self.stdout.write("\n")
+        self.stdout.write(str(datetime.now()))
 
         first_step_successes = self.merge_paths_with_one_neighbour()
-        print(f"└ {first_step_successes} merges")
+        self.stdout.write(f"└ {first_step_successes} merges")
         total_successes += first_step_successes
 
         second_step_successes = self.merge_paths_with_two_neighbours()
-        print(f"└ {second_step_successes} merges")
+        self.stdout.write(f"└ {second_step_successes} merges")
         total_successes += second_step_successes
 
         third_step_successes = self.merge_paths_with_three_neighbours()
-        print(f"└ {third_step_successes} merges")
+        self.stdout.write(f"└ {third_step_successes} merges")
         total_successes += third_step_successes
 
         paths_after = Path.include_invisible.count()
-        print(f"\n--- RAN {total_successes} MERGES - FROM {paths_before} TO {paths_after} PATHS ---\n")
-        print(datetime.now())
+        self.stdout.write(f"\n--- RAN {total_successes} MERGES - FROM {paths_before} TO {paths_after} PATHS ---\n")
+        self.stdout.write(str(datetime.now()))
diff --git a/geotrek/core/tests/test_commands.py b/geotrek/core/tests/test_commands.py
@@ -661,6 +661,8 @@ def setUpTestData(cls):
         cls.p14 = Path.objects.create(geom=geom_14)
         geom_15 = LineString((5, 3), (4, 1))
         cls.p15 = Path.objects.create(geom=geom_15)
+        geom_16 = LineString((7, 5), (8, 5), (8, 6), (7, 6), (7, 5))
+        cls.p16 = Path.objects.create(geom=geom_16)
 
     @override_settings(PATH_SNAPPING_DISTANCE=0, PATH_MERGE_SNAPPING_DISTANCE=0)
     def test_find_and_merge_paths(self):
@@ -670,16 +672,16 @@ def test_find_and_merge_paths(self):
         #                |                             |
         #                |  p4                         |  p13
         #                |                             |
-        #                +                             +
-        #                |
-        #                |  p10
-        #          p11   |
-        #         +------+------+ p15
+        #                +                             +-------
+        #                |                             |       |
+        #                |  p10                        |   p16 |
+        #          p11   |                             |       |
+        #         +------+------+ p15                  --------
         #                |
         #                |  p12
         #                |
         #                +
-        self.assertEqual(Path.objects.count(), 15)
+        self.assertEqual(Path.objects.count(), 16)
         output = StringIO()
         call_command('merge_segmented_paths', stdout=output)
         # After call
@@ -688,13 +690,13 @@ def test_find_and_merge_paths(self):
         #                |                             |
         #                |  p4                         |  p13
         #                |                             |
-        #                +                             +
-        #                |
-        #                |  p10
-        #          p11   |
-        #         +------+------+ p15
+        #                +                             +-------
+        #                |                             |       |
+        #                |  p10                        |   p16 |
+        #          p11   |                             |       |
+        #         +------+------+ p15                  --------
         #                |
         #                |  p12
         #                |
         #                +
-        self.assertEqual(Path.objects.count(), 9)
+        self.assertEqual(Path.objects.count(), 10)