feat: add partition take and nthrest functions

7 files changed, 187 insertions, 8 deletions

diff --git a/.dir-locals.el b/.dir-locals.el
index dee5d6a..32da19e 100644
--- a/.dir-locals.el
+++ b/.dir-locals.el
@@ -41,4 +41,7 @@
                  (eval . (put 'if-let 'fennel-indent-function 1))
                  (eval . (put 'fn* 'fennel-indent-function 'defun))
                  (eval . (put 'fn* 'fennel-doc-string-elt 2))
-                 (eval . (put 'defmulti 'fennel-doc-string-elt 2)))))
+                 (eval . (put 'defmulti 'fennel-doc-string-elt 2))
+                 (eval . (put 'try 'fennel-indent-function 0))
+                 (eval . (put 'catch 'fennel-indent-function 1))
+                 (eval . (put 'finally 'fennel-indent-function 0)))))
diff --git a/Makefile b/Makefile
index 25fd154..43682e5 100644
--- a/Makefile
+++ b/Makefile
@@ -8,7 +8,7 @@ LUATESTS = $(FNLTESTS:.fnl=.lua)
 LUA_EXECUTABLES ?= lua luajit
 FENNELDOC := $(shell command -v fenneldoc)
 
-.PHONY: build clean distclean help test luacov luacov-console fenneldoc $(LUA_EXECUTABLES)
+.PHONY: build clean distclean help test luacov luacov-console doc $(LUA_EXECUTABLES)
 
 build: $(LUASOURCES)
 
@@ -48,11 +48,12 @@ luacov-console: luacov
 	luacov-console .
 	@$(foreach test, $(LUATESTS), mv $(test).tmp $(test);)
 
-fenneldoc:
+doc:
 	fenneldoc $(FNLDOCS)
 
 help:
-	@echo "make                -- run tests and create lua library" >&2
+	@echo "make                -- create lua library" >&2
+	@echo "make doc            -- generate documentation files (requires fenneldoc)" >&2
 	@echo "make test           -- run tests" >&2
 	@echo "make clean          -- remove lua files" >&2
 	@echo "make distclean      -- remove all unnecessary files" >&2
