ANSI escape codes don't count towards a cell's width (#103)

Fixes #102
babashka · Jul 29, 2024 · de2569f · de2569f
1 parent 836b15f
commit de2569f
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ For breaking changes, check [here](#breaking-changes).
 
 [Babashka CLI](https://github.com/babashka/cli): turn Clojure functions into CLIs!
 
+## Unreleased
+
+- Fix [#102](https://github.com/babashka/cli/issues/102): `format-table` correctly pads cells containing ANSI escape codes
+
 ## v0.8.60 (2024-07-23)
 
 - Fix [#98](https://github.com/babashka/cli/issues/98): internal options should not interfere with :restrict

diff --git a/src/babashka/cli.cljc b/src/babashka/cli.cljc
@@ -521,17 +521,24 @@
 (defn- kw->str [kw]
   (subs (str kw) 1))
 
-(defn pad [len s] (str s (apply str (repeat (- len (count s)) " "))))
+(defn- str-width
+  "Width of `s` when printed, i.e. without ANSI escape codes."
+  [s]
+  (let [strip-escape-codes #(str/replace %
+                                         (re-pattern "(\\x9B|\\x1B\\[)[0-?]*[ -\\/]*[@-~]") "")]
+    (count (strip-escape-codes s))))
+
+(defn pad [len s] (str s (apply str (repeat (- len (str-width s)) " "))))
 
 (defn pad-cells [rows]
-  (let [widths (reduce
-                (fn [widths row]
-                  (map max (map count row) widths)) (repeat 0) rows)
+  (let [widths  (reduce
+                 (fn [widths row]
+                   (map max (map str-width row) widths)) (repeat 0) rows)
         pad-row (fn [row]
-                  (map (fn [width col] (pad width col)) widths row))]
+                  (map (fn [width cell] (pad width cell)) widths row))]
     (map pad-row rows)))
 
-(defn format-table [{:keys [rows indent]}]
+(defn format-table [{:keys [rows indent] :or {indent 2}}]
   (let [rows (pad-cells rows)
         fmt-row (fn [leader divider trailer row]
                   (str leader

diff --git a/test/babashka/cli_test.cljc b/test/babashka/cli_test.cljc
@@ -487,6 +487,21 @@
                                           :desc "Barbarbar" :default-desc "Mos def"}}}))
              :indent 2})))))
 
+(deftest format-table-test
+  (let [contains-row-matching (fn [re table]
+                                (let [rows (str/split-lines table)]
+                                  (is (some #(re-find re %) rows)
+                                      (str "expected " (pr-str rows)
+                                           " to contain a row matching " (pr-str re)))))]
+    (testing "ANSI escape codes don't count towards a cell's width"
+      (let [table (cli/format-table {:rows [["widest" "<- sets column width to 6"]
+                                            ["\033[31mfoo\033[0m" "<- needs 3+1 padding"]
+                                            ["bar" "<- needs 3+1 padding"]]})]
+        (contains-row-matching #"\033\[31mfoo\033\[0m    <-"
+                               table)
+        (contains-row-matching #"bar    <-"
+                               table)))))
+
 (deftest require-test
   (is (thrown-with-msg?
        Exception #"Required option: :bar"