Allow location to multiply a list of shapes

This commit is contained in:
Dmytry Lavrov 2025-09-28 12:49:54 -05:00
parent fed77612c0
commit 6b8409ab33
2 changed files with 8 additions and 8 deletions

View file

@ -1706,8 +1706,8 @@ class Location:
def __mul__(self, other: Iterable[Location]) -> list[Location]: ...
def __mul__(
self, other: Shape | Location | Iterable[Location]
) -> Shape | Location | list[Location]:
self, other: Shape | Location | Iterable[Location | Shape | Iterable]
) -> Shape | Location | list[Location | Shape | list]:
"""Combine locations"""
if self.wrapped is None:
raise ValueError("Cannot move a shape at an empty location")
@ -1742,14 +1742,10 @@ class Location:
raise ValueError("Can't multiply by empty location")
return Location(self.wrapped * other.wrapped)
# other is a list of Locations
# other is a list
if isinstance(other, Iterable):
others = list(other)
if not all(isinstance(o, Location) for o in others):
raise ValueError("other must be a list of Locations")
if any(o.wrapped is None for o in others):
raise ValueError("Can't multiple by empty Locations")
return [Location(self.wrapped * loc.wrapped) for loc in others]
return [self * loc for loc in others]
raise ValueError(f"Invalid input {other}")

View file

@ -795,6 +795,10 @@ class AlgebraTests(unittest.TestCase):
class LocationTests(unittest.TestCase):
def test_list_mul(self):
a = Rot(45) * (GridLocations(10, 10, 2, 2) * Cylinder(2, 3))
self.assertEqual(len(a), 4)
def test_wheel(self):
plane = Plane.ZX