From 2bbf33cf3b548ae36466a28076d7cfdfd0ae841c Mon Sep 17 00:00:00 2001 From: Gonghan-Princess <268298391+Gonghan-Princess@users.noreply.github.com> Date: Sun, 13 Sep 2026 20:52:38 +0800 Subject: [PATCH] feat: filter iterkeys by tag (fixes #341) --- diskcache/core.py | 45 +++++++++++++------------- docs/tutorial.rst | 14 +++++++++ tests/test_core.py | 78 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 116 insertions(+), 21 deletions(-) diff --git a/diskcache/core.py b/diskcache/core.py index 7a3d23b..6d62b07 100644 --- a/diskcache/core.py +++ b/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 diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 2eb454d..c32f72e 100644 --- a/docs/tutorial.rst +++ b/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. diff --git a/tests/test_core.py b/tests/test_core.py index 788afef..9d05a6b 100644 --- a/tests/test_core.py +++ b/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