Add tag filtering to Cache.iterkeys

diskcache/core.py

@@ -2205,7 +2205,7 @@ def _select_delete(
 
         return count
 
-    def iterkeys(self, reverse=False):
+    def iterkeys(self, reverse=False, tag=ENOVAL):
         """Iterate Cache keys in database sort order.
 
         >>> cache = Cache()
@@ -2216,36 +2216,39 @@ def iterkeys(self, reverse=False):
         >>> list(cache.iterkeys(reverse=True))
         [4, 3, 2, 1, 0]
 
+        If `tag` is given, only keys with a matching tag are returned. Use
+        ``tag=None`` to iterate untagged keys. Omitting `tag` returns all keys.
+        Iteration includes expired items and does not read cached values.
+
         :param bool reverse: reverse sort order (default False)
+        :param tag: tag to match (default ENOVAL, matching all tags)
         :return: iterator of Cache keys
 
         """
         sql = self._sql
         limit = 100
         _disk_get = self._disk.get
+        select = 'SELECT key, raw FROM Cache'
+        args = ()
+
+        if tag is None:
+            select += ' WHERE tag IS NULL'
+        elif tag is not ENOVAL:
+            select += ' WHERE tag = ?'
+            args = (tag,)
+
+        iterate = select + (' WHERE (' if tag is ENOVAL else ' AND (')
 
         if reverse:
-            select = (
-                'SELECT key, raw FROM Cache'
-                ' ORDER BY key DESC, raw DESC LIMIT 1'
-            )
-            iterate = (
-                'SELECT key, raw FROM Cache'
-                ' WHERE key = ? AND raw < ? OR key < ?'
-                ' ORDER BY key DESC, raw DESC LIMIT ?'
-            )
+            iterate += 'key = ? AND raw < ? OR key < ?)'
+            order = ' ORDER BY key DESC, raw DESC'
         else:
-            select = (
-                'SELECT key, raw FROM Cache'
-                ' ORDER BY key ASC, raw ASC LIMIT 1'
-            )
-            iterate = (
-                'SELECT key, raw FROM Cache'
-                ' WHERE key = ? AND raw > ? OR key > ?'
-                ' ORDER BY key ASC, raw ASC LIMIT ?'
-            )
+            iterate += 'key = ? AND raw > ? OR key > ?)'
+            order = ' ORDER BY key ASC, raw ASC'
 
-        row = sql(select).fetchall()
+        select += order + ' LIMIT 1'
+        iterate += order + ' LIMIT ?'
+        row = sql(select, args).fetchall()
 
         if row:
             ((key, raw),) = row
@@ -2255,7 +2258,7 @@ def iterkeys(self, reverse=False):
         yield _disk_get(key, raw)
 
         while True:
-            rows = sql(iterate, (key, raw, key, limit)).fetchall()
+            rows = sql(iterate, args + (key, raw, key, limit)).fetchall()
 
             if not rows:
                 break

docs/tutorial.rst

@@ -311,6 +311,20 @@ keys will be serialized which is likely to have a meaningless sorted order.
     >>> cache.peekitem(last=False)
     ('c', None)
 
+To iterate keys with a matching tag, pass `tag` to :meth:`iterkeys
+<.Cache.iterkeys>`. Filtering happens in the database without reading cached
+values. Omitting `tag` includes all keys; ``tag=None`` selects untagged keys.
+As with other iteration methods, expired items are included.
+
+    >>> _ = cache.set('a', None, tag='letter')
+    >>> _ = cache.set('c', None, tag='letter')
+    >>> list(cache.iterkeys(tag='letter'))
+    ['a', 'c']
+    >>> list(cache.iterkeys(reverse=True, tag='letter'))
+    ['c', 'a']
+    >>> list(cache.iterkeys(tag=None))
+    ['b']
+
 If only the first or last item in insertion order is desired then
 :meth:`peekitem <.Cache.peekitem>` is more efficient than using iteration.
 

tests/test_core.py

@@ -1034,6 +1034,84 @@ def test_iterkeys(cache):
     assert list(cache.iterkeys()) == []
 
 
+@pytest.mark.parametrize('reverse', [False, True])
+@pytest.mark.parametrize('tag', [None, '', b'', 0, False, 1.5, 'counter'])
+def test_iterkeys_tag(cache, reverse, tag):
+    assert list(cache.iterkeys(reverse=reverse, tag=tag)) == []
+    cache.set('a', 1, tag=tag)
+    cache.set('b', 2, tag='other')
+    cache.set('c', 3, tag=tag)
+    cache.set('d', 4)
+
+    expected = ['a', 'c', 'd'] if tag is None else ['a', 'c']
+    assert list(cache.iterkeys(reverse=reverse, tag=tag)) == (
+        expected[::-1] if reverse else expected
+    )
+    assert list(cache.iterkeys(reverse=reverse, tag='missing')) == []
+    assert list(cache.iterkeys(reverse)) == (
+        ['d', 'c', 'b', 'a'] if reverse else ['a', 'b', 'c', 'd']
+    )
+
+
+@pytest.mark.parametrize('reverse', [False, True])
+@pytest.mark.parametrize('tag_index', [False, True])
+def test_iterkeys_tag_pagination(cache, reverse, tag_index):
+    if tag_index:
+        cache.create_tag_index()
+
+    expected = []
+    for convert in (
+        int,
+        lambda num: 'key-%03d' % num,
+        lambda num: bytes([num]),
+    ):
+        for num in range(240):
+            key = convert(num)
+            tag = 'counter' if num % 2 else 'other'
+            cache.set(key, None, tag=tag)
+            if num % 2:
+                expected.append(key)
+
+    assert list(cache.iterkeys(reverse=reverse, tag='counter')) == (
+        expected[::-1] if reverse else expected
+    )
+
+
+@pytest.mark.parametrize('reverse', [False, True])
+def test_iterkeys_tag_raw_collision(cache, reverse):
+    pairs = []
+    for num in range(105):
+        key = (num,)
+        raw_key = bytes(cache.disk.put(key)[0])
+        cache.set(key, None, tag='counter')
+        cache.set(raw_key, None, tag='counter' if num % 2 else 'other')
+        pairs.append((raw_key, key, num % 2))
+
+    expected = []
+    for raw_key, key, matching in sorted(pairs):
+        expected.append(key)
+        if matching:
+            expected.append(raw_key)
+
+    assert list(cache.iterkeys(reverse=reverse, tag='counter')) == (
+        expected[::-1] if reverse else expected
+    )
+
+
+def test_iterkeys_tag_does_not_fetch_values(cache):
+    cache.set('key', {'value': 1}, tag='counter')
+    with mock.patch.object(cache.disk, 'fetch') as fetch:
+        assert list(cache.iterkeys(tag='counter')) == ['key']
+    fetch.assert_not_called()
+
+
+def test_iterkeys_tag_includes_expired(cache):
+    cache.reset('cull_limit', 0)
+    cache.set('expired', 1, tag='counter', expire=-1)
+    cache.set('live', 2, tag='counter')
+    assert list(cache.iterkeys(tag='counter')) == ['expired', 'live']
+
+
 def test_pickle(cache):
     for num, val in enumerate('abcde'):
         cache[val] = num