update & group-by core functions (#298)

* update & group-by core functions

* Fix select-keys test

Author: chrisrink10

src/basilisp/core/__init__.lpy

@@ -730,6 +730,13 @@
   ([m k default]
    (basilisp.lang.runtime/get m k default)))
 
+(defn update
+  "Updates the value for key k in associative data structure m with the return value
+  from calling (f old-v & args). If m is nil, use an empty map. If k is not in m,
+  old-v will be nil."
+  [m k f & args]
+  (apply basilisp.lang.runtime/update m k f args))
+
 (defn key
   "Return the key from a map entry."
   [entry]
@@ -980,6 +987,21 @@
   [pred coll]
   [(take-while pred coll) (drop-while pred coll)])
 
+(defn group-by
+  "Return a map whose keys are the result of calling f on each element
+  in coll and whose values are vectors of the values which produced the
+  corresponding key, in the order they were added."
+  [f coll]
+  (if-not (seq coll)
+    {}
+    (reduce (fn [m v]
+              (let [group (f v)]
+                (if (contains? m group)
+                  (update m group conj v)
+                  (assoc m group [v]))))
+            {}
+            coll)))
+
 (defn interpose
   "Return a lazy sequence of elements of coll separated by sep. If
   coll is empty, return an empty sequence."
@@ -1794,12 +1816,14 @@
 (defn select-keys
   "Return a map with only the keys of m which are in ks."
   [m ks]
-  (reduce (fn [new-map k]
-            (if (contains? m k)
-              (assoc new-map k (get m k))
-              new-map))
-          {}
-          ks))
+  (if-not (seq ks)
+    {}
+    (reduce (fn [new-map k]
+              (if (contains? m k)
+                (assoc new-map k (get m k))
+                new-map))
+            {}
+            ks)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Destructuring Support ;;
@@ -2097,4 +2121,3 @@
                       (mapcat destructure))]
     `(let* [~@bindings]
        ~@body)))
-

src/basilisp/lang/runtime.py

@@ -613,6 +613,19 @@ def assoc(m, *kvs):
     raise TypeError(f"Object of type {type(m)} does not implement Associative interface")
 
 
+def update(m, k, f, *args):
+    """Updates the value for key k in associative data structure m with the return value from
+    calling f(old_v, *args). If m is None, use an empty map. If k is not in m, old_v will be
+    None."""
+    if m is None:
+        return lmap.Map.empty().assoc(k, f(None, *args))
+    if isinstance(m, lassoc.Associative):
+        old_v = m.entry(k)
+        new_v = f(old_v, *args)
+        return m.assoc(k, new_v)
+    raise TypeError(f"Object of type {type(m)} does not implement Associative interface")
+
+
 def conj(coll, *xs):
     """Conjoin xs to collection. New elements may be added in different positions
     depending on the type of coll. conj returns the same type as coll. If coll

tests/basilisp/core_test.py

@@ -719,6 +719,11 @@ def test_split_with():
     assert vec.v(llist.List.empty(), llist.l(2, 4, 6, 8)) == core.split_with(core.odd__Q__, vec.v(2, 4, 6, 8))
 
 
+def test_group_by():
+    assert lmap.Map.empty() == core.group_by(core.inc, vec.Vector.empty())
+    assert lmap.map({True: vec.v(1, 3), False: vec.v(2, 4)}) == core.group_by(core.odd__Q__, vec.v(1, 2, 3, 4))
+
+
 def test_interpose():
     assert llist.List.empty() == core.interpose(",", vec.Vector.empty())
     assert llist.l("hi") == core.interpose(",", vec.v("hi"))
@@ -813,7 +818,7 @@ def test_re_seq():
 
 def test_select_keys():
     assert lmap.Map.empty() == core.select_keys(
-        lmap.Map.empty(), vec.v(vec.Vector.empty()))
+        lmap.Map.empty(), vec.Vector.empty())
     assert lmap.Map.empty() == core.select_keys(
         lmap.Map.empty(), vec.v(kw.keyword('a'), kw.keyword('b')))
     assert lmap.map({kw.keyword('a'): "a", kw.keyword('b'): "b"}) == core.select_keys(

tests/basilisp/runtime_test.py

@@ -181,6 +181,24 @@ def test_assoc():
         runtime.assoc(llist.List.empty(), 1, "a")
 
 
+def test_update():
+    assert lmap.map({"a": 1}) == runtime.update(None, "a", lambda _: 1)
+    assert lmap.map({"a": 50}) == runtime.update(None, "a", lambda _, x: x, 50)
+
+    assert lmap.map({"a": 2}) == runtime.update(lmap.map({"a": 1}), "a", lambda x: x + 1)
+    assert lmap.map({"a": 4}) == runtime.update(lmap.map({"a": 1}), "a", lambda x, y: x * y + 1, 3)
+
+    assert lmap.map({"a": 1, "b": "string"}) == runtime.update(lmap.map({"a": 1}), "b", lambda _: "string")
+    assert lmap.map({"a": 1, "b": "string"}) == runtime.update(lmap.map({"a": 1, "b": 583}), "b", lambda _: "string")
+
+    assert vec.v("a") == runtime.update(vec.Vector.empty(), 0, lambda _: "a")
+    assert vec.v("yay", "b") == runtime.update(vec.v("a", "b"), 0, lambda x: f"y{x}y")
+    assert vec.v("a", "boy") == runtime.update(vec.v("a", "b"), 1, lambda x, y: f"{x}{y}", "oy")
+
+    with pytest.raises(TypeError):
+        runtime.update(llist.List.empty(), 1, lambda _: "y")
+
+
 def test_conj():
     assert llist.l(1) == runtime.conj(None, 1)
     assert llist.l(3, 2, 1) == runtime.conj(None, 1, 2, 3)