Allow iteration of all the custom itertools types.issue9

3 files changed, 104 insertions, 9 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 941e432..359b865 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -20,7 +20,9 @@ Changes
 
 - Fix `issue 9
   <https://github.com/zopefoundation/zope.security/issues/9>`_:
-  iteration of ``itertools.groupby`` objects is now allowed by default.
+  iteration of ``itertools.groupby`` objects is now allowed by
+  default. In addition, iteration of all the custom iterator types
+  defined in itertools are also allowed by default.
 
 4.1.1 (2017-05-17)
 ------------------
diff --git a/src/zope/security/checker.py b/src/zope/security/checker.py
index c781d5f..e5e23c3 100644
--- a/src/zope/security/checker.py
+++ b/src/zope/security/checker.py
@@ -860,7 +860,7 @@ _fixup_zope_interface()
 del _fixup_zope_interface
 
 
-def _fixup_itertools_groupby():
+def _fixup_itertools():
     # itertools.groupby is a built-in custom iterator type introduced
     # in python2.4. It should have the same checker as other built-in
     # iterators.
@@ -869,9 +869,9 @@ def _fixup_itertools_groupby():
     # iterator. Its type is not exposed by name, but can be accessed
     # like so: type(list(itertools.groupby([0]))[0][1])
 
-    from itertools import groupby
+    import itertools
 
-    group = groupby([0])
+    group = itertools.groupby([0])
     type_group = type(group)
     if type_group not in _default_checkers:
         _default_checkers[type_group] = _iteratorChecker
@@ -880,8 +880,53 @@ def _fixup_itertools_groupby():
     if type_grouper not in _default_checkers:
         _default_checkers[type_grouper] = _iteratorChecker
 
-_fixup_itertools_groupby()
-del _fixup_itertools_groupby
+    # There are also many other custom types in itertools that need the
+    # same treatment. See a similar list in test_checker.py:test_itertools_checkers
+    pred = lambda x: x
+    iterable = (1, 2, 3)
+    pred_iterable = (pred, iterable)
+    missing_in_py3 = {'ifilter', 'ifilterfalse', 'imap',
+                      'izip', 'izip_longest'}
+    missing_in_py2 = {'zip_longest', 'accumulate', 'compress',
+                      'combinations', 'combinations_with_replacement'}
+    missing = missing_in_py3 if sys.version_info[0] >= 3 else missing_in_py2
+    for func, args in (('count', ()),
+                       ('cycle', ((),)),
+                       ('dropwhile', pred_iterable),
+                       ('ifilter', pred_iterable),
+                       ('ifilterfalse', pred_iterable),
+                       ('imap', pred_iterable),
+                       ('islice', (iterable, 2)),
+                       ('izip', (iterable,)),
+                       ('izip_longest', (iterable,)),
+                       ('permutations', (iterable,)),
+                       ('product', (iterable,)),
+                       ('repeat', (1, 2)),
+                       ('starmap', pred_iterable),
+                       ('takewhile', pred_iterable),
+                       ('tee', (iterable,)),
+                       # Python 3 additions
+                       ('zip_longest', (iterable,)),
+                       ('accumulate', (iterable,)),
+                       ('compress', (iterable, ())),
+                       ('combinations', (iterable, 1)),
+                       ('combinations_with_replacement', (iterable, 1)),
+    ):
+        try:
+            func = getattr(itertools, func)
+        except AttributeError:
+            if func in missing:
+                continue
+            raise
+        result = func(*args)
+        if func == itertools.tee:
+            result = result[0]
+        tresult = type(result)
+        if tresult not in _default_checkers:
+            _default_checkers[tresult] = _iteratorChecker
+
+_fixup_itertools()
+del _fixup_itertools
 
 def _clear():
     _checkers.clear()
diff --git a/src/zope/security/tests/test_checker.py b/src/zope/security/tests/test_checker.py
index 155671b..692cbde 100644
--- a/src/zope/security/tests/test_checker.py
+++ b/src/zope/security/tests/test_checker.py
@@ -14,6 +14,7 @@
 """Tests for zope.security.checker
 """
 import unittest
+from zope.security import checker
 
 def _skip_if_not_Py2(testfunc):
     import sys
@@ -201,8 +202,7 @@ _marker = []
 class CheckerTestsBase(object):
 
     def _getTargetClass(self):
-        from zope.security.checker import Checker
-        return Checker
+        raise NotImplementedError("Subclasses must define")
 
     def _makeOne(self, get_permissions=_marker, set_permissions=_marker):
         if get_permissions is _marker:
@@ -937,6 +937,52 @@ class _SelectCheckerBase(object):
         _checkers[Foo] = _factory_factory
         self.assertTrue(self._callFUT(Foo()) is checker)
 
+    def test_itertools_checkers(self):
+        from zope.security.checker import _iteratorChecker
+        import sys
+        import itertools
+        pred = lambda x: x
+        iterable = (1, 2, 3)
+        pred_iterable = (pred, iterable)
+        missing_in_py3 = {'ifilter', 'ifilterfalse', 'imap',
+                          'izip', 'izip_longest'}
+        missing_in_py2 = {'zip_longest', 'accumulate', 'compress',
+                          'combinations', 'combinations_with_replacement'}
+        missing = missing_in_py3 if sys.version_info[0] >= 3 else missing_in_py2
+        for func, args in (('count', ()),
+                           ('cycle', ((),)),
+                           ('dropwhile', pred_iterable),
+                           ('ifilter', pred_iterable),
+                           ('ifilterfalse', pred_iterable),
+                           ('imap', pred_iterable),
+                           ('islice', (iterable, 2)),
+                           ('izip', (iterable,)),
+                           ('izip_longest', (iterable,)),
+                           ('permutations', (iterable,)),
+                           ('product', (iterable,)),
+                           ('repeat', (1, 2)),
+                           ('starmap', pred_iterable),
+                           ('takewhile', pred_iterable),
+                           ('tee', (iterable,)),
+                           # Python 3 additions
+                           ('zip_longest', (iterable,)),
+                           ('accumulate', (iterable,)),
+                           ('compress', (iterable, ())),
+                           ('combinations', (iterable, 1)),
+                           ('combinations_with_replacement', (iterable, 1)),
+        ):
+            try:
+                func = getattr(itertools, func)
+            except AttributeError:
+                if func in missing:
+                    continue
+                raise
+            __traceback_info__ = func
+            result = func(*args)
+            if func == itertools.tee:
+                result = result[0]
+
+            self.assertIs(self._callFUT(result), _iteratorChecker)
 
 
 class Test_selectCheckerPy(_SelectCheckerBase, unittest.TestCase):
@@ -947,6 +993,8 @@ class Test_selectCheckerPy(_SelectCheckerBase, unittest.TestCase):
 
 
 
+@unittest.skipIf(checker.selectChecker is checker.selectCheckerPy,
+                 "Pure Python")
 class Test_selectChecker(_SelectCheckerBase, unittest.TestCase):
 
     def _callFUT(self, obj):
@@ -2148,7 +2196,7 @@ class TestMixinDecoratedChecker(unittest.TestCase):
         from zope.security.checker import Checker
         return Checker(self.decorationGetMap, self.decorationSetMap)
 
-class TestCombinedChecker(TestMixinDecoratedChecker, unittest.TestCase):
+class TestCombinedCheckerMixin(TestMixinDecoratedChecker, unittest.TestCase):
 
     def setUp(self):
         unittest.TestCase.setUp(self)