Merge pull request #287 from roryk/master

Add assoc_in function.

AUTHORS.md

@@ -25,3 +25,5 @@ Nikolaos-Digenis Karagiannis                    [@digenis](https://github.com/di
 [Antonio Lima](https://twitter.com/themiurgo)   [@themiurgo](https://github.com/themiurgo/)
 
 Joe Jevnik                                      [@llllllllll](https://github.com/llllllllll)
+
+Rory Kirchner                                      [@roryk](https://github.com/roryk)

doc/source/api.rst

@@ -76,6 +76,7 @@ Dicttoolz
 .. autosummary::
    assoc
    dissoc
+   assoc_in
    get_in
    keyfilter
    keymap

toolz/dicttoolz.py

@@ -5,7 +5,7 @@
 
 __all__ = ('merge', 'merge_with', 'valmap', 'keymap', 'itemmap',
            'valfilter', 'keyfilter', 'itemfilter',
-           'assoc', 'dissoc', 'update_in', 'get_in')
+           'assoc', 'dissoc', 'assoc_in', 'update_in', 'get_in')
 
 
 def _get_factory(f, kwargs):
@@ -218,6 +218,22 @@ def dissoc(d, *keys):
     return d2
 
 
+def assoc_in(d, keys, value, factory=dict):
+    """
+    Return a new dict with new, potentially nested, key value pair
+
+    >>> purchase = {'name': 'Alice',
+    ...             'order': {'items': ['Apple', 'Orange'],
+    ...                       'costs': [0.50, 1.25]},
+    ...             'credit card': '5555-1234-1234-1234'}
+    >>> assoc_in(purchase, ['order', 'costs'], [0.25, 1.00]) # doctest: +SKIP
+    {'credit card': '5555-1234-1234-1234',
+     'name': 'Alice',
+     'purchase': {'costs': [0.25, 1.00], 'items': ['Apple', 'Orange']}}
+    """
+    return update_in(d, keys, lambda x: value, value, factory)
+
+
 def update_in(d, keys, func, default=None, factory=dict):
     """ Update value in a (potentially) nested dictionary
 

toolz/tests/test_dicttoolz.py

@@ -1,7 +1,7 @@
 from collections import defaultdict as _defaultdict
 from toolz.dicttoolz import (merge, merge_with, valmap, keymap, update_in,
                              assoc, dissoc, keyfilter, valfilter, itemmap,
-                             itemfilter)
+                             itemfilter, assoc_in)
 from toolz.utils import raises
 from toolz.compatibility import PY3
 
@@ -103,6 +103,20 @@ def test_dissoc(self):
         assert d is oldd
         assert d2 is not oldd
 
+    def test_assoc_in(self):
+        D, kw = self.D, self.kw
+        assert assoc_in(D({"a": 1}), ["a"], 2, **kw) == D({"a": 2})
+        assert (assoc_in(D({"a": D({"b": 1})}), ["a", "b"], 2, **kw) ==
+                D({"a": D({"b": 2})}))
+        assert assoc_in(D({}), ["a", "b"], 1, **kw) == D({"a": D({"b": 1})})
+
+        # Verify immutability:
+        d = D({'x': 1})
+        oldd = d
+        d2 = assoc_in(d, ['x'], 2, **kw)
+        assert d is oldd
+        assert d2 is not oldd
+
     def test_update_in(self):
         D, kw = self.D, self.kw
         assert update_in(D({"a": 0}), ["a"], inc, **kw) == D({"a": 1})
@@ -236,4 +250,3 @@ class TestCustomMapping(TestDict):
     """
     D = CustomMapping
     kw = {'factory': lambda: CustomMapping()}
-