Add tree walking functions in basilisp.walk

src/basilisp/lang/reader.py

@@ -34,13 +34,14 @@
 import basilisp.lang.symbol as symbol
 import basilisp.lang.util as langutil
 import basilisp.lang.vector as vector
-import basilisp.walker as walk
 from basilisp.lang.interfaces import (
     ILispObject,
     ILookup,
     IMeta,
     IPersistentList,
+    IPersistentMap,
     IPersistentSet,
+    IPersistentVector,
     IRecord,
     IType,
 )
@@ -808,6 +809,28 @@ def _read_meta(ctx: ReaderContext) -> IMeta:
         )
 
 
+def _walk(inner_f, outer_f, form):
+    """Walk an arbitrary, possibly nested data structure, applying inner_f to each
+    element of form and then applying outer_f to the resulting form."""
+    if isinstance(form, IPersistentList):
+        return outer_f(llist.list(map(inner_f, form)))
+    elif isinstance(form, IPersistentVector):
+        return outer_f(vector.vector(map(inner_f, form)))
+    elif isinstance(form, IPersistentMap):
+        return outer_f(lmap.from_entries(map(inner_f, form)))
+    elif isinstance(form, IPersistentSet):
+        return outer_f(lset.set(map(inner_f, form)))
+    else:
+        return outer_f(form)
+
+
+def _postwalk(f, form):
+    """"Walk form using depth-first, post-order traversal, applying f to each form
+    and replacing form with its result."""
+    inner_f = functools.partial(_postwalk, f)
+    return _walk(inner_f, f, form)
+
+
 @_with_loc
 def _read_function(ctx: ReaderContext) -> llist.List:
     """Read a function reader macro from the input stream."""
@@ -841,7 +864,7 @@ def identify_and_replace(f):
                     return sym_replacement(arg_num)
         return f
 
-    body = walk.postwalk(identify_and_replace, form) if len(form) > 0 else None
+    body = _postwalk(identify_and_replace, form) if len(form) > 0 else None
 
     arg_list: List[symbol.Symbol] = []
     numbered_args = sorted(map(int, filter(lambda k: k != "rest", arg_set)))

src/basilisp/walk.lpy

