Trac #32315: support enumerated sets

src/sage/sets/condition_set.py

@@ -16,6 +16,7 @@
 from sage.structure.parent import Parent, Set_generic
 from sage.structure.unique_representation import UniqueRepresentation
 from sage.categories.sets_cat import Sets
+from sage.categories.enumerated_sets import EnumeratedSets
 from sage.misc.cachefunc import cached_method
 from sage.misc.misc import _stable_uniq
 from sage.symbolic.expression import is_Expression
@@ -24,50 +25,28 @@
 
 from .set import Set, Set_base, Set_boolean_operators, Set_add_sub_operators
 
-class ConditionSet(Set_generic, Set_base, Set_boolean_operators, Set_add_sub_operators,
-                   UniqueRepresentation):
+class ConditionSet_generic(Set_generic, Set_base, Set_boolean_operators,
+                           Set_add_sub_operators, UniqueRepresentation):
     r"""
     Set of elements of a universe that satisfy given predicates
 
+    .. NOTE::
+
+        This is not intended to be created directly by the user. Please use
+        :meth:`ConditionSet` instead.
+
     INPUT:
 
     - ``universe`` -- a set
 
     - ``*predicates`` -- callables
 
-    - ``vars`` or ``names`` -- (default: inferred from ``predicates`` if any predicate is
-      an element of a :class:`~sage.symbolic.callable.CallableSymbolicExpressionRing_class`)
-      variables or names of variables
+    - ``names`` -- a tuple of names of variables
 
-    - ``category`` -- (default: inferred from ``universe``) a category
+    - ``category`` -- a category
 
     EXAMPLES::
 
-        sage: Evens = ConditionSet(ZZ, is_even); Evens
-        { x ∈ Integer Ring : <function is_even at 0x...>(x) }
-        sage: 2 in Evens
-        True
-        sage: 3 in Evens
-        False
-        sage: 2.0 in Evens
-        True
-
-        sage: Odds = ConditionSet(ZZ, is_odd); Odds
-        { x ∈ Integer Ring : <function is_odd at 0x...>(x) }
-        sage: EvensAndOdds = Evens | Odds; EvensAndOdds
-        Set-theoretic union of
-         { x ∈ Integer Ring : <function is_even at 0x...>(x) } and
-         { x ∈ Integer Ring : <function is_odd at 0x...>(x) }
-        sage: 5 in EvensAndOdds
-        True
-        sage: 7/2 in EvensAndOdds
-        False
-
-        sage: var('y')
-        y
-        sage: SmallOdds = ConditionSet(ZZ, is_odd, abs(y) <= 11, vars=[y]); SmallOdds
-        { y ∈ Integer Ring : abs(y) <= 11, <function is_odd at 0x...>(y) }
-
         sage: P = polytopes.cube(); P
         A 3-dimensional polyhedron in ZZ^3 defined as the convex hull of 8 vertices
         sage: P.rename("P")
@@ -78,6 +57,8 @@ class ConditionSet(Set_generic, Set_base, Set_boolean_operators, Set_add_sub_ope
         sage: vector([1, 1, 1]) in P_inter_B
         False
 
+    ::
+
         sage: predicate(x, y, z) = sqrt(x^2 + y^2 + z^2) < 1.2; predicate
         (x, y, z) |--> sqrt(x^2 + y^2 + z^2) < 1.20000000000000
         sage: P_inter_B_again = ConditionSet(P, predicate); P_inter_B_again
@@ -87,92 +68,12 @@ class ConditionSet(Set_generic, Set_base, Set_boolean_operators, Set_add_sub_ope
         sage: vector([1, 1, 1]) in P_inter_B_again
         False
 
-    Using ``ConditionSet`` without predicates provides a way of attaching variable names
-    to a set::
-
-        sage: Z3 = ConditionSet(ZZ^3, vars=['x', 'y', 'z']); Z3
-        { (x, y, z) ∈ Ambient free module of rank 3 over the principal ideal domain Integer Ring }
-        sage: Z3.variable_names()
-        ('x', 'y', 'z')
-        sage: Z3.arguments()
-        (x, y, z)
-
-        sage: Q4.<a, b, c, d> = ConditionSet(QQ^4); Q4
-        { (a, b, c, d) ∈ Vector space of dimension 4 over Rational Field }
-        sage: Q4.variable_names()
-        ('a', 'b', 'c', 'd')
-        sage: Q4.arguments()
-        (a, b, c, d)
-
     TESTS::
 
         sage: TestSuite(P_inter_B).run(skip='_test_pickling')  # cannot pickle lambdas
         sage: TestSuite(P_inter_B_again).run()
-    """
-    @staticmethod
-    def __classcall_private__(cls, universe, *predicates, vars=None, names=None, category=None):
-        r"""
-        Normalize init arguments.
-
-        TESTS::
-
-            sage: ConditionSet(ZZ, names=["x"]) is ConditionSet(ZZ, names=x)
-            True
-            sage: ConditionSet(RR, x > 0, names=x) is ConditionSet(RR, (x > 0).function(x))
-            True
-        """
-        if category is None:
-            category = Sets()
-        if isinstance(universe, Parent):
-            if universe in Sets().Finite():
-                category = category & Sets().Finite()
-
-        if vars is not None:
-            if names is not None:
-                raise ValueError('cannot use names and vars at the same time; they are aliases')
-            names, vars = vars, None
-
-        if names is not None:
-            names = normalize_names(-1, names)
-
-        callable_symbolic_predicates = []
-        other_predicates = []
-
-        for predicate in predicates:
-            if is_CallableSymbolicExpression(predicate):
-                if names is None:
-                    names = tuple(str(var) for var in predicate.args())
-                elif len(names) != len(predicate.args()):
-                    raise TypeError('mismatch in number of arguments')
-                if vars is None:
-                    vars = predicate.args()
-                callable_symbolic_predicates.append(predicate)
-            elif is_Expression(predicate):
-                if names is None:
-                    raise TypeError('use callable symbolic expressions or provide variable names')
-                if vars is None:
-                    vars = tuple(SR.var(name) for name in names)
-                callable_symbolic_predicates.append(predicate.function(*vars))
-            else:
-                other_predicates.append(predicate)
-
-        predicates = list(_stable_uniq(callable_symbolic_predicates + other_predicates))
-
-        if not other_predicates and not callable_symbolic_predicates:
-            if names is None and category is None:
-                # No conditions, no variable names, no category, just use Set.
-                return Set(universe)
-
-        if any(predicate.args() != vars
-               for predicate in callable_symbolic_predicates):
-            # TODO: Implement safe renaming of the arguments of a callable symbolic expressions
-            raise NotImplementedError('all callable symbolic expressions must use the same arguments')
-
-        if names is None:
-            names = ("x",)
-        return super().__classcall__(cls, universe, *predicates,
-                                     names=names, category=category)
 
