s/keys+, an s/keys variant that allows inline respec-ing of a key to narrow the range of its type

keys_plus.cljc

(ns com.breezeehr.specs.keys-plus
  "Variants of clojure.spec/keys and keys* that allow additional inline spec-ing."
  (:refer-clojure :exclude [keys])
  (:require [clojure.core :as c]
            [clojure.spec.alpha :as s]
            [clojure.spec.gen.alpha :as gen]
            [clojure.walk :as walk])
  #?(:cljs (:require-macros [com.breezeehr.specs.keys-plus]))
  #?(:clj (:import (java.util UUID))))

;; Get around private-fn restrictions
#?(:cljs
   (do (defn s-specize [s] (s/specize s))
       (defn s-pvalid?
         ([pred x] (s/pvalid? pred x))
         ([pred x form] (s/pvalid? pred x form)))
       (defn s-explain-1 [form pred path via in v]
         (s/explain-1 form pred path via in v))
       (defn s-inck [m k] (s/inck m k))
       (defn s-gensub [spec overrides path rmap form]
         (s/gensub spec overrides path rmap form))
       (defn s-recur-limit? [rmap id path k]
         (s/recur-limit? rmap id path k)))
   :default
   (do (def s-specize #'s/specize)
       (def s-pvalid? #'s/pvalid?)
       (def s-explain-1 #'s/explain-1)
       (def s-inck #'s/inck)
       (def s-gensub #'s/gensub)
       (def s-recur-limit? #'s/recur-limit?)))


(deftype Reform [spec form gfn ___meta]
  #?@(:cljs
      [IMeta (-meta [_] ___meta)
       IWithMeta (-with-meta [_ meta] (Reform. spec form gfn meta))]

      :default
      [clojure.lang.IMeta (meta [_] ___meta)
       clojure.lang.IObj (withMeta [_ meta] (Reform. spec form gfn meta))])

  s/Specize
  (specize* [self] self)
  (specize* [self _] self)

  s/Spec
  (conform* [_ x] (s/conform* spec x))
  (unform* [_ x] (s/unform* spec x))
  (explain* [_ path via in x] (s/explain* spec path via in x))
  (gen* [_ a b c] (if gfn (gfn) (s/gen* spec a b c)))
  (with-gen* [_ gfn] (Reform. spec form gfn ___meta))
  (describe* [_] form))

(defmacro with-describe
  "Return a spec like `s` except its s/describe is `form`."
  [s form]
  `(let [s# (s-specize ~s)]
     (->Reform s#
       #?(;; self-hosted cljs
          :cljs '~(s/res &env form)
          :clj  '~(#'s/res form))

       nil (meta s#))))

(defmacro together
  "Evaluates expressions one at a time, left to right. Returns true if all forms
  are falsey, or the last truthy value if all forms are truthy. Short-circuits
  and returns something falsey if the expressions are a mix of truthy and falsey.
  (together) returns true."
  ([] true)
  ([x] x)
  ([x y] `(if ~x ~y (not ~y)))
  ([x y & rest]
   `(if ~x
      (and ~y ~@rest)
      (not (or ~y ~@rest)))))

(defn map+-spec-impl
  [{:keys [req-un opt-un keys-pred pred-exprs opt-keys req-specs req req-keys
           opt-specs pred-forms conform-override conform-override-form opt gfn]
    :as   argm}]
  (let [conform-override (or conform-override {})
        conform-override-form (or conform-override-form {})
        k->s             (zipmap (concat req-keys opt-keys) (concat req-specs opt-specs))
        keys->specnames  #(or (k->s %) %)
        id #?(:clj       (UUID/randomUUID)
              :cljs      (random-uuid))]
    (reify
      s/Specize
      (specize* [s] s)
      (specize* [s _] s)

      s/Spec
      (conform* [_ m]
        (if (keys-pred m)
          (let [reg (s/registry)]
            (loop [ret m, [[k v] & ks :as keys] m]
              (if keys
                (let [sname (keys->specnames k)
                      cspec (conform-override sname)
                      s     (get reg sname)]
                  (cond
                    cspec
                    (if (or (not s) (s/valid? s v))
                      (let [cv (s/conform cspec v)]
                        (if (s/invalid? cv)
                          ::s/invalid
                          (recur (if (identical? cv v) ret (assoc ret k cv))
                            ks)))
                      ::s/invalid)

                    s
                    (let [cv (s/conform s v)]
                      (if (s/invalid? cv)
                        ::s/invalid
                        (recur (if (identical? cv v) ret (assoc ret k cv))
                          ks)))

                    :else
                    (recur ret ks)))
                ret)))
          ::s/invalid))
      (unform* [_ m]
        (let [reg (s/registry)]
          (loop [ret m, [k & ks :as keys] (c/keys m)]
            (if keys
              (let [sname (keys->specnames k)]
                (if (contains? reg sname)
                  (let [cv (get m k)
                        v  (s/unform (or (conform-override sname) sname) cv)]
                    (recur (if (identical? cv v) ret (assoc ret k v))
                      ks))
                  (recur ret ks)))
              ret))))
      (explain* [_ path via in x]
        (if-not (map? x)
          [{:path path :pred 'map? :val x :via via :in in}]
          (let [reg (s/registry)]
            (apply concat
              (when-let [probs (->> (map (fn [pred form] (when-not (pred x) form))
                                      pred-exprs pred-forms)
                                    (keep identity)
                                    seq)]
                (map
                  #(do {:path path :pred % :val x :via via :in in})
                  probs))
              (map
                (fn [[k v]]
                  (let [sname (keys->specnames k)]
                    (if-not (or (not (contains? reg sname))
                              (s-pvalid? sname v k))
                      (s-explain-1 sname sname (conj path k) via (conj in k) v)
                      (if-not (or (not (contains? conform-override sname))
                                (s-pvalid? (conform-override sname) v (conform-override-form sname)))
                        (s-explain-1 (conform-override-form sname) (conform-override sname) (conj path k) via (conj in k) v)))))
                (seq x))))))
      (gen* [_ overrides path rmap]
        (if gfn
          (gfn)
          (let [rmap      (s-inck rmap id)
                gen       (fn [k s]
                            (s-gensub (or (conform-override k) s) overrides
                              (conj path k) rmap (or (conform-override-form k) k)))
                ogen      (fn [k s]
                           (when-not (s-recur-limit? rmap id path k)
                             [k (gen/delay (gen k s))]))
                req-gens  (map gen req-keys req-specs)
                opt-gens  (remove nil? (map ogen opt-keys opt-specs))]
            (when (every? identity (concat req-gens opt-gens))
              (let [reqs (zipmap req-keys req-gens)
                    opts (into {} opt-gens)]
                (gen/bind (gen/choose 0 (count opts))
                  #(let [args (concat (seq reqs) (when (seq opts) (shuffle (seq opts))))]
                     (->> args
                          (take (+ % (count reqs)))
                          (apply concat)
                          (apply gen/hash-map)))))))))
      (with-gen* [_ gfn] (map+-spec-impl (assoc argm :gfn gfn)))
      (describe* [_] (cons `keys
                       (cond-> []
                         req (conj :req req)
                         opt (conj :opt opt)
                         req-un (conj :req-un req-un)
                         opt-un (conj :opt-un opt-un)
                         conform-override-form (conj :reconform conform-override-form)))))))

(defmacro keys+
  "Like s/keys, but accepts an additional :conf map from spec keywords to a
  predicate, conformer, spec, etc. :conf will override the key's spec for
  conforming, but the key's spec will still be checked for validity."
  [& {:keys [req req-un opt opt-un gen conf]}]
  (let [unk          #(-> % name keyword)
        req-keys     (filterv keyword? (flatten req))
        req-un-specs (filterv keyword? (flatten req-un))
        _            (assert (every? #(and (keyword? %) (namespace %)) (concat req-keys req-un-specs opt opt-un))
                       "all keys must be namespace-qualified keywords")
        req-specs    (into req-keys req-un-specs)
        req-keys     (into req-keys (map unk req-un-specs))
        opt-keys     (into (vec opt) (map unk opt-un))
        _            (assert (every? (into (set req-keys) opt-keys) (c/keys conf))
                       "Every key in conf must be metioned by :req, :req-un, :opt, or :opt-un")
        opt-specs    (into (vec opt) opt-un)
        gx           (gensym)
        parse-req    (fn [rk f]
                       (map (fn [x]
                              (if (keyword? x)
                                `(contains? ~gx ~(f x))
                                (walk/postwalk
                                  (fn [y] (if (keyword? y) `(contains? ~gx ~(f y)) y))
                                  x)))
                         rk))
        pred-exprs   [`(map? ~gx)]
        pred-exprs   (into pred-exprs (parse-req req identity))
        pred-exprs   (into pred-exprs (parse-req req-un unk))
        keys-pred    `(fn* [~gx] (and ~@pred-exprs))
        pred-exprs   (mapv (fn [e] `(fn* [~gx] ~e)) pred-exprs)
        pred-forms   (walk/postwalk
                       #?(;; self-hosted cljs
                          :cljs #(s/res &env %)
                          ;; clj targeting cljs or clj
                          :clj  #(#'s/res %))
                       pred-exprs)]
    `(map+-spec-impl {:req                   '~req :opt '~opt
                      :req-un                '~req-un :opt-un '~opt-un
                      :req-keys              '~req-keys :req-specs '~req-specs
                      :opt-keys              '~opt-keys :opt-specs '~opt-specs
                      :pred-forms            '~pred-forms
                      :pred-exprs            ~pred-exprs
                      :keys-pred             ~keys-pred
                      :conform-override      ~conf
                      :conform-override-form '~conf
                      :gfn                   ~gen})))

(s/spec (s/merge :either/either (keys+ :req-un [:either/left :either/right] :conf {:either/left number? :either/right string?}))