scan_items/scan_values honour a reader subclass overriding __iter__/__getitem__ (#9)

dynamodol/base.py

@@ -654,12 +654,32 @@ def load_sample_data():
 # TODO class DynamoDbStore(DynamoDbBasePersister, Store): ...
 
 
+_SCAN_FAST_PATH_METHODS = ("__iter__", "__getitem__")
+
+
+def _overrides_reader_protocol(reader) -> bool:
+    """Whether ``type(reader)`` overrides ``__iter__``/``__getitem__`` outside dynamodol.
+
+    The single-scan fast path reads the table itself, so it would skip an override that
+    narrows the keys or maps the values (dynamodol's own subclasses, e.g. the prefix and
+    partition readers, are what the scan understands). Only those two methods are
+    checked: a subclass that narrows ``_scan_records`` or ``items()`` is not detected.
+    """
+    package = __name__.partition(".")[0]
+    for name in _SCAN_FAST_PATH_METHODS:
+        owner = next((c for c in type(reader).__mro__ if name in vars(c)), None)
+        if owner is not None and (owner.__module__ or "").partition(".")[0] != package:
+            return True
+    return False
+
+
 def _leaf_reader(store, *, fn):
     """``store`` if it is an unwrapped dynamodol reader (the single-scan fast path),
-    ``None`` if it is any other mapping (a dol-wrapped store: use its own ``.items()`` /
-    ``.values()``, which apply the wrapper's scope and codecs); ``TypeError`` otherwise."""
+    ``None`` if it is any other mapping (a dol-wrapped store, or a reader subclass that
+    overrides ``__iter__``/``__getitem__``: use its own ``.items()`` / ``.values()``, which
+    apply its scope and codecs); ``TypeError`` otherwise."""
     if isinstance(store, DynamoDbBaseReader):
-        return store
+        return None if _overrides_reader_protocol(store) else store
     if isinstance(store, Mapping):
         return None
     raise TypeError(
@@ -674,7 +694,8 @@ def scan_items(store):
     A bare dynamodol reader is read in one scan: the opt-in fast path that
     ``store.items()`` deliberately is not, since ``items()`` batches ``BatchGetItem``
     (100 keys per call) so that it is correct under any dol wrapper. A *wrapped* store
-    (key codec, ``filt_iter``, value codec...) has its own mapping, scope and decoding,
+    (key codec, ``filt_iter``, value codec...), or a reader subclass that overrides
+    ``__iter__``/``__getitem__``, has its own mapping, scope and decoding,
     which a free function cannot reproduce from the table: for it this is simply
     ``iter(store.items())``. Anything that is not a mapping raises ``TypeError``.
 

tests/test_wrapped_views.py

@@ -123,3 +123,45 @@ def test_base_persister_scan():
         data_fields=DATA_FIELDS,
     )
     assert list(dynamodol.scan_items(s)) == [("k1", "A")]
+
+
+# --- a reader subclass that overrides __iter__/__getitem__ (dynamodol#9) ----------
+
+
+class OnlyAlice(DynamoDbPrefixReader):
+    def __iter__(self):
+        return (k for k in super().__iter__() if k.startswith("alice/"))
+
+
+class Shouting(DynamoDbPrefixReader):
+    def __getitem__(self, k):
+        return super().__getitem__(k).lower()
+
+
+def test_scan_honours_a_subclass_that_overrides_iter():
+    s = _leaf(OnlyAlice)
+    assert list(s.items()) == [("alice/notes.txt", "ALICE-SECRET")]
+    assert list(dynamodol.scan_items(s)) == list(s.items())
+    assert list(dynamodol.scan_values(s)) == ["ALICE-SECRET"]  # no BOB-SECRET
+
+
+def test_scan_honours_a_subclass_that_overrides_getitem():
+    s = _leaf(Shouting)
+    assert list(dynamodol.scan_items(s)) == list(s.items())
+    assert list(dynamodol.scan_values(s)) == ["alice-secret", "bob-secret"]
+
+
+def test_scan_of_an_override_inherited_through_a_subclass_delegates():
+    class Deeper(OnlyAlice):
+        pass
+
+    assert list(dynamodol.scan_values(_leaf(Deeper))) == ["ALICE-SECRET"]
+
+
+def test_dynamodols_own_subclasses_keep_the_single_scan(monkeypatch):
+    s = _leaf()  # DynamoDbPrefixReader: overrides __iter__, but inside dynamodol
+    calls = []
+    monkeypatch.setattr(
+        DynamoDbPrefixReader, "_scan_records", lambda self, projection: calls.append(1) or []
+    )
+    assert list(dynamodol.scan_items(s)) == [] and len(calls) == 1