+    """
     def __init__(self, universe, *predicates, names=None, category=None):
         r"""
         TESTS::
@@ -187,7 +88,8 @@ def __init__(self, universe, *predicates, names=None, category=None):
         if isinstance(universe, Parent):
             facade = universe
         super().__init__(facade=facade, category=category,
-                         names=names, normalize=False) # names already normalized by classcall
+                         names=names, normalize=False) # names already normalized
+                                                       # by factory method
 
     def _first_ngens(self, n):
         r"""
@@ -470,8 +372,186 @@ def intersection(self, X):
             sage: SmallMirrorUniverse & SmallOblongUniverse
             { (y, x) ∈ Vector space of dimension 2 over Rational Field : 3*x^2 + y^2 <= 42 }
         """
-        if isinstance(X, ConditionSet):
+        if isinstance(X, ConditionSet_generic):
             return ConditionSet(self.ambient().intersection(X.ambient()),
                                 *(self._predicates + X._predicates),
                                 vars=self.arguments())
         return super().intersection(X)
+
+class ConditionSet_enumerated(ConditionSet_generic):
+    r"""
+    Enumerated set of elements of an enumerated universe that satisfy given
+    predicates
+
+    .. NOTE::
+
+        This is not intended to be created directly by the user. Please use
+        :meth:`ConditionSet` instead.
+
+    INPUT:
+
+    - ``universe`` -- a set
+
+    - ``*predicates`` -- callables
+
+    - ``names`` -- a tuple of names of variables
+
+    - ``category`` -- a category
+
+    EXAMPLES::
+
+        sage: R = IntegerModRing(8)
+        sage: R_primes = ConditionSet(R, is_prime); R_primes
+        { x ∈ Ring of integers modulo 8 : <function is_prime at 0x...>(x) }
+        sage: R_primes.is_finite()
+        True
+        sage: list(R_primes)
+        [2, 6]
+
+    ::
+
+        sage: odds = ConditionSet(ZZ, is_odd); odds
+        { x ∈ Integer Ring : <function is_odd at 0x...>(x) }
+        sage: list(odds.iterator_range(stop=6))
+        [1, -1, 3, -3, 5, -5]
+
+    """
+    def __iter__(self):
+        r"""
+        Return the iterator of ``self``.
+
+        TESTS::
+
+            sage: Evens = ConditionSet(ZZ, is_even); Evens
+            { x ∈ Integer Ring : <function is_even at 0x...>(x) }
+            sage: list(Evens.iterator_range(stop=5))
+            [0, 2, -2, 4, -4]
+
+        """
+        for x in self._universe:
+            if x in self:
+                yield x
+
+def ConditionSet(universe, *predicates, vars=None, names=None, category=None):
+    r"""
+    INPUT:
+
+    - ``universe`` -- a set
+
+    - ``*predicates`` -- callables
+
+    - ``vars`` or ``names`` -- (default: inferred from ``predicates`` if any predicate is
+      an element of a :class:`~sage.symbolic.callable.CallableSymbolicExpressionRing_class`)
+      variables or names of variables
+
+    - ``category`` -- (default: inferred from ``universe``) a category
+
+    EXAMPLES::
+
+        sage: Evens = ConditionSet(ZZ, is_even); Evens
+        { x ∈ Integer Ring : <function is_even at 0x...>(x) }
+        sage: 2 in Evens
+        True
+        sage: 3 in Evens
+        False
+        sage: 2.0 in Evens
+        True
+
+    ::
+
+        sage: Odds = ConditionSet(ZZ, is_odd); Odds
+        { x ∈ Integer Ring : <function is_odd at 0x...>(x) }
+        sage: EvensAndOdds = Evens | Odds; EvensAndOdds
+        Set-theoretic union of
+         { x ∈ Integer Ring : <function is_even at 0x...>(x) } and
+         { x ∈ Integer Ring : <function is_odd at 0x...>(x) }
+        sage: 5 in EvensAndOdds
+        True
+        sage: 7/2 in EvensAndOdds
+        False
+
+    ::
+
+        sage: var('y')
+        y
+        sage: SmallOdds = ConditionSet(ZZ, is_odd, abs(y) <= 11, vars=[y]); SmallOdds
+        { y ∈ Integer Ring : abs(y) <= 11, <function is_odd at 0x...>(y) }
+
+    Using ``ConditionSet`` without predicates provides a way of attaching variable names
+    to a set::
+
+        sage: Z3 = ConditionSet(ZZ^3, vars=['x', 'y', 'z']); Z3
+        { (x, y, z) ∈ Ambient free module of rank 3 over the principal ideal domain Integer Ring }
+        sage: Z3.variable_names()
+        ('x', 'y', 'z')
+        sage: Z3.arguments()
+        (x, y, z)
+
+    ::
+
+        sage: Q4.<a, b, c, d> = ConditionSet(QQ^4); Q4
+        { (a, b, c, d) ∈ Vector space of dimension 4 over Rational Field }
+        sage: Q4.variable_names()
+        ('a', 'b', 'c', 'd')
+        sage: Q4.arguments()
+        (a, b, c, d)
+
+    """
+    if category is None:
+        category = Sets()
+    if isinstance(universe, Parent):
+        if universe in Sets().Finite():
+            category = category & Sets().Finite()
+
+    if vars is not None:
+        if names is not None:
+            raise ValueError(
+                'cannot use names and vars at the same time; they are aliases')
+        names, vars = vars, None
+
+    if names is not None:
+        names = normalize_names(-1, names)
+
+    callable_symbolic_predicates = []
+    other_predicates = []
+
+    for predicate in predicates:
+        if is_CallableSymbolicExpression(predicate):
+            if names is None:
+                names = tuple(str(var) for var in predicate.args())
+            elif len(names) != len(predicate.args()):
+                raise TypeError('mismatch in number of arguments')
+            if vars is None:
+                vars = predicate.args()
+            callable_symbolic_predicates.append(predicate)
+        elif is_Expression(predicate):
+            if names is None:
+                raise TypeError('use callable symbolic expressions or provide variable names')
+            if vars is None:
+                vars = tuple(SR.var(name) for name in names)
+            callable_symbolic_predicates.append(predicate.function(*vars))
+        else:
+            other_predicates.append(predicate)
+
+    predicates = list(_stable_uniq(callable_symbolic_predicates + other_predicates))
+
+    if not other_predicates and not callable_symbolic_predicates:
+        if names is None and category is None:
+            # No conditions, no variable names, no category, just use Set.
+            return Set(universe)
+
+    if any(predicate.args() != vars
+           for predicate in callable_symbolic_predicates):
+        # TODO: Implement safe renaming of the arguments of a callable symbolic expressions
+        raise NotImplementedError('all callable symbolic expressions must use the same arguments')
+
+    if names is None:
+        names = ("x",)
+
+    if universe in EnumeratedSets():
+        category &= EnumeratedSets()
+        return ConditionSet_enumerated(universe, *predicates, names=names,
+                                       category=category)
+
+    return ConditionSet_generic(universe, *predicates, names=names,
+                                category=category)