Repository: tobgu/pyrsistent

Issue:
Allow predicates on values in transform paths for pvectors and pmaps
I often want to write transformations like:

```
# Double even elements of the list
xs = pvector([1, 3, 2, 4])
assert xs.transform(lambda x: x % 2 == 0, lambda x: x * 2) == pvector([1, 3, 4, 8])
```

Particular with more complex paths:

```
xs = pmap({"foo": pvector([1, 3, 2, 4]), "bar": pvector([1, 2, 3, 4])})
assert xs.transform(
    ["foo", lambda x: x % 2 == 0],
    lambda x: x * 2,
) == pmap({"foo": pvector([1, 3, 4, 8]), "bar": pvector([1, 2, 3, 4])})
```

I can often cobble together something that deals with the above two scenarios.  It usually involves duplicating some information from the path but it's kind of manageable.

Where I really have trouble is when I want to use this missing feature twice or more within a single path:

```
xs = pmap({"foo": pvector([1, 3, 2, 4]), "bar": pvector([1, 2])})
assert xs.transform(
    [lambda (k, v): len(v) == 4, lambda x: x % 2 == 0],
    lambda x: x * 2,
) == pmap({"foo": pvector([1, 3, 4, 8]), "bar": pvector([1, 2])})
```

Discussion:

Patch:
diff --git a/README.rst b/README.rst
index ccbd38a..2fb4135 100644
--- a/README.rst
+++ b/README.rst
@@ -444,6 +444,8 @@ treated as matchers. If the matcher returns True for a specific key it is consid
     pvector([1, 2, 4, 4, 5])
     >>> v1.transform([lambda ix: 0 < ix < 4], 8)
     pvector([1, 8, 8, 8, 5])
+    >>> v1.transform([lambda ix, v: ix == 0 or v == 5], 0)
+    pvector([0, 2, 3, 4, 0])
 
     # The (a)ny matcher can be used to match anything
     >>> v1.transform([ny], 8)
diff --git a/pyrsistent/_transformations.py b/pyrsistent/_transformations.py
index c4c7e10..6fe31cf 100644
--- a/pyrsistent/_transformations.py
+++ b/pyrsistent/_transformations.py
@@ -1,5 +1,13 @@
 import re
 import six
+try:
+    from inspect import Parameter, signature
+except ImportError:
+    signature = None
+    try:
+        from inspect import getfullargspec as getargspec
+    except ImportError:
+        from inspect import getargspec
 
 
 def inc(x):
@@ -74,11 +82,40 @@ def _get(structure, key, default):
 def _get_keys_and_values(structure, key_spec):
     from pyrsistent._pmap import pmap
     if callable(key_spec):
-        return [(k, v) for k, v in _items(structure) if key_spec(k)]
-
+        # Support predicates as callable objects in the path
+        arity = _get_arity(key_spec)
+        if arity == 1:
+            # Unary predicates are called with the "key" of the path
+            # - eg a key in a mapping, an index in a sequence.
+            return [(k, v) for k, v in _items(structure) if key_spec(k)]
+        elif arity == 2:
+            # Binary predicates are called with the key and the corresponding
+            # value.
+            return [(k, v) for k, v in _items(structure) if key_spec(k, v)]
+        else:
+            # Other arities are an error.
+            raise ValueError(
+                "callable in transform path must take 1 or 2 arguments"
+            )
+
+    # Non-callables are used as-is as a key.
     return [(key_spec, _get(structure, key_spec, pmap()))]
 
 
+if signature is None:
+    def _get_arity(f):
+        argspec = getargspec(f)
+        return len(argspec.args) - len(argspec.defaults or ())
+else:
+    def _get_arity(f):
+        return sum(
+            1
+            for p
+            in signature(f).parameters.values()
+            if p.default is Parameter.empty
+            and p.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD)
+        )
+
 def _update_structure(structure, kvs, path, command):
     e = structure.evolver()
     if not path and command is discard:


Test patch:
diff --git a/tests/transform_test.py b/tests/transform_test.py
index 490ce9b..32e1772 100644
--- a/tests/transform_test.py
+++ b/tests/transform_test.py
@@ -3,7 +3,7 @@ from pyrsistent import freeze, inc, discard, rex, ny, field, PClass
 
 def test_callable_command():
     m = freeze({'foo': {'bar': {'baz': 1}}})
-    m.transform(['foo', 'bar', 'baz'], inc) == {'foo': {'bar': {'baz': 2}}}
+    assert m.transform(['foo', 'bar', 'baz'], inc) == {'foo': {'bar': {'baz': 2}}}
 
 
 def test_predicate():
@@ -11,6 +11,32 @@ def test_predicate():
     assert m.transform(['foo', lambda x: x.startswith('b'), 'baz'], inc) == {'foo': {'bar': {'baz': 2}, 'qux': {'baz': 1}}}
 
 
+def test_broken_predicate():
+    broken_predicates = [
+        lambda: None,
+        lambda a, b, c: None,
+        lambda a, b, c, d=None: None,
+        lambda *args: None,
+        lambda **kwargs: None,
+    ]
+    for pred in broken_predicates:
+        try:
+            freeze({}).transform([pred], None)
+            assert False
+        except ValueError as e:
+            assert str(e) == "callable in transform path must take 1 or 2 arguments"
+
+
+def test_key_value_predicate():
+    m = freeze({
+        'foo': 1,
+        'bar': 2,
+    })
+    assert m.transform([
+        lambda k, v: (k, v) == ('foo', 1),
+    ], lambda v: v * 3) == {"foo": 3, "bar": 2}
+
+
 def test_remove():
     m = freeze({'foo': {'bar': {'baz': 1}}})
     assert m.transform(['foo', 'bar', 'baz'], discard) == {'foo': {'bar': {}}}