diff --git a/cadquery/selectors.py b/cadquery/selectors.py index 72362c8..ada9bd5 100644 --- a/cadquery/selectors.py +++ b/cadquery/selectors.py @@ -40,6 +40,18 @@ class Selector(object): """ return objectList + def __and__(self, other): + return AndSelector(self, other) + + def __add__(self, other): + return SumSelector(self, other) + + def __sub__(self, other): + return SubtractSelector(self, other) + + def __neg__(self): + return InverseSelector(self) + class NearestToPointSelector(Selector): """ Selects object nearest the provided point. @@ -299,6 +311,57 @@ class DirectionMinMaxSelector(Selector): else: return [ min(objectList,key=distance) ] +class BinarySelector(Selector): + """ + Base class for selectors that operates with two other + selectors. Subclass must implement the :filterResults(): method. + """ + def __init__(self, left, right): + self.left = left + self.right = right + + def filter(self, objectList): + return self.filterResults(self.left.filter(objectList), + self.right.filter(objectList)) + + def filterResults(self, r_left, r_right): + raise NotImplementedError + +class AndSelector(BinarySelector): + """ + Intersection selector. Returns objects that is selected by both selectors. + """ + def filterResults(self, r_left, r_right): + # return intersection of lists + return list(set(r_left) & set(r_right)) + +class SumSelector(BinarySelector): + """ + Union selector. Returns the sum of two selectors results. + """ + def filterResults(self, r_left, r_right): + # return the union (no duplicates) of lists + return list(set(r_left + r_right)) + +class SubtractSelector(BinarySelector): + """ + Difference selector. Substract results of a selector from another + selectors results. + """ + def filterResults(self, r_left, r_right): + return list(set(r_left) - set(r_right)) + +class InverseSelector(Selector): + """ + Inverts the selection of given selector. In other words, selects + all objects that is not selected by given selector. + """ + def __init__(self, selector): + self.selector = selector + + def filter(self, objectList): + # note that Selector() selects everything + return SubtractSelector(Selector(), self.selector).filter(objectList) class StringSyntaxSelector(Selector): """ diff --git a/tests/TestCQSelectors.py b/tests/TestCQSelectors.py index 678766a..53f5780 100644 --- a/tests/TestCQSelectors.py +++ b/tests/TestCQSelectors.py @@ -274,6 +274,63 @@ class TestCQSelectors(BaseTest): fl = c.faces(selectors.BoxSelector((-0.1, 0.4, -0.1), (1.1, 1.1, 1.1), True)).vals() self.assertEqual(1, len(fl)) + def testAndSelector(self): + c = CQ(makeUnitCube()) + + S = selectors.StringSyntaxSelector + BS = selectors.BoxSelector + + el = c.edges(selectors.AndSelector(S('|X'), BS((-2,-2,0.1), (2,2,2)))).vals() + self.assertEqual(2, len(el)) + + # test 'and' (intersection) operator + el = c.edges(S('|X') & BS((-2,-2,0.1), (2,2,2))).vals() + self.assertEqual(2, len(el)) + + def testSumSelector(self): + c = CQ(makeUnitCube()) + + S = selectors.StringSyntaxSelector + + fl = c.faces(selectors.SumSelector(S(">Z"), S("Z") + S("X"))).vals() + self.assertEqual(3, len(fl)) + + # test the subtract operator + fl = c.faces(S("#Z") - S(">X")).vals() + self.assertEqual(3, len(fl)) + + def testInverseSelector(self): + c = CQ(makeUnitCube()) + + S = selectors.StringSyntaxSelector + + fl = c.faces(selectors.InverseSelector(S('>Z'))).vals() + self.assertEqual(5, len(fl)) + el = c.faces('>Z').edges(selectors.InverseSelector(S('>X'))).vals() + self.assertEqual(3, len(el)) + + # test invert operator + fl = c.faces(-S('>Z')).vals() + self.assertEqual(5, len(fl)) + el = c.faces('>Z').edges(-S('>X')).vals() + self.assertEqual(3, len(el)) + def testFaceCount(self): c = CQ(makeUnitCube()) self.assertEqual( 6, c.faces().size() )