Added sort_by lambda Issue#485

This commit is contained in:
gumyr 2025-01-23 09:53:15 -05:00
parent 4274f78f6a
commit f077d72819
2 changed files with 24 additions and 2 deletions

View file

@ -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)

View file

@ -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)