Updating `require`s with rewrite-clj

I've seen Emacs functionality that updates a Clojure file's require form to add a library based on the aliases used in the file, but I never tried to get something similar in Neovim because adding require forms manually doesn't bothered me very much. However, writing more scratch code lately, I wanted to move a little faster and thought I'd make a script to update requires. It was also a good excuse to finally try out rewrite-clj, which comes packaged with Babashka.

At a high level, the script reads stdin, parses it as Clojure forms, modifies it with rewrite-clj's zipper functions, then prints the modified code. I added a command to Neovim which pipes the current file through the script and updates the cursor position based on how much longer the new file is than the original.

In the Clojure script, let's parse the contents of stdin and get an S-expression of the code:

(require '[rewrite-clj.zip :as z])

(let [zloc (z/of-string (slurp *in*))
      sexpr (-> zloc z/up z/sexpr)])

To find the namespaces that are used, we can walk the S-expression with tree-seq and collect results:

(defn namespaces [form]
  (into #{}
        (comp
          (filter (every-pred symbol? namespace))
          (map (comp symbol namespace)))
        (tree-seq coll? seq form)))

(let [zloc (z/of-string (slurp *in*))
      sexpr (-> zloc z/up z/sexpr)
      nses (namespaces sexpr)])

The namespaces function isn't as robust as I'd like. It picks up symbol namespaces that don't need to be required, such as System and java.util.UUID, as well as namespaces on quoted symbols. I haven't though of a simple way to address that yet, short of specifically listing namespaces to ignore. If you have an idea, email me.

To support short aliases for libraries I commonly use, I created an import-aliases.edn file that looks like this:

{cli   clojure.tools.cli
 curl  babashka.curl
 deps  babashka.deps
 io    clojure.java.io
 json  cheshire.core
 pods  babashka.pods
 shell clojure.java.shell
 str   clojure.string
 walk  clojure.walk}

Let's create vectors to pass to require. We'll look up the namespace as an alias, and when no alias is found fall back to the original, presumably fully-qualified namespace name:

(def aliases
  (edn/read-string (slurp "import-aliases.edn")))

(defn require-vector [lib-or-alias]
  (let [lib (get aliases lib-or-alias lib-or-alias)
        a lib-or-alias]
    [lib :as a]))

(let [zloc (z/of-string (slurp *in*))
      sexpr (-> zloc z/up z/sexpr)
      nses (namespaces sexpr)
      needed (map require-vector nses)])

Now that we know what namespaces need to be required, there are four scenarios to handle:

the file has an ns form with a :require form inside it
the file has an ns form but no :require form inside it
the file has a top-level require form
the file has neither a ns form nor a require form

Here's a function that uses rewrite-clj's zipper functions to find list forms and determines which scenario to handle, where to insert new forms, and what libraries are already required:

(defn find-invoke [zloc prec]
  (z/find zloc
          (fn [zloc]
            (let [form (z/sexpr zloc)]
              (and (list? form)
                   (pred (first form)))))))

(defn existing-requires [[_ & forms]]
  (map (fn [form]
         (if (and (list? form)
                  (= 'quote (first form)))
           (second form)
           form))
       forms))

(defn find-insert-point [zloc]
  (if-let [ns-loc (find-invoke zloc #{'ns})]
    (if-let [require-loc (-> ns-loc z/down (find-invoke #{:require}))]
      [:ns-require-update
       require-loc
       (existing-requires (z/sexpr require-loc))]
      [:ns-require-create
       (-> ns-loc z/down z/rightmost)])
    (if-let [require-loc (find-invoke zloc #{'require})]
      [:require-update
       require-loc
       (existing-requires (z/sexpr require-loc))]
      [:require-create
       zloc])))

To generate a new set of require forms, we'll concatenate what we have with what's needed, normalize their forms, deduplicate, and sort:

(defn concat-requires [existing needed]
  (->> (concat existing needed)
       (map #(if (symbol? %) [% :as %] %))
       distinct
       (sort-by first)))

(let [zloc (z/of-string (slurp *in*))
      sexpr (-> zloc z/up z/sexpr)
      nses (namespaces sexpr)
      needed (map require-vector nses)
      [kind zloc existing] (find-insert-point zloc)
      new-requires (concat-requires existing needed)])

If you want a particular ordering, such as sorting all clojure libraries to be first, you can define a custom comparator and use it in sort-by:

(defn clojure-first [a b]
  (cond
    (and (str/starts-with? a "clojure.")
         (not (str/starts-with? b "clojure.")))
    , -1
    (and (not (str/starts-with? a "clojure."))
         (str/starts-with? b "clojure."))
    , 1
    :else
    , (compare a b)))

The final step is to edit the require form depending on which scenario we're in. Let's use a multimethod:

(defmulti edit-requires (fn [_ kind _] kind))

(let [zloc (z/of-string (slurp *in*))
      sexpr (-> zloc z/up z/sexpr)
      nses (namespaces sexpr)
      needed (map require-vector nses)
      [kind zloc existing] (find-insert-point zloc)
      new-requires (concat-requires existing needed)]
  (-> zloc
      (edit-requires kind new-requires)
      z/root-string
      println))

The four different scenarios need to be handled slightly differently, but they all need a list of forms interposed with a newline and some indentation. Here's a function that constructs the sequence of rewrite-clj nodes, with an optional transformation:

(require '[rewrite-clj.node :as n])

(defn lib-forms
  ([forms indent]
   (lib-forms forms identity indent))
  ([forms transform indent]
   (transduce
     (comp
       (map format-lib-form)
       (map transform)
       (interpose (n/whitespace-node (str "\n" indent))))
     conj forms)))

For the case where there's an existing ns form but it doesn't have a :require clause, we can insert the (:require ...) form to the right of the zipper's location, then insert the whitespace that should precede the form:

(defmethod edit-requires :ns-require-create [zloc _ forms]
  (let [nodes (lib-forms forms "            ")]
    (-> zloc
        (z/insert-right* (n/list-node (list* :require (n/spaces 1) nodes)))
        (z/insert-right* (n/whitespace-node "\n  ")))))

That's the reverse of the way you'd type it, adding the newline and whitespace first then all the new forms, but doing it this order avoids having to explicitly move right to the added node before adding the sibling node.

When updating an existing require form, we can replace the entire form with rewrite-clj.zip/replace*:

(defmethod edit-requires :ns-require-update [zloc _ forms]
  (let [nodes (lib-forms forms "            ")]
      (z/replace* zloc (n/list-node (list* :require (n/spaces 1) nodes)))))

When there is no ns form but there is a require form, we replace it the same way:

(defmethod edit-requires :require-update [zloc _ forms]
  (let [nodes (lib-forms forms n/quote-node "         ")]
    (z/replace* zloc (n/list-node (list* 'require (n/spaces 1) nodes)))))

And finally, when there's neither a ns nor a require form, we add a require form at the top of the code:

(defmethod edit-requires :require-create [zloc _ forms]
  (let [nodes (lib-forms forms n/quote-node "         ")]
    (-> zloc
        (z/insert-left* (n/list-node (list* 'require (n/spaces 1) nodes)))
        (z/insert-left* (n/newlines 2)))))

If you prefer, you could add a ns form, but for me it was simpler not to choose a name for the namespace.

The total script is only about 100 lines long, thanks largely to rewrite-clj. I'm still fiddling with how it works on real-world code, so if you have any suggestions, please email me.

Updating requires with rewrite-clj

Updating `require`s with rewrite-clj