Make assoc a generic function

tests/test_benchmarks.py

@@ -1,8 +1,6 @@
 import pytest
 
-from toolz import assoc
-
-from unification import unify, reify, var, isvar
+from unification import unify, reify, var, isvar, assoc
 from unification.utils import transitive_get as walk
 
 from tests.utils import gen_long_chain

tests/test_core.py

@@ -6,12 +6,33 @@
 from collections import OrderedDict
 
 from unification import var
-from unification.core import isground, reify, unground_lvars, unify
+from unification.core import isground, reify, unground_lvars, unify, assoc
 from unification.utils import freeze
 
 from tests.utils import gen_long_chain
 
 
+def test_assoc():
+    d = {"a": 1, 2: 2}
+    assert assoc(d, "c", 3) is not d
+    assert assoc(d, "c", 3) == {"a": 1, 2: 2, "c": 3}
+    assert assoc(d, 2, 3) == {"a": 1, 2: 3}
+    assert assoc(d, "a", 0) == {"a": 0, 2: 2}
+    assert d == {"a": 1, 2: 2}
+
+    def assoc_OrderedDict(s, u, v):
+        s[u] = v
+        return s
+
+    assoc.add((OrderedDict, object, object), assoc_OrderedDict)
+
+    x = var()
+    d2 = OrderedDict(d)
+    assert assoc(d2, x, 3) is d2
+    assert assoc(d2, x, 3) == {"a": 1, 2: 2, x: 3}
+    assert assoc(d, x, 3) is not d
+
+
 def test_reify():
     x, y, z = var(), var(), var()
     s = {x: 1, y: 2, z: (x, y)}

unification/__init__.py

@@ -1,4 +1,4 @@
-from .core import unify, reify
+from .core import unify, reify, assoc
 from .more import unifiable
 from .variable import var, isvar, vars, variables, Var
 

unification/core.py

@@ -1,4 +1,4 @@
-from toolz import assoc
+from copy import copy
 from operator import length_hint
 from functools import partial
 from collections import OrderedDict, deque
@@ -14,6 +14,17 @@
 construction_sentinel = object()
 
 
+@dispatch(Mapping, object, object)
+def assoc(s, u, v):
+    """Add an entry to a `Mapping` and return it."""
+    if hasattr(s, "copy"):
+        s = s.copy()
+    else:
+        s = copy(s)  # pragma: no cover
+    s[u] = v
+    return s
+
+
 def stream_eval(z, res_filter=None):
     """Evaluate a stream of `_reify`/`_unify` results.