1 use {Async, Poll};
2 use stream::Stream;
3
4 /// A stream combinator which returns a maximum number of elements.
5 ///
6 /// This structure is produced by the `Stream::take` method.
7 #[derive(Debug)]
8 #[must_use = "streams do nothing unless polled"]
9 pub struct Take<S> {
10 stream: S,
11 remaining: u64,
12 }
13
new<S>(s: S, amt: u64) -> Take<S> where S: Stream,14 pub fn new<S>(s: S, amt: u64) -> Take<S>
15 where S: Stream,
16 {
17 Take {
18 stream: s,
19 remaining: amt,
20 }
21 }
22
23 impl<S> Take<S> {
24 /// Acquires a reference to the underlying stream that this combinator is
25 /// pulling from.
get_ref(&self) -> &S26 pub fn get_ref(&self) -> &S {
27 &self.stream
28 }
29
30 /// Acquires a mutable reference to the underlying stream that this
31 /// combinator is pulling from.
32 ///
33 /// Note that care must be taken to avoid tampering with the state of the
34 /// stream which may otherwise confuse this combinator.
get_mut(&mut self) -> &mut S35 pub fn get_mut(&mut self) -> &mut S {
36 &mut self.stream
37 }
38
39 /// Consumes this combinator, returning the underlying stream.
40 ///
41 /// Note that this may discard intermediate state of this combinator, so
42 /// care should be taken to avoid losing resources when this is called.
into_inner(self) -> S43 pub fn into_inner(self) -> S {
44 self.stream
45 }
46 }
47
48 // Forwarding impl of Sink from the underlying stream
49 impl<S> ::sink::Sink for Take<S>
50 where S: ::sink::Sink + Stream
51 {
52 type SinkItem = S::SinkItem;
53 type SinkError = S::SinkError;
54
start_send(&mut self, item: S::SinkItem) -> ::StartSend<S::SinkItem, S::SinkError>55 fn start_send(&mut self, item: S::SinkItem) -> ::StartSend<S::SinkItem, S::SinkError> {
56 self.stream.start_send(item)
57 }
58
poll_complete(&mut self) -> Poll<(), S::SinkError>59 fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
60 self.stream.poll_complete()
61 }
62
close(&mut self) -> Poll<(), S::SinkError>63 fn close(&mut self) -> Poll<(), S::SinkError> {
64 self.stream.close()
65 }
66 }
67
68 impl<S> Stream for Take<S>
69 where S: Stream,
70 {
71 type Item = S::Item;
72 type Error = S::Error;
73
poll(&mut self) -> Poll<Option<S::Item>, S::Error>74 fn poll(&mut self) -> Poll<Option<S::Item>, S::Error> {
75 if self.remaining == 0 {
76 Ok(Async::Ready(None))
77 } else {
78 let next = try_ready!(self.stream.poll());
79 match next {
80 Some(_) => self.remaining -= 1,
81 None => self.remaining = 0,
82 }
83 Ok(Async::Ready(next))
84 }
85 }
86 }
87