diff --git a/pkgs/racket-doc/scribblings/reference/sequences.scrbl b/pkgs/racket-doc/scribblings/reference/sequences.scrbl index 2db99dadaf..9f88fd2ec5 100644 --- a/pkgs/racket-doc/scribblings/reference/sequences.scrbl +++ b/pkgs/racket-doc/scribblings/reference/sequences.scrbl @@ -1122,6 +1122,11 @@ stream, but plain lists can be used as streams, and functions such as the resulting stream. } +@defproc[(stream-take [s stream?] [i exact-nonnegative-integer?]) + stream?]{ + Returns a list of the first @racket[i] elements of @racket[s]. +} + @defproc[(stream-append [s stream?] ...) stream?]{ Returns a stream that contains all elements of each stream in the diff --git a/pkgs/racket-test-core/tests/racket/stream.rktl b/pkgs/racket-test-core/tests/racket/stream.rktl index e8311d02c2..525d3fc35f 100644 --- a/pkgs/racket-test-core/tests/racket/stream.rktl +++ b/pkgs/racket-test-core/tests/racket/stream.rktl @@ -47,6 +47,7 @@ (test 3 stream-length '(1 2 3)) (test 1 stream-ref '(1 2 3) 0) (test '(2 3) stream-tail '(1 2 3) 1) +(test '(1 2) stream->list (stream-take '(1 2 3) 2)) (test '(1 2 3 4 5) stream->list (stream-append '(1 2 3) '(4 5))) (test '(1 2 3) stream->list (stream-map values '(1 2 3))) (test #f stream-andmap even? '(1 2 3)) @@ -67,4 +68,18 @@ (test '(0 1 2 3 4 5) stream->list (for/stream ([i (in-naturals)]) (define ii (sqr i)) #:break (> ii 30) i)) +;; stream-take works on infinite streams with lazy-delayed errors later +(test '(1 4/3 4/2 4/1) stream->list + (stream-take (let loop ([i 4]) + (stream-cons (/ 4 i) (loop (sub1 i)))) + 4)) +;; stream-take preserves laziness, doesn't evaluate elements too early +(define (guarded-second s) + (if (stream-ref s 0) (stream-ref s 1) #f)) +(define (div a b) + (stream (not (zero? b)) (/ a b))) +(test #f guarded-second (stream-take (div 1 0) 2)) +(test 3/4 guarded-second (stream-take (div 3 4) 2)) +(err/rt-test (stream->list (stream-take (stream 1 2) 3)) exn:fail:contract? "stream-take") + (report-errs) diff --git a/racket/collects/racket/stream.rkt b/racket/collects/racket/stream.rkt index ea770509c9..a176f38e4a 100644 --- a/racket/collects/racket/stream.rkt +++ b/racket/collects/racket/stream.rkt @@ -33,6 +33,7 @@ stream-length stream-ref stream-tail + stream-take stream-append stream-map stream-andmap @@ -115,6 +116,23 @@ [else (loop (sub1 n) (stream-rest s))]))) + + +(define (stream-take st i) + (unless (stream? st) (raise-argument-error 'stream-take "stream?" st)) + (unless (exact-nonnegative-integer? i) + (raise-argument-error 'stream-take "exact-nonnegative-integer?" i)) + (let loop ([n i] [s st]) + (cond + [(zero? n) empty-stream] + [(stream-empty? s) + (raise-arguments-error 'stream-take + "stream ended before index" + "index" i + "stream" st)] + [else + (stream* (stream-first s) (loop (sub1 n) (stream-rest s)))]))) + (define (stream-append . l) (for ([s (in-list l)]) (unless (stream? s) (raise-argument-error 'stream-append "stream?" s)))