@@ -0,0 +1,116 @@
+(ns basilisp.walk)
+
+(defn walk
+  "Walk an arbitrary, possibly nested data structure, applying inner to each
+  element of form and then applying outer to the resulting form.
+
+  All built in data structures are supported.
+
+  Lazy sequences will be completely consumed (and thus may not be infinite)."
+  [inner outer form]
+  (cond
+    (list? form)
+    (outer (apply list (map inner form)))
+
+    (map-entry? form)
+    (outer (map-entry (inner (key form)) (inner (val form))))
+
+    (seq? form)
+    (outer (doall (map inner form)))
+
+    (vector? form)
+    (outer (apply vector (map inner form)))
+
+    (map? form)
+    (outer (apply hash-map (mapcat inner form)))
+
+    (set? form)
+    (outer (apply hash-set (map inner form)))
+
+    (record? form)
+    (outer (reduce (fn [rec field]
+                     (conj rec (inner field)))
+                   form
+                   form))
+
+    :else
+    (outer form)))
+
+(defn postwalk
+  "Walk form using depth-first, post-order traversal, applying f to each form
+  and replacing form with its result.
+
+  All built in data structures are supported.
+
+  Lazy sequences will be completely consumed (and thus may not be infinite)."
+  [f form]
+  (walk (partial postwalk f) f form))
+
+(defn prewalk
+  "Walk form using depth-first, pre-order traversal, applying f to each form
+  and replacing form with its result.
+
+  All built in data structures are supported.
+
+  Lazy sequences will be completely consumed (and thus may not be infinite)."
+  [f form]
+  (walk (partial prewalk f) identity (f form)))
+
+(defn postwalk-replace
+  "Recursively walk through form as by postwalk, replacing elements appearing
+  as keys in replacements with the corresponding values."
+  [replacements form]
+  (postwalk #(if-let [newv (get replacements %)]
+               newv
+               %)
+            form))
+
+(defn prewalk-replace
+  "Recursively walk through form as by prewalk, replacing elements appearing
+  as keys in replacements with the corresponding values."
+  [replacements form]
+  (prewalk #(if-let [newv (get replacements %)]
+              newv
+              %)
+           form))
+
+(defn postwalk-demo
+  "Print each element as it is walked as by postwalk."
+  [form]
+  (postwalk #(do (println (str "Walked: " %)) %) form))
+
+(defn prewalk-demo
+  "Print each element as it is walked as by postwalk."
+  [form]
+  (prewalk #(do (println (str "Walked: " %)) %) form))
+
+(defn keywordize-keys
+  "Recursively walk form, transforming string keys into keywords in any maps."
+  [form]
+  (postwalk (fn [v]
+              (if (map? v)
+                (->> v
+                     (mapcat (fn [[k v]] [(cond-> k (string? k) (keyword)) v]))
+                     (apply hash-map))
+                v))
+            form))
+
+(defn stringify-keys
+  "Recursively walk form, transforming keyword keys into strings in any maps."
+  [form]
+  (postwalk (fn [v]
+              (if (map? v)
+                (->> v
+                     (mapcat (fn [[k v]] [(cond-> k (keyword? k) (name)) v]))
+                     (apply hash-map))
+                v))
+            form))
+
+(defn macroexpand-all
+  "Recursively macroexpand all eligible forms contained in form."
+  [form]
+  (prewalk (fn [v]
+             (if (seq? v)
+               (macroexpand v)
+               v))
+           form))

tests/basilisp/test_walk.lpy

@@ -0,0 +1,112 @@
+(ns tests.basilisp.walk-test
+  (:require
+   [basilisp.test :refer [deftest are is testing]]
+   [basilisp.walk :as walk]))
+
+(deftest postwalk-replace-test
+  (is (= [:c :d] (walk/postwalk-replace {:a 1 :b 2} [:c :d])))
+  (is (= [1 2] (walk/postwalk-replace {:a 1 :b 2} [:a :b])))
+  (is (= [1 2 :c] (walk/postwalk-replace {:a 1 :b 2} [:a :b :c])))
+  (is (= [1 2 [1 2] :c] (walk/postwalk-replace {:a 1 :b 2} [:a :b [:a :b] :c])))
+  (is (= {:NIL 4, :a 1, :b :NIL, :c 3}
+         (walk/postwalk-replace {nil :NIL} {:a 1, :b nil, :c 3, nil 4}))))
+
+(deftest prewalk-replace-test
+  (is (= [:c :d] (walk/prewalk-replace {:a 1 :b 2} [:c :d])))
+  (is (= [1 2] (walk/prewalk-replace {:a 1 :b 2} [:a :b])))
+  (is (= [1 2 :c] (walk/prewalk-replace {:a 1 :b 2} [:a :b :c])))
+  (is (= [1 2 [1 2] :c] (walk/prewalk-replace {:a 1 :b 2} [:a :b [:a :b] :c])))
+  (is (= {:NIL 4, :a 1, :b :NIL, :c 3}
+         (walk/postwalk-replace {nil :NIL} {:a 1, :b nil, :c 3, nil 4}))))
+
+(deftest keywordize-keys-test
+  (testing "maps and nested maps"
+    (is (= {} (walk/keywordize-keys {})))
+    (is (= {:a 1 :b 2} (walk/keywordize-keys {"a" 1 "b" 2})))
+    (is (= {:a 1 :b 2 :c 3} (walk/keywordize-keys {"a" 1 "b" 2 :c 3})))
+    (is (= {:a 1 :c 3 4 :d}
+           (walk/keywordize-keys {"a" 1 :c 3 4 :d})))
+    (is (= {:a 1 :c 3 4 :d :nested {:e 5 'b 6}}
+           (walk/keywordize-keys {"a" 1 :c 3 4 :d "nested" {"e" 5 'b 6}}))))
+
+  (testing "maps in lists and seqs"
+    (is (= '({:a 1 :b 2}) (walk/keywordize-keys '({"a" 1 "b" 2}))))
+    (is (= '({:a 1 :b 2 :c 3}) (walk/keywordize-keys '({"a" 1 "b" 2 :c 3}))))
+    (is (= '({:a 1 :c 3 4 :d})
+           (walk/keywordize-keys '({"a" 1 :c 3 4 :d}))))
+    (is (= '({:a 1 :c 3 4 :d :nested {:e 5 'b 6}})
+           (walk/keywordize-keys '({"a" 1 :c 3 4 :d "nested" {"e" 5 'b 6}}))))
+    (is (= '({:a 1} {:a 2} {:a 3})
+           (walk/keywordize-keys (map #(hash-map "a" %) (range 1 4))))))
+
+  (testing "maps in sets"
+    (is (= #{{:a 1 :b 2}} (walk/keywordize-keys #{{"a" 1 "b" 2}})))
+    (is (= #{{:a 1 :b 2 :c 3}} (walk/keywordize-keys #{{"a" 1 "b" 2 :c 3}})))
+    (is (= #{{:a 1 :c 3 4 :d}}
+           (walk/keywordize-keys #{{"a" 1 :c 3 4 :d}})))
+    (is (= #{{:a 1 :c 3 4 :d :nested {:e 5 'b 6}}}
+           (walk/keywordize-keys #{{"a" 1 :c 3 4 :d "nested" {"e" 5 'b 6}}}))))
+
+  (testing "maps in vectors"
+    (is (= [{:a 1 :b 2}] (walk/keywordize-keys [{"a" 1 "b" 2}])))
+    (is (= [{:a 1 :b 2 :c 3}] (walk/keywordize-keys [{"a" 1 "b" 2 :c 3}])))
+    (is (= [{:a 1 :c 3 4 :d}]
+           (walk/keywordize-keys [{"a" 1 :c 3 4 :d}])))
+    (is (= [{:a 1 :c 3 4 :d :nested {:e 5 'b 6}}]
+           (walk/keywordize-keys [{"a" 1 :c 3 4 :d "nested" {"e" 5 'b 6}}])))))
+
+(deftest stringify-keys-test
+  (testing "maps and nested maps"
+    (is (= {} (walk/stringify-keys {})))
+    (is (= {"a" 1 "b" 2} (walk/stringify-keys {:a 1 :b 2})))
+    (is (= {"a" 1 "b" 2 "c" 3} (walk/stringify-keys {"a" 1 "b" 2 :c 3})))
+    (is (= {"a" 1 "c" 3 4 :d}
+           (walk/stringify-keys {"a" 1 :c 3 4 :d})))
+    (is (= {"a" 1 "c" 3 4 :d "nested" {"e" 5 'b 6}}
+           (walk/stringify-keys {"a" 1 :c 3 4 :d "nested" {:e 5 'b 6}}))))
+
+  (testing "maps in lists and seqs"
+    (is (= '({"a" 1 "b" 2}) (walk/stringify-keys '({:a 1 :b 2}))))
+    (is (= '({"a" 1 "b" 2 "c" 3}) (walk/stringify-keys '({"a" 1 "b" 2 :c 3}))))
+    (is (= '({"a" 1 "c" 3 4 :d})
+           (walk/stringify-keys '({"a" 1 :c 3 4 :d}))))
+    (is (= '({"a" 1 "c" 3 4 :d "nested" {"e" 5 'b 6}})
+           (walk/stringify-keys '({"a" 1 :c 3 4 :d "nested" {:e 5 'b 6}}))))
+    (is (= '({"a" 1} {"a" 2} {"a" 3})
+           (walk/stringify-keys (map #(hash-map :a %) (range 1 4))))))
+
+  (testing "maps in sets"
+    (is (= #{{"a" 1 "b" 2}} (walk/stringify-keys #{{:a 1 :b 2}})))
+    (is (= #{{"a" 1 "b" 2 "c" 3}} (walk/stringify-keys #{{"a" 1 "b" 2 :c 3}})))
+    (is (= #{{"a" 1 "c" 3 4 :d}}
+           (walk/stringify-keys #{{"a" 1 :c 3 4 :d}})))
+    (is (= #{{"a" 1 "c" 3 4 :d "nested" {"e" 5 'b 6}}}
+           (walk/stringify-keys #{{"a" 1 :c 3 4 :d "nested" {:e 5 'b 6}}}))))
+
+  (testing "maps in vectors"
+    (is (= [{"a" 1 "b" 2}] (walk/stringify-keys [{:a 1 :b 2}])))
+    (is (= [{"a" 1 "b" 2 "c" 3}] (walk/stringify-keys [{"a" 1 "b" 2 :c 3}])))
+    (is (= [{"a" 1 "c" 3 4 :d}]
+           (walk/stringify-keys [{"a" 1 :c 3 4 :d}])))
+    (is (= [{"a" 1 "c" 3 4 :d "nested" {"e" 5 'b 6}}]
+           (walk/stringify-keys [{"a" 1 :c 3 4 :d "nested" {:e 5 'b 6}}])))))
+
+(defmacro plus [n1 n2]
+  `(+ ~n1 ~n2))
+
+(defmacro pl [p1 p2]
+  `(plus ~p1 ~p2))
+
+(defmacro minus [m1 m2]
+  `(- ~m1 ~m2))
+
+(defmacro calc [c1 c2]
+  `(pl ~c1 (minus ~c1 ~c2)))
+
+(deftest macroexpand-all-test
+  (is (= '(tests.basilisp.walk-test/pl 20 (tests.basilisp.walk-test/minus 20 30))
+         (macroexpand-1 '(calc 20 30))))
+  (is (= '(basilisp.core/+ 20 (tests.basilisp.walk-test/minus 20 30))
+         (macroexpand '(calc 20 30))))
+  (is (= '(basilisp.core/+ 20 (basilisp.core/- 20 30))
+         (walk/macroexpand-all '(calc 20 30)))))