diff --git a/doc/cljlib.md b/doc/cljlib.md
index 61765d4..a486bdf 100644
--- a/doc/cljlib.md
+++ b/doc/cljlib.md
@@ -101,6 +101,9 @@ non-ASCII strings.
 - [`not-any?`](#not-any)
 - [`range`](#range)
 - [`reverse`](#reverse)
+- [`take`](#take)
+- [`nthrest`](#nthrest)
+- [`partition`](#partition)
 - [`identity`](#identity)
 - [`comp`](#comp)
 - [`complement`](#complement)
@@ -904,6 +907,73 @@ Function signature:
 
 Returns table with same items as in `tbl` but in reverse order.
 
+## `take`
+Function signature:
+
+```
+(take ([n col]))
+```
+
+Returns a sequence of the first `n` items in `col`, or all items if
+there are fewer than `n`.
+
+## `nthrest`
+Function signature:
+
+```
+(nthrest ([col n]))
+```
+
+Returns the nth rest of `col`, `col` when `n` is 0.
+
+### Examples
+
+``` fennel
+(assert-eq (nthrest [1 2 3 4] 3) [4])
+(assert-eq (nthrest [1 2 3 4] 2) [3 4])
+(assert-eq (nthrest [1 2 3 4] 1) [2 3 4])
+(assert-eq (nthrest [1 2 3 4] 0) [1 2 3 4])
+```
+
+## `partition`
+Function signature:
+
+```
+(partition ([n col]) ([n step col]) ([n step pad col]))
+```
+
+Returns a sequence of sequences of `n` items each, at offsets step
+apart. If `step` is not supplied, defaults to `n`, i.e. the partitions
+do not overlap. If a `pad` collection is supplied, use its elements as
+necessary to complete last partition up to `n` items. In case there
+are not enough padding elements, return a partition with less than `n`
+items.
+
+### Examples
+Partition sequence into sub-sequences of size 3:
+
+``` fennel
+(assert-eq (partition 3 [1 2 3 4 5 6]) [[1 2 3] [4 5 6]])
+```
+
+When collection doesn't have enough elements, partition will not include those:
+
+``` fennel
+(assert-eq (partition 3 [1 2 3 4]) [[1 2 3]])
+```
+
+Partitions can overlap if step is supplied:
+
+``` fennel
+(assert-eq (partition 2 1 [1 2 3 4]) [[1 2] [2 3] [3 4]])
+```
+
+Additional padding can be used to supply insufficient elements:
+
+``` fennel
+(assert-eq (partition 3 3 [3 2 1] [1 2 3 4]) [[1 2 3] [4 3 2]])
+```
+
 ## `identity`
 Function signature:
 
@@ -1188,5 +1258,5 @@ Copyright (C) 2020-2021 Andrey Listopadov
 License: [MIT](https://gitlab.com/andreyorst/fennel-cljlib/-/raw/master/LICENSE)
 
 
-<!-- Generated with Fenneldoc 0.1.2
+<!-- Generated with Fenneldoc v0.1.3
      https://gitlab.com/andreyorst/fenneldoc -->
diff --git a/doc/macros.md b/doc/macros.md
index d834c0a..e347e96 100644
--- a/doc/macros.md
+++ b/doc/macros.md
@@ -594,5 +594,5 @@ Copyright (C) 2020-2021 Andrey Listopadov
 License: [MIT](https://gitlab.com/andreyorst/fennel-cljlib/-/raw/master/LICENSE)
 
 
-<!-- Generated with Fenneldoc 0.1.2
+<!-- Generated with Fenneldoc v0.1.3
      https://gitlab.com/andreyorst/fenneldoc -->
diff --git a/doc/tests/test.md b/doc/tests/test.md
index 2f5f87a..cb97a81 100644
--- a/doc/tests/test.md
+++ b/doc/tests/test.md
@@ -93,5 +93,5 @@ Assert `expr` for not truth. Generates more verbose message if
   `msg` is not set. Works the same as [`assert-is`](#assert-is).
 
 
-<!-- Generated with Fenneldoc 0.1.2
+<!-- Generated with Fenneldoc v0.1.3
      https://gitlab.com/andreyorst/fenneldoc -->
diff --git a/init.fnl b/init.fnl
index c5a3460..5621127 100644
--- a/init.fnl
+++ b/init.fnl
@@ -843,10 +843,87 @@ Basic `zipmap' implementation:
   (when-some [tbl (seq tbl)]
     (reduce consj (empty []) tbl)))
 
+(fn* core.take
+  "Returns a sequence of the first `n' items in `col', or all items if
+there are fewer than `n'."
+  [n col]
+  (if (= n 0)
+      []
+      (pos-int? n)
+      (if-let [s (seq col)]
+        (cons (first s) (take (dec n) (rest s)))
+        nil)
+      (error "expected positive integer as first argument" 2)))
+
+(fn* core.nthrest
+  "Returns the nth rest of `col', `col' when `n' is 0.
+
+# Examples
+
+``` fennel
+(assert-eq (nthrest [1 2 3 4] 3) [4])
+(assert-eq (nthrest [1 2 3 4] 2) [3 4])
+(assert-eq (nthrest [1 2 3 4] 1) [2 3 4])
+(assert-eq (nthrest [1 2 3 4] 0) [1 2 3 4])
+```
+"
+  [col n]
+  [(_unpack col (inc n))])
+
+(fn* core.partition
+  "Returns a sequence of sequences of `n' items each, at offsets step
+apart. If `step' is not supplied, defaults to `n', i.e. the partitions
+do not overlap. If a `pad' collection is supplied, use its elements as
+necessary to complete last partition up to `n' items. In case there
+are not enough padding elements, return a partition with less than `n'
+items.
+
+# Examples
+Partition sequence into sub-sequences of size 3:
+
+``` fennel
+(assert-eq (partition 3 [1 2 3 4 5 6]) [[1 2 3] [4 5 6]])
+```
+
+When collection doesn't have enough elements, partition will not include those:
+
+``` fennel
+(assert-eq (partition 3 [1 2 3 4]) [[1 2 3]])
+```
+
+Partitions can overlap if step is supplied:
+
+``` fennel
+(assert-eq (partition 2 1 [1 2 3 4]) [[1 2] [2 3] [3 4]])
+```
+
+Additional padding can be used to supply insufficient elements:
+
+``` fennel
+(assert-eq (partition 3 3 [3 2 1] [1 2 3 4]) [[1 2 3] [4 3 2]])
+```"
+  ([n col]
+   (partition n n col))
+  ([n step col]
+   (if-let [s (seq col)]
+     (let [p (take n s)]
+       (if (= n (length p))
+           (cons p (partition n step (nthrest s step)))
+           nil))
+     nil))
+  ([n step pad col]
+   (if-let [s (seq col)]
+     (let [p (take n s)]
+       (if (= n (length p))
+           (cons p (partition n step pad (nthrest s step)))
+           [(take n (concat p pad))]))
+     nil)))
+
 (local sequence-doc-order
        [:vector :seq :kvseq :first :rest :last :butlast
         :conj :disj :cons :concat :reduce :reduced :reduce-kv
-        :mapv :filter :every? :some :not-any? :range :reverse])
+        :mapv :filter :every? :some :not-any? :range :reverse :take
+        :nthrest :partition])
 
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; Equality ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
diff --git a/tests/core.fnl b/tests/core.fnl
index f15de7e..429df46 100644
--- a/tests/core.fnl
+++ b/tests/core.fnl
@@ -209,6 +209,8 @@
   (testing "kvseq"
     (assert-not (pcall kvseq))
     (assert-not (pcall kvseq [] []))
+    (assert-eq (kvseq nil) nil)
+    (assert-eq (kvseq []) nil)
     (assert-eq (kvseq {123 456}) [[123 456]])
     (assert-eq (kvseq {:a 1}) [[:a 1]])
     (assert-eq (kvseq [0 0 0 10]) [[1 0] [2 0] [3 0] [4 10]])
@@ -241,6 +243,32 @@
                 "Alice Watson works as chief officer at Coffee With You"])
     (assert-eq (table.concat (mapv string.upper "vaiv")) "VAIV"))
 
+  (testing "partition"
+    (assert-not (pcall partition))
+    (assert-not (pcall partition 1))
+    (assert-eq (partition 1 [1 2 3 4]) [[1] [2] [3] [4]])
+    (assert-eq (partition 1 2 [1 2 3 4]) [[1] [3]])
+    (assert-eq (partition 3 2 [1 2 3 4 5]) [[1 2 3] [3 4 5]])
+    (assert-eq (partition 3 3 [0 -1 -2 -3] [1 2 3 4]) [[1 2 3] [4 0 -1]]))
+
+  (testing "nthrest"
+    (assert-not (pcall nthrest))
+    (assert-not (pcall nthrest []))
+    (assert-eq (nthrest [1 2 3] 0) [1 2 3])
+    (assert-eq (nthrest [1 2 3] 1) [2 3])
+    (assert-eq (nthrest [1 2 3] 2) [3])
+    (assert-eq (nthrest [1 2 3] 3) []))
+
+  (testing "take"
+    (assert-not (pcall take))
+    (assert-not (pcall take []))
+    (assert-not (pcall take :a []))
+    (assert-not (pcall take -1 []))
+    (assert-eq (take 0 [1 2 3]) [])
+    (assert-eq (take 1 {:a 1}) [[:a 1]])
+    (assert-eq (take 10 [1 2 3]) [1 2 3])
+    (assert-eq (take 1 [1 2 3]) [1]))
+
   (testing "reduce"
     (fn* add
       ([] 0)