poem_openapi/payload/
response.rs

1use poem::{
2    http::{header::HeaderName, HeaderMap, HeaderValue, StatusCode},
3    Error, IntoResponse,
4};
5
6use crate::{
7    registry::{MetaResponses, Registry},
8    ApiResponse,
9};
10
11/// A response type wrapper.
12///
13/// Use it to modify the status code and HTTP headers.
14///
15/// # Examples
16///
17/// ```
18/// use poem::{
19///     error::BadRequest,
20///     http::{Method, StatusCode, Uri},
21///     test::TestClient,
22///     Body, IntoEndpoint, Request, Result,
23/// };
24/// use poem_openapi::{
25///     payload::{Json, Response},
26///     OpenApi, OpenApiService,
27/// };
28/// use tokio::io::AsyncReadExt;
29///
30/// struct MyApi;
31///
32/// #[OpenApi]
33/// impl MyApi {
34///     #[oai(path = "/test", method = "get")]
35///     async fn test(&self) -> Response<Json<i32>> {
36///         Response::new(Json(100)).header("foo", "bar")
37///     }
38/// }
39///
40/// let api = OpenApiService::new(MyApi, "Demo", "0.1.0");
41///
42/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
43/// let resp = TestClient::new(api).get("/test").send().await;
44/// resp.assert_status_is_ok();
45/// resp.assert_header("foo", "bar");
46/// resp.assert_text("100").await;
47/// # });
48/// ```
49pub struct Response<T> {
50    inner: T,
51    status: Option<StatusCode>,
52    headers: HeaderMap,
53}
54
55impl<T> Response<T> {
56    /// Create a response object.
57    #[must_use]
58    pub fn new(resp: T) -> Self {
59        Self {
60            inner: resp,
61            status: None,
62            headers: HeaderMap::new(),
63        }
64    }
65
66    /// Sets the HTTP status for this response.
67    #[must_use]
68    pub fn status(self, status: StatusCode) -> Self {
69        Self {
70            status: Some(status),
71            ..self
72        }
73    }
74
75    /// Appends a header to this response.
76    #[must_use]
77    pub fn header<K, V>(mut self, key: K, value: V) -> Self
78    where
79        K: TryInto<HeaderName>,
80        V: TryInto<HeaderValue>,
81    {
82        let key = key.try_into();
83        let value = value.try_into();
84        if let (Ok(key), Ok(value)) = (key, value) {
85            self.headers.append(key, value);
86        }
87        self
88    }
89}
90
91impl<T: IntoResponse> IntoResponse for Response<T> {
92    fn into_response(self) -> poem::Response {
93        let mut resp = self.inner.into_response();
94        if let Some(status) = self.status {
95            resp.set_status(status);
96        }
97        resp.headers_mut().extend(self.headers);
98        resp
99    }
100}
101
102impl<T: ApiResponse> ApiResponse for Response<T> {
103    const BAD_REQUEST_HANDLER: bool = T::BAD_REQUEST_HANDLER;
104
105    fn meta() -> MetaResponses {
106        T::meta()
107    }
108
109    fn register(registry: &mut Registry) {
110        T::register(registry);
111    }
112
113    fn from_parse_request_error(err: Error) -> Self {
114        Self::new(T::from_parse_request_error(err))
115    }
116}