feat: on_overwrite(policy, same) (#161)

dol/__init__.py

@@ -142,6 +142,7 @@ def ihead(store, n=1):
     FirstArgIsMapping,  # mark a transform whose first arg is the store (self), not data
     confirm_overwrite,  # wrap_kvs preset: confirm before overwriting a differing value
     mk_confirm_overwrite_preset,  # factory for a customizable confirm_overwrite preset
+    on_overwrite,  # overwrite policy for a store: warn/raise/overwrite + same-value check
     filt_iter,  # filter store keys (and contains ready to use filters as attributes)
     cached_keys,  # cache store keys
     add_decoder,  # add a decoder (i.e. outcomming value transformer) to a store

dol/tests/test_on_overwrite.py

@@ -0,0 +1,161 @@
+"""Tests for ``dol.on_overwrite`` (issue #161)."""
+
+import sys
+import warnings
+
+import pytest
+
+from dol import on_overwrite
+from dol.errors import OverWritesNotAllowedError
+
+
+def test_raise_and_same_value_silent():
+    d = on_overwrite(dict(a=1), policy="raise")
+    d["a"] = 1
+    d["b"] = 2
+    with pytest.raises(OverWritesNotAllowedError):
+        d["a"] = 3
+    assert d["a"] == 1
+
+
+def test_warn_writes_and_warns_once():
+    d = on_overwrite(dict(a=1))
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        d["a"] = 1
+        assert not w
+        d["a"] = 2
+    assert len(w) == 1 and d["a"] == 2
+    assert w[0].filename == __file__  # points at the caller
+
+
+def test_overwrite_is_silent():
+    d = on_overwrite(dict(a=1), policy="overwrite")
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        d["a"] = 2
+    assert d["a"] == 2
+
+
+def test_custom_same_and_class_decorator():
+    @on_overwrite(policy="raise", same=lambda a, b: a.__name__ == b.__name__)
+    class Reg(dict):
+        pass
+
+    r = Reg()
+    r["f"] = len
+    r["f"] = len
+    with pytest.raises(OverWritesNotAllowedError):
+        r["f"] = print
+
+
+def test_bad_policy():
+    with pytest.raises(ValueError):
+        on_overwrite(dict(), policy="nope")
+
+
+def test_overwrite_policy_returns_store_unwrapped():
+    d = dict(a=1)
+    assert on_overwrite(d, policy="overwrite") is d
+
+
+class _AmbiguousBool:
+    def __bool__(self):
+        raise ValueError("truth value is ambiguous")
+
+
+class _ArrayLike:
+    """numpy-like: ``==`` returns something with an ambiguous truth value."""
+
+    def __eq__(self, other):
+        return _AmbiguousBool()
+
+    __hash__ = object.__hash__
+
+
+class _EqRaises:
+    def __eq__(self, other):
+        raise TypeError("no comparison")
+
+    __hash__ = object.__hash__
+
+
+@pytest.mark.parametrize("cls", [_ArrayLike, _EqRaises])
+def test_default_same_does_not_crash(cls):
+    x = cls()
+    d = on_overwrite(dict(a=x), policy="raise")
+    d["a"] = x  # identical object: same
+    with pytest.raises(OverWritesNotAllowedError):
+        d["a"] = cls()  # comparison fails -> treated as different
+
+
+def test_warning_points_at_setitem_line():
+    d = on_overwrite(dict(a=1))
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        lineno = sys._getframe().f_lineno + 1
+        d["a"] = 2
+    assert (w[0].filename, w[0].lineno) == (__file__, lineno)
+
+
+def test_factory_form():
+    d = on_overwrite(policy="raise")(dict(a=1))
+    with pytest.raises(OverWritesNotAllowedError):
+        d["a"] = 2
+
+
+def test_inner_store_mutated():
+    inner = dict(a=1)
+    d = on_overwrite(inner, policy="raise")
+    inner["a"] = 5
+    d["a"] = 5  # matches the current inner value: silent
+    with pytest.raises(OverWritesNotAllowedError):
+        d["a"] = 1
+
+
+def test_same_argument_order():
+    calls = []
+
+    def same(existing, new):
+        calls.append((existing, new))
+        return False
+
+    d = on_overwrite(dict(a="old"), same=same)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        d["a"] = "new"
+    assert calls == [("old", "new")]
+
+
+def test_callable_policy():
+    def no_shrink(k, existing, new):
+        if new < existing:
+            raise ValueError(f"{k}: {new} < {existing}")
+        return new
+
+    d = on_overwrite(dict(a=1), policy=no_shrink)
+    d["a"] = 3
+    with pytest.raises(ValueError, match="a: 2 < 3"):
+        d["a"] = 2
+    assert d["a"] == 3
+
+    keep = on_overwrite(dict(a=1), policy=lambda k, existing, new: existing)
+    keep["a"] = 9
+    keep["b"] = 2  # policy only consulted for a differing existing value
+    assert dict(keep) == {"a": 1, "b": 2}
+
+
+def test_update_goes_through_policy():
+    d = on_overwrite(dict(a=1), policy="raise")
+    d.update(b=2)
+    with pytest.raises(OverWritesNotAllowedError):
+        d.update(a=3)
+    assert dict(d) == {"a": 1, "b": 2}
+
+
+def test_delete_then_write_bypasses_raise():
+    """Documented: unlike ``disallow_overwrites``, deleting first lets a write through."""
+    d = on_overwrite(dict(a=1), policy="raise")
+    del d["a"]
+    d["a"] = 3
+    assert d["a"] == 3

dol/trans.py

@@ -2547,6 +2547,34 @@ def __call__(self, *args, **kwargs):
             raise TypeError(f"{self._obj=} is not callable")
 
 
+_ABSENT = object()  # sentinel: "no differing existing value"
+
+
+def _not_ne(a, b):
+    """``not (a != b)``: the historical sameness test of ``confirm_overwrite``."""
+    return not (a != b)
+
+
+def _same_value(a, b) -> bool:
+    """Default ``same`` for :func:`on_overwrite`: identity, else ``bool(a == b)``; any
+    error (e.g. an ambiguous array comparison) counts as "not the same"."""
+    if a is b:
+        return True
+    try:
+        return bool(a == b)
+    except Exception:
+        return False
+
+
+def _existing_if_different(self, k, v, same):
+    """Return the value ``self`` holds at ``k`` if there is one that is not
+    ``same(existing, v)``, else ``_ABSENT``."""
+    existing = self.get(k, _ABSENT)
+    if existing is not _ABSENT and not same(existing, v):
+        return existing
+    return _ABSENT
+
+
 def _dflt_overwrite_prompt(k, existing, new):
     return (
         f"The key {k} already exists and has value {existing}. "
@@ -2589,9 +2617,8 @@ def mk_confirm_overwrite_preset(
     def confirm_overwrite(self, k, v):
         """``preset`` that returns ``v`` unless ``k`` already holds a different value
         and the user does not confirm; then the existing value is kept."""
-        _sentinel = object()
-        existing = self.get(k, _sentinel)
-        if existing is not _sentinel and existing != v:
+        existing = _existing_if_different(self, k, v, _not_ne)
+        if existing is not _ABSENT:
             # resolve ``input`` at call time (not capture) so monkeypatching works
             ask = get_input if get_input is not None else input
             if ask(prompt(k, existing, v)) != str(v):
@@ -2614,6 +2641,121 @@ def confirm_overwrite(self, k, v):
 confirm_overwrite = mk_confirm_overwrite_preset()
 
 
+def _overwrite_msg(k):
+    return f"key {k!r} already exists with a different value"
+
+
+def _warn_on_overwrite(k, existing, new, *, stacklevel=3):
+    warn(_overwrite_msg(k) + "; overwriting it", stacklevel=stacklevel)
+    return new
+
+
+def _raise_on_overwrite(k, existing, new, *, stacklevel=None):
+    raise OverWritesNotAllowedError(_overwrite_msg(k))
+
+
+def _just_overwrite(k, existing, new, *, stacklevel=None):
+    return new
+
+
+#: Named overwrite strategies: ``(k, existing, new, *, stacklevel) -> value to write``.
+_OVERWRITE_POLICIES = {
+    "warn": _warn_on_overwrite,
+    "raise": _raise_on_overwrite,
+    "overwrite": _just_overwrite,
+}
+
+
+def _resolve_overwrite_policy(policy, *, stacklevel):
+    """Turn a policy name or a ``(k, existing, new) -> value`` callable into a strategy.
+
+    >>> _resolve_overwrite_policy("nope", stacklevel=3)
+    Traceback (most recent call last):
+      ...
+    ValueError: policy must be one of ['warn', 'raise', 'overwrite'] or a callable, got 'nope'
+    """
+    if callable(policy):
+        return policy
+    if policy not in _OVERWRITE_POLICIES:
+        raise ValueError(
+            f"policy must be one of {list(_OVERWRITE_POLICIES)} or a callable, "
+            f"got {policy!r}"
+        )
+    # stacklevel + 1: the preset adds one frame between the strategy and __setitem__
+    return partial(_OVERWRITE_POLICIES[policy], stacklevel=stacklevel + 1)
+
+
+def _mk_overwrite_policy_preset(*, policy="warn", same=_same_value, stacklevel=3):
+    """Make a ``wrap_kvs`` ``preset`` applying an overwrite ``policy``.
+    See :func:`on_overwrite`."""
+    strategy = _resolve_overwrite_policy(policy, stacklevel=stacklevel)
+
+    def overwrite_policy(self, k, v):
+        existing = _existing_if_different(self, k, v, same)
+        if existing is _ABSENT:
+            return v
+        return strategy(k, existing, v)
+
+    return overwrite_policy
+
+
+def on_overwrite(
+    store=None,
+    *,
+    policy: str | Callable[[Any, Any, Any], Any] = "warn",
+    same: Callable[[Any, Any], bool] = _same_value,
+    stacklevel: int = 3,
+):
+    """Apply an overwrite ``policy`` to a store (class or instance; with no ``store``,
+    return a decorator), with a pluggable ``same(existing, new)`` check so that
+    re-writing an equivalent value (e.g. a module reload) is silent.
+
+    ``policy`` is ``"warn"`` (warn, then write), ``"raise"``
+    (``OverWritesNotAllowedError``), ``"overwrite"`` (the store is returned unwrapped),
+    or a callable ``(k, existing, new) -> value_to_write``, called only when ``k``
+    already holds a value that is not ``same`` as ``new`` (it may also raise).
+    ``same`` defaults to ``is``-or-``==``, where a failing or ambiguous ``==`` (arrays)
+    counts as different. ``stacklevel`` is the ``"warn"`` warning's stack level, as seen
+    from the ``store[k] = v`` line; raise it by one per extra layer you call it through.
+
+    Notes: each checked write first reads the old value (a round-trip on remote stores);
+    the read-compare-write is not atomic across threads; ``del s[k]`` then ``s[k] = v``
+    bypasses ``"raise"`` (unlike ``disallow_overwrites``); and the comparison is made on
+    the objects at this layer (what ``store[k]`` returns and ``store[k] = v`` receives,
+    not the encoded data of inner layers).
+
+    >>> d = on_overwrite(dict(a=1), policy="raise")
+    >>> d['a'] = 1  # same value: silent
+    >>> d['b'] = 2  # new key: fine
+    >>> d['a'] = 3
+    Traceback (most recent call last):
+      ...
+    dol.errors.OverWritesNotAllowedError: key 'a' already exists with a different value
+
+    With ``"warn"`` (the default) the write goes through, with a warning:
+
+    >>> import warnings
+    >>> w = on_overwrite(dict(a=1))
+    >>> with warnings.catch_warnings(record=True) as caught:
+    ...     warnings.simplefilter("always")
+    ...     w['a'] = 2
+    >>> len(caught), w['a']
+    (1, 2)
+
+    ``same`` decides what counts as "the same value":
+
+    >>> s = on_overwrite(dict(f=len), policy="raise", same=lambda a, b: a.__name__ == b.__name__)
+    >>> s['f'] = len  # same by name: silent
+    """
+    strategy = _resolve_overwrite_policy(policy, stacklevel=stacklevel)  # validate early
+    if store is None:  # factory / parametrized class decorator form
+        return partial(on_overwrite, policy=policy, same=same, stacklevel=stacklevel)
+    if isinstance(policy, str) and policy == "overwrite":
+        return store  # nothing to check: skip the per-write read
+    preset = _mk_overwrite_policy_preset(policy=strategy, same=same)
+    return wrap_kvs(store, preset=preset)
+
+
 def add_aliases(obj, **aliases):
     """A function that wraps the object instance and adds aliases.