async_utils/
traits.rs

1// Copyright 2020 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Traits that are useful for working with async code, but do not fit into a more specific
6//! category. These are often extension traits for types that are not defined by `async_utils`.
7
8use core::task::Poll;
9
10/// An extension trait for `core::task::Poll` that provides convenient adapters for extracting
11/// output values.
12pub trait PollExt<T> {
13    /// Returns the value contained in a Poll::Ready value or panics.
14    fn unwrap(self) -> T;
15
16    /// Returns the value contained in a Poll::Ready value or panics with a custom message.
17    fn expect(self, msg: &str) -> T;
18
19    /// Expect a Poll to be Pending, or panics with a custom message.
20    fn expect_pending(&self, msg: &str);
21
22    /// Turns a Poll into a Result, mapping Poll::Ready(value) to Ok(value) and
23    /// Poll::Pending to Err(error)
24    fn ready_or<E>(self, error: E) -> Result<T, E>;
25
26    /// Turns a Poll into a Result, mapping Poll::Ready(value) to Ok(value) and
27    /// Poll::Pending to Err(error())
28    fn ready_or_else<E, F>(self, error: F) -> Result<T, E>
29    where
30        F: FnOnce() -> E;
31}
32
33impl<T> PollExt<T> for Poll<T> {
34    #[inline]
35    #[track_caller]
36    fn unwrap(self) -> T {
37        match self {
38            Poll::Ready(val) => val,
39            Poll::Pending => panic!("called `Poll::unwrap()` on a `Pending` value"),
40        }
41    }
42
43    #[inline]
44    #[track_caller]
45    fn expect(self, msg: &str) -> T {
46        match self {
47            Poll::Ready(val) => val,
48            Poll::Pending => panic!("{}", msg),
49        }
50    }
51
52    #[inline]
53    #[track_caller]
54    fn expect_pending(&self, msg: &str) {
55        if self.is_ready() {
56            panic!("{}", msg);
57        }
58    }
59
60    fn ready_or<E>(self, error: E) -> Result<T, E> {
61        self.ready_or_else(|| error)
62    }
63
64    fn ready_or_else<E, F>(self, error: F) -> Result<T, E>
65    where
66        F: FnOnce() -> E,
67    {
68        match self {
69            Poll::Ready(val) => Ok(val),
70            Poll::Pending => Err(error()),
71        }
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn poll_unwrap_ready_returns_value() {
81        let p = Poll::Ready("value");
82        assert_eq!(p.unwrap(), "value");
83    }
84
85    #[test]
86    #[should_panic]
87    fn poll_unwrap_pending_panics() {
88        let p: Poll<()> = Poll::Pending;
89        p.unwrap();
90    }
91
92    #[test]
93    fn poll_expect_ready_returns_value() {
94        let p = Poll::Ready("value");
95        assert_eq!(p.expect("missing value"), "value");
96    }
97
98    #[test]
99    fn poll_expect_pending() {
100        let p: Poll<()> = Poll::Pending;
101        p.expect_pending("is pending");
102    }
103
104    #[test]
105    #[should_panic]
106    fn poll_expect_pending_on_ready() {
107        let p = Poll::Ready("value");
108        p.expect_pending("value is not pending");
109    }
110
111    #[test]
112    #[should_panic(expected = "missing value")]
113    fn poll_expect_pending_panics_with_message() {
114        let p: Poll<()> = Poll::Pending;
115        p.expect("missing value");
116    }
117
118    #[test]
119    fn poll_ready_or_ready_returns_ok() {
120        let p = Poll::Ready("value");
121        assert_eq!(p.ready_or(()), Ok("value"));
122    }
123
124    #[test]
125    fn poll_ready_or_pending_returns_error() {
126        let p: Poll<()> = Poll::Pending;
127        assert_eq!(p.ready_or(()), Err(()));
128    }
129
130    #[test]
131    fn poll_ready_or_else_ready_returns_ok() {
132        let p = Poll::Ready("value");
133        assert_eq!(p.ready_or_else(|| ()), Ok("value"));
134    }
135
136    #[test]
137    fn poll_ready_or_else_pending_returns_error() {
138        let p: Poll<()> = Poll::Pending;
139        assert_eq!(p.ready_or_else(|| ()), Err(()));
140    }
141}