feat: scan_items/scan_values accept wrapped stores (#6)

dynamodol/base.py

@@ -649,33 +649,40 @@ def load_sample_data():
 
 
 def _leaf_reader(store, *, fn):
-    if not isinstance(store, DynamoDbBaseReader):
-        raise TypeError(
-            f"`{fn}` needs an unwrapped dynamodol reader, got "
-            f"{type(store).__name__}. A wrapped store (key codec, filt_iter, value "
-            f"codec...) must be iterated with its own `.items()` / `.values()`, which "
-            f"apply the wrapper's scope and codecs (one `get` per key)."
-        )
-    return store
+    """``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."""
+    if isinstance(store, DynamoDbBaseReader):
+        return store
+    if isinstance(store, Mapping):
+        return None
+    raise TypeError(
+        f"`{fn}` needs a dynamodol reader (bare, or wrapped with a key codec, "
+        f"filt_iter, value codec...), got {type(store).__name__}."
+    )
 
 
 def scan_items(store):
-    """Iterate ``(key, value)`` pairs of an *unwrapped* dynamodol reader in one scan.
+    """Iterate ``(key, value)`` pairs of a dynamodol store.
 
-    The opt-in fast path that ``store.items()`` deliberately is not: ``items()`` goes
-    through ``__iter__`` and ``__getitem__`` (one get per key) so that it is correct
-    under any dol wrapper. Passing a wrapped store raises ``TypeError``, because a
-    free function cannot apply the wrapper's key mapping, scope or value decoding.
+    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,
+    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``.
 
     >>> import dynamodol
     >>> dynamodol.scan_items(object())
     Traceback (most recent call last):
       ...
-    TypeError: `scan_items` needs an unwrapped dynamodol reader, got object. A wrapped store (key codec, filt_iter, value codec...) must be iterated with its own `.items()` / `.values()`, which apply the wrapper's scope and codecs (one `get` per key).
+    TypeError: `scan_items` needs a dynamodol reader (bare, or wrapped with a key codec, filt_iter, value codec...), got object.
     """
     reader = _leaf_reader(
         store, fn="scan_items"
     )  # eager: fail at the call, not at first next()
+    if reader is None:
+        return iter(store.items())
     records = reader._scan_records(reader._keys_values_expression)
     return (
         (reader.format_get_key(record), reader.format_get_item(record))
@@ -684,7 +691,10 @@ def scan_items(store):
 
 
 def scan_values(store):
-    """Iterate the values of an *unwrapped* dynamodol reader in one scan (see `scan_items`)."""
+    """Iterate the values of a dynamodol store (see `scan_items`: a wrapped store is
+    iterated with its own ``.values()``)."""
     reader = _leaf_reader(store, fn="scan_values")
+    if reader is None:
+        return iter(store.values())
     records = reader._scan_records(reader._values_expression)
     return (reader.format_get_item(record) for record in records)

tests/test_wrapped_views.py

@@ -60,11 +60,39 @@ def test_scan_fast_path_on_unwrapped_reader():
     assert list(dynamodol.scan_values(s)) == list(s.values())
 
 
-def test_scan_refuses_wrapped_store():
-    with pytest.raises(TypeError, match="unwrapped"):
-        list(dynamodol.scan_items(alice_only(_leaf())))
-    with pytest.raises(TypeError, match="unwrapped"):
-        list(dynamodol.scan_values(alice_only(_leaf())))
+@pytest.mark.parametrize(
+    "mk", [lambda: alice_only(_leaf()), lambda: _leaf(alice_only(DynamoDbPrefixReader))]
+)
+def test_scan_on_wrapped_store_honours_the_wrapper(mk):
+    s = mk()
+    assert list(dynamodol.scan_items(s)) == [("notes.txt", "ALICE-SECRET")]
+    assert list(dynamodol.scan_values(s)) == ["ALICE-SECRET"]  # no BOB-SECRET
+
+
+def test_scan_on_value_codec_wrapped_store_decodes():
+    s = wrap_kvs(_leaf(), obj_of_data=lambda v: {"text": v})
+    assert dict(dynamodol.scan_items(s)) == dict(s)
+    assert list(dynamodol.scan_values(s)) == [{"text": "ALICE-SECRET"}, {"text": "BOB-SECRET"}]
+
+
+def test_scan_on_wrapped_store_is_eager_about_errors_and_lazy_about_data():
+    s = alice_only(_leaf())
+    assert iter(dynamodol.scan_items(s)) is not None
+    for fn in (dynamodol.scan_items, dynamodol.scan_values):
+        with pytest.raises(TypeError, match="needs a dynamodol reader"):
+            fn(object())  # at the call, not at the first next()
+        with pytest.raises(TypeError, match="got NoneType"):
+            fn(None)
+
+
+def test_scan_of_a_bare_reader_still_uses_a_single_scan(monkeypatch):
+    s = _leaf()
+    calls = []
+    monkeypatch.setattr(
+        type(s), "_scan_records", lambda self, projection: calls.append(1) or []
+    )
+    assert list(dynamodol.scan_items(s)) == [] and list(dynamodol.scan_values(s)) == []
+    assert len(calls) == 2
 
 
 def test_removed_methods_point_to_replacement():