enhance partition

docs/language/core.rst

@@ -671,21 +671,37 @@ Returns ``True`` if *x* is odd. Raises ``TypeError`` if
    => (odd? 0)
    False
 
-
 .. _partition-fn:
 
 partition
 ---------
 
-Usage: ``(partition n coll)``
+Usage: ``(partition coll [n] [step] [fillvalue])``
 
-Chunks coll into tuples of length *n*. The remainder, if any, is not included.
+Chunks *coll* into *n*-tuples (pairs by default).
 
 .. code-block:: hy
 
-   => (list (partition 3 (range 10)))
+   => (list (partition (range 10)))  ; n=2
-   [(0, 1, 2), (3, 4, 5), (6, 7, 8)]
+   [(, 0 1) (, 2 3) (, 4 5) (, 6 7) (, 8 9)]
 
+The *step* defaults to *n*, but can be more to skip elements, or less for a sliding window with overlap.
+
+.. code-block:: hy
+
+   => (list (partition (range 10) 2 3))
+   [(, 0 1) (, 3 4) (, 6 7)]
+   => (list (partition (range 5) 2 1))
+   [(, 0 1) (, 1 2) (, 2 3) (, 3 4)])
+
+The remainder, if any, is not included unless a *fillvalue* is specified.
+
+.. code-block:: hy
+
+   => (list (partition (range 10) 3))
+   [(, 0 1 2) (, 3 4 5) (, 6 7 8)]
+   => (list (partition (range 10) 3 :fillvalue "x"))
+   [(, 0 1 2) (, 3 4 5) (, 6 7 8) (, 9 "x" "x")]
 
 .. _pos?-fn:
 

hy/core/language.hy

@@ -314,9 +314,17 @@
   (_numeric-check n)
   (= (% n 2) 1))
 
-(defn partition [n coll]
+(def -sentinel (object))
-  "Chunks coll into tuples of length n. The remainder, if any, is not included."
+(defn partition [coll &optional [n 2] step [fillvalue -sentinel]]
-  (apply zip (* n (, (iter coll)))))
+  "Chunks coll into n-tuples (pairs by default). The remainder, if any, is not
+   included unless a fillvalue is specified. The step defaults to n, but can be
+   more to skip elements, or less for a sliding window with overlap."
+  (setv
+   step (or step n)
+   slices (genexpr (itertools.islice coll start nil step) [start (range n)]))
+  (if (is fillvalue -sentinel)
+    (apply zip slices)
+    (apply zip-longest slices {"fillvalue" fillvalue})))
 
 (defn pos? [n]
   "Return true if n is > 0"

tests/native_tests/core.hy

@@ -473,14 +473,27 @@
 (defn test-partition []
   "NATIVE: testing the partition function"
   (setv ten (range 10))
-  (assert-equal (list (partition 3 ten))
+  ;; no remainder
+  (assert-equal (list (partition ten 3))
                 [(, 0 1 2) (, 3 4 5) (, 6 7 8)])
-  (assert-equal (list (partition 2 ten))
+  ;; pair by default
+  (assert-equal (list (partition ten))
                 [(, 0 1) (, 2 3) (, 4 5) (, 6 7) (, 8 9)])
-  (assert-equal (list (partition 1 ten))
+  ;; length 1 is valid
+  (assert-equal (list (partition ten 1))
                 [(, 0) (, 1) (, 2) (, 3) (, 4) (, 5) (, 6) (, 7) (, 8) (, 9)])
-  (assert-equal (list (partition 0 ten)) [])
+  ;; tuples of length < 1 don't crash
-  (assert-equal (list (partition -1 ten)) []))
+  (assert-equal (list (partition ten 0)) [])
+  (assert-equal (list (partition ten -1)) [])
+  ;; keep remainder with a fillvalue
+  (assert-equal (list (partition ten 3 :fillvalue "x"))
+                [(, 0 1 2) (, 3 4 5) (, 6 7 8) (, 9 "x" "x")])
+  ;; skip elements with step > n
+  (assert-equal (list (partition ten 2 3))
+                [(, 0 1) (, 3 4) (, 6 7)])
+  ;; overlap with step < n
+  (assert-equal (list (partition (range 5) 2 1))
+                [(, 0 1) (, 1 2) (, 2 3) (, 3 4)]))
 
 (defn test-pos []
   "NATIVE: testing the pos? function"