mirrors/janet
1
0
Fork 0
mirror of https://github.com/janet-lang/janet synced 2024-12-26 00:10:27 +00:00
Code Issues 1 Releases Wiki Activity

Merge pull request #756 from subsetpark/take-streams

Browse Source 
Make take/take-while/take-until fiber-aware
This commit is contained in:
Calvin Rose 2021-08-06 15:13:35 -05:00 committed by GitHub
commit 3a787afec6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 50 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
src/boot/boot.janet
View File

@ -1031,30 +1031,60 @@
(set k (next ind k)))
ret)
(defn take
"Take the first n elements of an indexed or bytes type. Returns a new tuple or string, respectively."
[n ind]
(def use-str (bytes? ind))
(def f (if use-str string/slice tuple/slice))
(defn- resume-n
[n fib]
(def res @[])
(var taken 0)
(while (and (< taken n) (fiber/can-resume? fib))
(let [elem (resume fib)]
(+= taken 1)
(array/push res elem)))
res)
(defn- slice-n
[f n ind]
(def len (length ind))
# make sure end is in [0, len]
(def m (if (> n 0) n 0))
(def end (if (> m len) len m))
(f ind 0 end))
(defn take-until
"Same as `(take-while (complement pred) ind)`."
[pred ind]
(def use-str (bytes? ind))
(def f (if use-str string/slice tuple/slice))
(defn take
"Take the first n elements of a fiber, indexed or bytes type. Returns a new array, tuple or string, respectively."
[n ind]
(cond
(fiber? ind) (resume-n n ind)
(bytes? ind) (slice-n string/slice n ind)
(slice-n tuple/slice n ind)))
(defn- resume-until
[pred fib]
(def res @[])
(while (fiber/can-resume? fib)
(let [elem (resume fib)]
(if (pred elem)
(break)
(array/push res elem))))
res)
(defn- slice-until
[f pred ind]
(def len (length ind))
(def i (find-index pred ind))
(def end (if (nil? i) len i))
(f ind 0 end))
(defn take-until
"Same as `(take-while (complement pred) ind)`."
[pred ind]
(cond
(fiber? ind) (resume-until pred ind)
(bytes? ind) (slice-until string/slice pred ind)
(slice-until tuple/slice pred ind)))
(defn take-while
`Given a predicate, take only elements from an indexed or bytes type that satisfy
the predicate, and abort on first failure. Returns a new tuple or string, respectively.`
`Given a predicate, take only elements from a fiber, indexed or bytes type that satisfy
the predicate, and abort on first failure. Returns a new array, tuple or string, respectively.`
[pred ind]
(take-until (complement pred) ind))

8
test/suite0005.janet
View File

@ -51,6 +51,10 @@
(assert (deep= (take 0 [1 2 3 4 5]) []) "take 3")
(assert (deep= (take 10 [1 2 3]) [1 2 3]) "take 4")
(assert (deep= (take -1 [:a :b :c]) []) "take 5")
(assert (deep= (take 3 (generate [x :in [1 2 3 4 5]] x)) @[1 2 3]) "take from fiber")
# NB: repeatedly resuming a fiber created with `generate` includes a `nil` as
# the final element. Thus a generate of 2 elements will create an array of 3.
(assert (= (length (take 4 (generate [x :in [1 2]] x))) 3) "take from short fiber")
(assert-error :invalid-type (take 3 {}) "take 6")
# take-until
@ -61,6 +65,8 @@
(assert (deep= (take-until pos? @[-1 -2 3]) [-1 -2]) "take-until 4")
(assert (deep= (take-until pos? @[-1 1 -2]) [-1]) "take-until 5")
(assert (deep= (take-until |(= $ 115) "books") "book") "take-until 6")
(assert (deep= (take-until |(= $ 115) (generate [x :in "books"] x))
@[98 111 111 107]) "take-until from fiber")
# take-while
@ -69,6 +75,8 @@
(assert (deep= (take-while neg? @[-1 -2 -3]) [-1 -2 -3]) "take-while 3")
(assert (deep= (take-while neg? @[-1 -2 3]) [-1 -2]) "take-while 4")
(assert (deep= (take-while neg? @[-1 1 -2]) [-1]) "take-while 5")
(assert (deep= (take-while neg? (generate [x :in @[-1 1 -2]] x))
@[-1]) "take-while from fiber")
# drop