Add decorator to run function on buffered las

pdaltools/las_add_buffer.py

@@ -2,7 +2,9 @@
 import logging
 import os
 import tempfile
-from typing import List
+from functools import wraps
+from pathlib import Path
+from typing import Callable, List
 
 import pdal
 
@@ -90,7 +92,7 @@ def las_merge_and_crop(
         tile_filename (str): full path to the queried LIDAR tile
         bounds (List): 2D bounding box to crop to : provided as ([xmin, xmax], [ymin, ymax])
         output_filename (str): full path to the saved cropped tile
-        spatial_ref (_type_, optional): spatial reference for the writer. Defaults to "EPSG:2154".
+        spatial_ref (str, optional): spatial reference for the writer. Defaults to "EPSG:2154".
         tile_width (int, optional): width of tiles in meters (usually 1000m). Defaults to 1000.
         tile_coord_scale (int, optional): scale used in the filename to describe coordinates in meters.
         Defaults to 1000.
@@ -157,6 +159,82 @@ def remove_points_from_buffer(input_file: str, output_file: str):
         remove_dimensions_from_las(tmp_las.name, dimensions=[ORIGINAL_TILE_TAG], output_las=output_file)
 
 
+def run_on_buffered_las(
+    buffer_width: int, spatial_ref: str, tile_width: int = 1000, tile_coord_scale: int = 1000
+) -> Callable:
+    """Decorator to apply a function that takes a las/laz as input and returns a las/laz output
+    on an input with an additional buffer, then remove the buffer points from the output
+
+    The first argument of the decorated function must be an input path
+    The second argument of the decorated function must be an output path
+
+    The buffer is added by merging lidar tiles around the queried tile and crop them based
+    on their filenames
+
+    Args:
+        buffer_width (int): width of the border to add to the tile (in meters)
+        spatial_ref (str): spatial reference for the writer. Example: "EPSG:2154".
+        tile_width (int, optional): width of tiles in meters (usually 1000m). Defaults to 1000.
+        tile_coord_scale (int, optional): scale used in the filename to describe coordinates in meters.
+        Defaults to 1000.
+
+    Raises:
+        FileNotFoundError: when the first argument of the decorated function is not an existing
+            file
+        FileNotFoundError:  when the second argument of the decorated function is not a path
+            with an existing parent folder
+
+    Returns:
+        Callable: decorated function
+    """
+    """Decorator to run a function that takes a las as input and returns a las output
+    on a las with an additional buffer, then remove the buffer points from the buffer points
+
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            input_file = args[0]
+            output_file = args[1]
+            if not Path(input_file).is_file():
+                raise FileNotFoundError(
+                    f"File {args[0]} not found. The first argument of a function decorated by "
+                    "'run_on_buffered_las' is expected to be the path to an existing input file."
+                )
+
+            if not Path(output_file).parent.is_dir():
+                raise FileNotFoundError(
+                    f"Parent folder for file {args[1]} not found. The second argument of a function "
+                    "decorated by 'run_on_buffered_las' is expected to be the path to an output "
+                    "file in an existing folder"
+                )
+
+            with (
+                tempfile.NamedTemporaryFile(suffix="_buffered_input.laz", dir=".") as buf_in,
+                tempfile.NamedTemporaryFile(suffix="_buffered_output.laz", dir=".") as buf_out,
+            ):
+                create_las_with_buffer(
+                    Path(input_file).parent,
+                    input_file,
+                    buf_in.name,
+                    buffer_width=buffer_width,
+                    spatial_ref=spatial_ref,
+                    tile_width=tile_width,
+                    tile_coord_scale=tile_coord_scale,
+                    tag_original_tile=True,
+                )
+                func(buf_in.name, buf_out.name, *args[2:], **kwargs)
+
+                remove_points_from_buffer(buf_out.name, output_file)
+
+            return
+
+        return wrapper
+
+    return decorator
+
+
 def parse_args():
     parser = argparse.ArgumentParser("Add a buffer to a las tile by stitching with its neighbors")
     parser.add_argument(

test/test_las_add_buffer.py

@@ -10,7 +10,11 @@
 from pdaltools.count_occurences.count_occurences_for_attribute import (
     compute_count_one_file,
 )
-from pdaltools.las_add_buffer import create_las_with_buffer, remove_points_from_buffer
+from pdaltools.las_add_buffer import (
+    create_las_with_buffer,
+    remove_points_from_buffer,
+    run_on_buffered_las,
+)
 
 TEST_PATH = os.path.dirname(os.path.abspath(__file__))
 TMP_PATH = os.path.join(TEST_PATH, "tmp")
@@ -127,7 +131,6 @@ def test_create_las_with_buffer_with_tag():
     # check number of points with the additional tag
     assert get_nb_points(output_file) > input_nb_points
     count_points_from_original = compute_count_one_file(output_file, attribute="is_in_original")
-    print(count_points_from_original)
     assert count_points_from_original["1"] == input_nb_points
 
     # Check boundaries
@@ -170,6 +173,7 @@ def test_remove_points_from_buffer():
 
     remove_points_from_buffer(buffered_file, output_file)
     assert os.path.isfile(output_file)
+    assert get_nb_points(buffered_file) > get_nb_points(input_file)
     assert get_nb_points(output_file) == get_nb_points(input_file)
 
     # Check output dimensions are the same as input dimensions
@@ -183,6 +187,28 @@ def test_remove_points_from_buffer():
     assert_header_info_are_similar(output_file, input_file)
 
 
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
-    test_create_las_with_buffer()
+def test_run_on_buffered_las():
+    # Dummy example with copy only
+    buffer_width = 5
+    tile_width = 50
+    tile_coord_scale = 10
+    spatial_ref = "EPSG:2154"
+
+    input_file = os.path.join(INPUT_DIR, "test_data_77055_627755_LA93_IGN69.laz")
+    output_dir = os.path.join(TMP_PATH, "run_with_buffer")
+    os.makedirs(output_dir)
+    output_file = os.path.join(output_dir, "copied.laz")
+    decorated_copy = run_on_buffered_las(
+        buffer_width, spatial_ref=spatial_ref, tile_width=tile_width, tile_coord_scale=tile_coord_scale
+    )(shutil.copy)
+
+    decorated_copy(input_file, output_file)
+
+    # Check output dimensions are the same as input dimensions
+    output_dimensions = tu.get_pdal_infos_summary(output_file)["summary"]["dimensions"]
+    input_dimensions = tu.get_pdal_infos_summary(input_file)["summary"]["dimensions"]
+    assert output_dimensions == input_dimensions
+
+    # Check las content
+    assert get_nb_points(output_file) == get_nb_points(input_file)
+    assert compute_count_one_file(output_file) == compute_count_one_file(input_file)