Cherry-pick "associative syntax alignment" attempt

Source: weavejester#77

README.md

@@ -133,6 +133,13 @@ selectively enabled or disabled:
   true if cljfmt should collapse consecutive blank lines. This will
   convert `(foo)\n\n\n(bar)` to `(foo)\n\n(bar)`. Defaults to true.
 
+* `:align-associative?` -
+  true if cljfmt should left align the values of maps and binding
+  special forms (let, loop, binding). This will convert
+  `{:foo 1\n:barbaz 2}` to `{:foo    1\n :barbaz 2}`
+  and `(let [foo 1\n barbaz 2])` to `(let [foo    1\n      barbaz 2])`.
+  Defaults to false.
+
 * `:remove-multiple-non-indenting-spaces?` -
   true if cljfmt should remove multiple non indenting spaces. This
   will convert <code>{:a 1&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;:b 2}</code> to `{:a 1 :b 2}`. Defaults to false.

cljfmt/src/cljfmt/core.cljc

@@ -2,7 +2,9 @@
   #?(:clj (:refer-clojure :exclude [reader-conditional?]))
   (:require #?(:clj [clojure.java.io :as io])
             [clojure.string :as str]
+            [clojure.pprint]
             [rewrite-clj.node :as n]
+            [clojure.zip :as zip]
             [rewrite-clj.parser :as p]
             [rewrite-clj.zip :as z])
   #?(:clj (:import java.util.regex.Pattern)
@@ -391,6 +393,119 @@
 (defn remove-trailing-whitespace [form]
   (transform form edit-all trailing-whitespace? z/remove*))
 
