Associative utility functions

src/basilisp/core.lpy

@@ -1272,13 +1272,55 @@
   ([m k default]
    (basilisp.lang.runtime/get m k default)))
 
+(defn assoc-in
+  "Associate value in a nested associative structure, with ks as a sequence of
+  keys and v as the new value. If no map exists for any key in ks, a new empty
+  map will be created."
+  [m ks v]
+  (let [fk  (first ks)
+        rks (rest ks)]
+    (if (seq rks)
+      (basilisp.lang.runtime/assoc m fk (assoc-in (get m fk) rks v))
+      (basilisp.lang.runtime/assoc m fk v))))
+
+(defn get-in
+  "Return the entry of an associative data structure addressed by the sequence of
+  keys ks or default (default: nil) if the value is not found."
+  ([m ks]
+   (get-in m ks nil))
+  ([m ks default]
+   (let [fk  (first ks)
+         rks (rest ks)]
+     (if (seq rks)
+       (if-let [child-map (basilisp.lang.runtime/get m fk)]
+         (get-in child-map rks default)
+         default)
+       (basilisp.lang.runtime/get m fk 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 update-in
+  "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 ks f & args]
+  (let [fk  (first ks)
+        rks (rest ks)]
+    (if (seq rks)
+      (basilisp.lang.runtime/assoc m
+                                   fk
+                                   (apply update-in
+                                          (or (basilisp.lang.runtime/get m fk) {})
+                                          rks
+                                          f
+                                          args))
+      (apply basilisp.lang.runtime/update m fk f args))))
+
 (defn map-entry
   "With one argument, coerce the input to a map entry. With two arguments, return a
   map entry containing key and value."
@@ -1292,6 +1334,12 @@
   ([k v]
    (basilisp.lang.map.MapEntry/of k v)))
 
+(defn find
+  "Find the map entry of k in m, if it exists. Return nil otherwise."
+  [m k]
+  (when-let [v (get m k)]
+    (map-entry k v)))
+
 (defn key
   "Return the key from a map entry."
   [entry]

tests/basilisp/core_test.py

@@ -1063,6 +1063,59 @@ def test_get(self):
         assert 1 == core.get(vec.v(1, 2, 3), -3)
         assert None is core.get(vec.v(1, 2, 3), -4)
 
+    def test_find(self):
+        assert None is core.find(None, "a")
+        assert core.map_entry("a", 1) == core.find(lmap.map({"a": 1}), "a")
+        assert None is core.find(lmap.map({"a": 1}), "b")
+        assert core.map_entry(0, 1) == core.find(vec.v(1, 2, 3), 0)
+        assert core.map_entry(1, 2) == core.find(vec.v(1, 2, 3), 1)
+        assert core.map_entry(2, 3) == core.find(vec.v(1, 2, 3), 2)
+        assert None is core.find(vec.v(1, 2, 3), 3)
+
+    def test_get_in(self):
+        assert 1 == core.get_in(lmap.map({"a": 1}), vec.v("a"))
+        assert None is core.get_in(lmap.map({"a": 1}), vec.v("b"))
+        assert 2 == core.get_in(lmap.map({"a": 1}), vec.v("b"), 2)
+
+        assert lmap.map({"b": lmap.map({"c": 3})}) == core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}), vec.v("a")
+        )
+        assert lmap.map({"c": 3}) == core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}), vec.v("a", "b")
+        )
+        assert 3 == core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}), vec.v("a", "b", "c")
+        )
+        assert None is core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}), vec.v("a", "b", "f")
+        )
+        assert None is core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}), vec.v("a", "e", "c")
+        )
+        assert "Not Found" == core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}),
+            vec.v("a", "b", "f"),
+            "Not Found",
+        )
+        assert "Not Found" == core.get_in(
+            lmap.map({"a": lmap.map({"b": lmap.map({"c": 3})})}),
+            vec.v("a", "e", "c"),
+            "Not Found",
+        )
+
+        assert "b" == core.get_in(vec.v("a", "b", "c"), vec.v(1))
+        assert "t" == core.get_in(
+            vec.v("a", vec.v("q", "r", "s", "t"), "c"), vec.v(1, 3)
+        )
+        assert "x" == core.get_in(
+            vec.v("a", vec.v("q", "r", "s", lmap.map({"w": "x"})), "c"),
+            vec.v(1, 3, "w"),
+        )
+        assert None is core.get_in(
+            vec.v("a", vec.v("q", "r", "s", lmap.map({"w": "x"})), "c"),
+            vec.v(1, 3, "v"),
+        )
+
     def test_keys(self):
         assert None is core.keys(lmap.map({}))
         assert llist.l("a") == core.keys(lmap.map({"a": 1}))