Datomic-ish pull for vanilla Clojure data structures

datapull.clj

(ns datapull
  "Datomic-ish pull for vanilla Clojure data structures. Usage examples in test below."
  (:require [clojure.test :as t]))

;; Alternatives:
;; - Juxt has a full-featured library for doing the same thing: https://github.com/juxt/pull
;; - Meander and Specter are both libraries for performing complex searches & transformations on Clojure data

(defn- seq-into
  "Like into, but returns nil if existing-seq is empty or contains only nils."
  [empty-coll existing-seq]
  (let [nil-free-existing-seq (remove nil? existing-seq)]
    (when (seq nil-free-existing-seq)
     (into empty-coll nil-free-existing-seq))))

(defn dpull*
  "Inner implementation of dpull"
  [d p]
  (when (and d p)
    (cond
      (sequential? p) ; seq containing keys for a collection
      (do
        (assert (or (map? d) (set? d) (sequential? d)) (str "Pull vector " p " can't be applied to a " (type d) ", only to a map or sequence."))
        (apply merge
               (seq-into []
                (map (partial dpull* d) p))))

      (map? p) ; specification for a key and its substructure
      (do
        (assert (= 1 (count p)) (str "Pull syntax map " p " can only contain one item."))
        (let [[k v] (first p)]
          (assert (vector? v) (str "Value in pull syntax map " p " must be a vector; " v " is a " (type v) "."))
          ;; k tells us what substructure to get from d; v describes what should
          ;; be pulled from that substructure.
          (when (contains? d k)
            (when-let [v-pull (dpull* (get d k) v)]
              {k v-pull}))))

      :else ; treat this element as a key which must be present in the current data structure
      (do
        (assert (or (map? d) (set? d) (sequential? d)) (str "Pull vector " p " can't be applied to a " (type d) ", only to a map or sequence."))
        (if (or (sequential? d) (set? d)) ; seq-ish of entities; apply pull to each
          (seq-into (empty d)
                    (remove nil?
                            (mapv #(dpull* % p) d)))
          (when-let [v (get d p)] {p v})))))) ; simple map retrieval

(defn dpull
  "Pull data from a nested data structure d, using Datomic pull structure p.
  Matches Datomic pull behavior in most ways, but it's unnecessary to use '* for
  'all attributes'; simply including :foo/bar without further specification is
  sufficient to return all substructure of :foo/bar."
  [d p]
  (assert (vector? p) "Top-level structure of pull specification must be a vector.")
  (dpull* d p))

(t/deftest dpull-test
  (let [ex {:user/id 12,
            :user/name "Bob",
            :user/address {:address/town "Springfield",
                           :address/zip "11111",
                           :address/georegion {:georegion/name "Central Valley",
                                               :georegion/ave-temp 80}}
            :user/pets [{:pet/name "Tony",
                         :pet/type "Tiger"},
                        {:pet/name "Louis",
                         :pet/type "Lion"}]
            1247 :excuse-for-non-kw-key
            :user/emails #{{:email/name "Gmail", :email/address "foo@gmail.com"}
                           {:email/name "Yahoo", :email/address "foo@yahoo.com"}}}]
    ;; TODO maybe add nested-vector case?
    (t/testing "simple cases"
      (t/is (= {:user/name "Bob"}
               (dpull ex [:user/name])))
      (t/is (= {:user/name "Bob", :user/id 12}
               (dpull ex [:user/name :user/id])))
      (t/is (= {1247 :excuse-for-non-kw-key}
               (dpull ex [1247]))))
    (t/testing "simple nested maps"
      (t/is (= {:user/name "Bob", :user/address {:address/zip "11111"}}
               (dpull ex [:user/name {:user/address [:address/zip]}])))
      (t/is (= {:user/address {:address/zip "11111"}}
               (dpull ex [{:user/address [:address/zip]}])))
      (t/is (= #:user{:address #:address{:georegion #:georegion{:ave-temp 80}}}
               (dpull ex [{:user/address [{:address/georegion [:georegion/ave-temp]}]}]))))
    (t/testing "underspecification of nested structures returns everything (so no need for '*)"
      (t/is (= #:user{:address
                      #:address{:town "Springfield",
                                :zip "11111",
                                :georegion #:georegion{:name "Central Valley", :ave-temp 80}}}
               (dpull ex [:user/address]))))
    (t/testing "sequences in data structure"
      (t/is (= {:user/pets [{:pet/type "Tiger"} {:pet/type "Lion"}]}
               (dpull ex [{:user/pets [:pet/type]}]))))
    (t/testing "sets"
      (t/is (= {:user/emails #{{:email/name "Gmail"}, {:email/name "Yahoo"}}}
               (dpull ex [{:user/emails [:email/name]}]))))
    (t/testing "missing"
      (t/is (= {:user/id 12}
               (dpull ex [:user/id :user/religion])))
      (t/is (= nil
               (dpull ex [{:user/emails [:email/administrator]}])))
      (t/is (= nil
               (dpull ex [{:user/pets [:pet/diseases]}]))))
    (t/testing "nil data structure (pull structure isn't allowed to be nil):"
      (t/is (= nil (dpull nil [:foo/bar]))))))