Implement ProductSet.__iter__ using product_cantor

Closes sympy/sympy#10497
skirpichev · Jun 15, 2016 · 6b8e1f4 · 6b8e1f4
1 parent 8c9d01c
commit 6b8e1f4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/sympy/sets/sets.py b/sympy/sets/sets.py
@@ -622,8 +622,9 @@ def is_iterable(self):
         return all(set.is_iterable for set in self.sets)
 
     def __iter__(self):
+        from sympy.utilities.iterables import cantor_product
         if self.is_iterable:
-            return itertools.product(*self.sets)
+            return cantor_product(*self.sets)
         else:
             raise TypeError("Not all constituent sets are iterable")
 

diff --git a/sympy/sets/tests/test_fancysets.py b/sympy/sets/tests/test_fancysets.py
@@ -279,3 +279,9 @@ def test_ImageSet_simplification():
     assert (imageset(Lambda(n, sin(n)),
                      imageset(Lambda(m, tan(m)), S.Integers)) ==
             imageset(Lambda(m, sin(tan(m))), S.Integers))
+
+
+def test_issue_10497():
+    # __iter__ for infinite product set
+    assert (list(itertools.islice(iter(S.Integers**2), 7)) ==
+            [(0, 0), (0, 1), (1, 0), (1, 1), (0, -1), (1, -1), (-1, 0)])