From f077d72819529cee0dfefb58dc6ee2ee624fa1ae Mon Sep 17 00:00:00 2001 From: gumyr Date: Thu, 23 Jan 2025 09:53:15 -0500 Subject: [PATCH] Added sort_by lambda Issue#485 --- src/build123d/topology/shape_core.py | 12 ++++++++++-- tests/test_direct_api/test_shape_list.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/build123d/topology/shape_core.py b/src/build123d/topology/shape_core.py index b30e596..42fb793 100644 --- a/src/build123d/topology/shape_core.py +++ b/src/build123d/topology/shape_core.py @@ -2624,7 +2624,9 @@ class ShapeList(list[T]): return ShapeList([s for shape in self for s in shape.solids()]) # type: ignore def sort_by( - self, sort_by: Axis | Edge | Wire | SortBy = Axis.Z, reverse: bool = False + self, + sort_by: Axis | Callable[[T], K] | Edge | Wire | SortBy = Axis.Z, + reverse: bool = False, ) -> ShapeList[T]: """sort by @@ -2639,7 +2641,11 @@ class ShapeList(list[T]): ShapeList: sorted list of objects """ - if isinstance(sort_by, Axis): + if callable(sort_by): + # If a callable is provided, use it directly as the key + objects = sorted(self, key=sort_by, reverse=reverse) + + elif isinstance(sort_by, Axis): if sort_by.wrapped is None: raise ValueError("Cannot sort by an empty axis") assert sort_by.location is not None @@ -2702,6 +2708,8 @@ class ShapeList(list[T]): key=lambda obj: obj.volume, # type: ignore reverse=reverse, ) + else: + raise ValueError("Invalid sort_by criteria provided") return ShapeList(objects) diff --git a/tests/test_direct_api/test_shape_list.py b/tests/test_direct_api/test_shape_list.py index 45ec7eb..98139d8 100644 --- a/tests/test_direct_api/test_shape_list.py +++ b/tests/test_direct_api/test_shape_list.py @@ -79,6 +79,20 @@ class TestShapeList(unittest.TestCase): faces = Solid.make_box(1, 2, 3).faces() < SortBy.AREA self.assertAlmostEqual(faces[-1].area, 2, 5) + def test_sort_by_lambda(self): + c = Solid.make_cone(2, 1, 2) + flat_faces = c.faces().filter_by(GeomType.PLANE) + sorted_flat_faces = flat_faces.sort_by(lambda f: f.area) + smallest = sorted_flat_faces[0] + largest = sorted_flat_faces[-1] + + self.assertAlmostEqual(smallest.area, math.pi * 1**2, 5) + self.assertAlmostEqual(largest.area, math.pi * 2**2, 5) + + def test_sort_by_invalid(self): + with self.assertRaises(ValueError): + Solid.make_box(1, 1, 1).faces().sort_by(">Z") + def test_filter_by_geomtype(self): non_planar_faces = ( Solid.make_cylinder(1, 1).faces().filter_by(GeomType.PLANE, reverse=True)