Merge pull request #98 from tomdl89/window

Add window function

Author: weavejester

src/medley/core.cljc

@@ -506,6 +506,35 @@
   ([pred coll]
    (partition-between (fn [_ x] (pred x)) coll)))
 
+(defn window
+  "A sliding window, returning a lazy sequence of partitions, containing each
+  element and n-1 preceeding elements, when present. Therefore partitions at the
+  start may contain fewer items than the rest. Returns a stateful transducer
+  when no collection is provided. For a sliding window containing each element
+  and n-1 _following_ elements, use `clojure.core/partition` with a `step` size
+  of 1."
+  {:added "1.9.0"}
+  ([n]
+   (fn [rf]
+     (let [part #?(:clj (java.util.ArrayList. n) :cljs (array))]
+       (fn
+         ([] (rf))
+         ([result] (rf result))
+         ([result x]
+          #?(:clj (.add part x) :cljs (.push part x))
+          (when (< n #?(:clj (.size part) :cljs (.-length part)))
+            #?(:clj (.remove part 0) :cljs (.shift part)))
+          (rf result (vec #?(:clj (.toArray part) :cljs (.slice part)))))))))
+  ([n coll]
+   (letfn [(part [part-n coll]
+             (let [run (doall (take part-n coll))]
+               (lazy-seq
+                (when (== part-n (count run))
+                  (cons run
+                        (part (min n (inc part-n))
+                              (if (== n part-n) (rest coll) coll)))))))]
+     (part (min 1 n) coll))))
+
 (defn indexed
   "Returns an ordered, lazy sequence of vectors `[index item]`, where item is a
   value in coll, and index its position starting from zero. Returns a stateful

test/medley/core_test.cljc

@@ -398,6 +398,28 @@
                       [1 2 3 2 8 9])
            [[1] [2] [3 2] [8]]))))
 
+(deftest test-window
+  (testing "sequences"
+    (is (= (m/window 2 nil) '()))
+    (is (= (m/window 2 [:a]) '((:a))))
+    (is (= (m/window 2 [:a :b]) '((:a) (:a :b))))
+    (is (= (m/window 2 [:a :b :c]) '((:a) (:a :b) (:b :c))))
+    (is (= (take 3 (m/window 0 [:a :b :c])) '(() () ())))
+    (is (= (m/window 10 [:a :b :c]) '((:a) (:a :b) (:a :b :c)))))
+  (testing "transducers"
+    (is (= (transduce (m/window 2) conj nil) []))
+    (is (= (transduce (m/window 2) conj [:a]) [[:a]]))
+    (is (= (transduce (m/window 2) conj [:a :b]) [[:a] [:a :b]]))
+    (is (= (transduce (m/window 2) conj [:a :b :c]) [[:a] [:a :b] [:b :c]]))
+    (is (= (transduce (m/window 0) conj [:a :b :c]) [[] [] []]))
+    (is (= (transduce (m/window 10) conj [:a :b :c]) [[:a] [:a :b] [:a :b :c]]))
+    (is (= (transduce (m/window 2)
+                      (completing (fn [coll x]
+                                    (cond-> (conj coll x) (= [:a :b] x) reduced)))
+                      []
+                      [:a :b :c :d])
+           [[:a] [:a :b]]))))
+
 (deftest test-indexed
   (testing "sequences"
     (is (= (m/indexed [:a :b :c :d])