+;; added
+
+(defn- append-newline-if-absent [zloc]
+  (if (or (-> zloc z/right* skip-whitespace line-break?)
+          (z/rightmost? zloc))
+    zloc
+    (do
+      (zip/insert-right zloc (n/newlines 1)))))
+
+(defn- map-odd-seq
+  "Applies f to all oddly-indexed nodes."
+  [f zloc]
+  (loop [loc (z/down zloc)
+         parent zloc]
+    (if-not (and loc (z/node loc))
+      parent
+      (if-let [v (f loc)]
+        (recur (z/right (z/right v)) (z/up v))
+        (recur (z/right (z/right loc)) parent)))))
+
+(defn- map-even-seq
+  "Applies f to all evenly-indexed nodes."
+  [f zloc]
+  (loop [loc   (z/right (z/down zloc))
+         parent zloc]
+    (if-not (and loc (z/node loc))
+      parent
+      (if-let [v (f loc)]
+        (recur (z/right (z/right v)) (z/up v))
+        (recur (z/right (z/right loc)) parent)))))
+
+(defn- add-map-newlines [zloc]
+  (map-even-seq #(cond-> % (complement z/rightmost?)
+                         append-newline-if-absent) zloc))
+
+(defn- add-binding-newlines [zloc]
+  (map-even-seq append-newline-if-absent zloc))
+
+(defn- update-in-path [[node path :as loc] k f]
+  (let [v (get path k)]
+    (if (seq v)
+      (with-meta
+        [node (assoc path k (f v) :changed? true)]
+        (meta loc))
+      loc)))
+
+(defn- remove-right
+  [loc]
+  (update-in-path loc :r next))
+
+(defn- *remove-right-while
+  [zloc p?]
+  (loop [zloc zloc]
+    (if-let [rloc (z/right* zloc)]
+      (if (p? rloc)
+        (recur (remove-right zloc))
+        zloc)
+      zloc)))
+
+(defn- align-seq-value [zloc max-length]
+  (let [key-length (-> zloc z/sexpr str count)
+        width      (- max-length key-length)
+        zloc       (*remove-right-while zloc clojure-whitespace?)]
+    (zip/insert-right zloc (whitespace (inc width)))))
+
+(defn- align-map [zloc]
+  (let [key-list       (-> zloc z/sexpr keys)
+        max-key-length (apply max (map #(-> % str count) key-list))]
+    (map-odd-seq #(align-seq-value % max-key-length) zloc)))
+
+(defn- align-binding [zloc]
+  (let [vec-sexpr    (z/sexpr zloc)
+        odd-elements (take-nth 2 vec-sexpr)
+        max-length   (apply max (map #(-> % str count) odd-elements))]
+    (map-odd-seq #(align-seq-value % max-length) zloc)))
+
+(defn- align-elements [zloc]
+  (if (z/map? zloc)
+    (-> zloc align-map add-map-newlines)
+    (-> zloc align-binding add-binding-newlines)))
+
+(def ^:private binding-keywords
+  #{"doseq" "let" "loop" "binding" "with-open" "go-loop" "if-let" "when-some"
+    "if-some" "for" "with-local-vars" "with-redefs"})
+
+(defn- binding? [zloc]
+  (and (z/vector? zloc)
+       (-> zloc z/sexpr count even?)
+       (->> zloc
+            z/left
+            z/string
+            (contains? binding-keywords))))
+
+(defn- align-binding? [zloc]
+  (and (binding? zloc)
+       (-> zloc z/sexpr count (> 2))))
+
+(defn- empty-seq? [zloc]
+  (if (z/map? zloc)
+    (-> zloc z/sexpr empty?)
+    false))
+
+(defn- align-map? [zloc]
+  (and (z/map? zloc)
+       (not (empty-seq? zloc))))
+
+(defn- align-elements? [zloc]
+  (or (align-binding? zloc)
+      (align-map? zloc)))
+
+(defn align-collection-elements [form]
+  (transform form edit-all align-elements? align-elements))
+
 (defn- replace-with-one-space [zloc]
   (z/replace* zloc (whitespace 1)))
 
@@ -493,8 +608,9 @@
    :remove-trailing-whitespace?           true
    :split-keypairs-over-multiple-lines?   false
    :sort-ns-references?                   false
-   :indents   default-indents
-   :alias-map {}})
+   :indents                               default-indents
+   :align-associative?                    false
+   :alias-map                             {}})
 
 (defn reformat-form
   ([form]
@@ -512,6 +628,8 @@
            remove-surrounding-whitespace)
          (cond-> (:insert-missing-whitespace? opts)
            insert-missing-whitespace)
+         (cond-> (:align-associative? opts)
+           align-collection-elements)
          (cond-> (:remove-multiple-non-indenting-spaces? opts)
            remove-multiple-non-indenting-spaces)
          (cond-> (:indentation? opts)

cljfmt/src/cljfmt/main.clj

@@ -187,6 +187,9 @@
    [nil "--[no-]indentation"
     :default (:indentation? cljfmt/default-options)
     :id :indentation?]
+   [nil "--[no-]align-associative"
+    :default (:align-associative? cljfmt/default-options)
+    :id :align-associative?]
    [nil "--[no-]remove-multiple-non-indenting-spaces"
     :default (:remove-multiple-non-indenting-spaces? cljfmt/default-options)
     :id :remove-multiple-non-indenting-spaces?]

cljfmt/test/cljfmt/core_test.cljc

@@ -706,10 +706,7 @@
          ["(foo bar)"]))
     (is (reformats-to?
          ["[ 1 2 3 ]"]
-         ["[1 2 3]"]))
-    (is (reformats-to?
-         ["{  :x 1, :y 2 }"]
-         ["{:x 1, :y 2}"])))
+         ["[1 2 3]"])))
 
   (testing "surrounding newlines removed"
     (is (reformats-to?
@@ -956,6 +953,37 @@
         ""])
       "blank lines at end of string are not affected by :remove-consecutive-blank-lines?"))
 
+(deftest test-align-associative
+  (testing "binding alignment"
+    (is (= (reformat-string "(let [foo 1\n barbaz 2])" {:align-associative? true})
+           "(let [foo    1\n      barbaz 2])"))
+    (is (= (reformat-string "(let [foo 1\n barbaz 2 qux 3])")
+           "(let [foo    1\n      barbaz 2\n      qux    3])")))
+
+  (testing "binding alignment preserves comments"
+    (is (= (reformat-string "(let [foo 1 ;; test 1\n barbaz 2])" {:align-associative? true})
+           "(let [foo    1 ;; test 1\n      barbaz 2])")))
+
+  (testing "map alignment"
+    (is (= (reformat-string "{:foo 1\n:barbaz 2}" {:align-associative? true})
+           "{:foo    1\n :barbaz 2}"))
+    (is (= (reformat-string "{:foo\n  1\n:baz 2}" {:align-associative? true})
+           "{:foo 1\n :baz 2}"))
+    (is (= (reformat-string "{:bar\n {:qux 1\n  :quux 2}}" {:align-associative? true})
+           "{:bar {:qux  1\n       :quux 2}}"))
+    (is (= (reformat-string "{:foo 1\n (baz quux) 2}" {:align-associative? true})
+           "{:foo       1\n (baz quux) 2}"))
+    (is (= (reformat-string "{:foo (bar)\n :quux (baz)}" {:align-associative? true})
+           "{:foo  (bar)\n :quux (baz)}"))
+    (is (= (reformat-string "[{:foo 1 :bar 2}\n{:foo 1 :bar 2}]" {:align-associative? true})
+           "[{:foo 1\n  :bar 2}\n {:foo 1\n  :bar 2}]")))
+
+  (testing "map alignment preserves comments"
+    (is (= (reformat-string "{:foo 1 ;; test 1\n:barbaz 2}" {:align-associative? true})
+           "{:foo    1 ;; test 1\n :barbaz 2}"))
+    (is (= (reformat-string "{:foo 1 ;; test 1\n:barbaz 2\n:fuz 1}" {:align-associative? true})
+           "{:foo    1 ;; test 1\n :barbaz 2\n :fuz    1}"))))
+
 (deftest test-trailing-whitespace
   (testing "trailing-whitespace removed"
     (is (reformats-to?