add stream-take + docs (#2181)

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

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)

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)))