GREsau / okapi

OpenAPI (AKA Swagger) document generation for Rust projects
MIT License
578 stars 103 forks source link

How to specify image response content-type? #139

Closed ShayBox closed 6 months ago

ShayBox commented 6 months ago

I'm switching from Poem to Rocket and I need to respond with an image/png content-type

Poem:

use std::io::Cursor;

use color_processing::Color;
use image::{ImageBuffer, ImageFormat, Rgba};
use poem::{error::InternalServerError, http::StatusCode, Error, Result};
use poem_openapi::{param::Path, payload::Binary, ApiResponse, OpenApi};

#[derive(ApiResponse)]
enum ImageResponse {
    #[oai(content_type = "image/png", status = 200)]
    Png(Binary<Vec<u8>>),
}

pub struct Hex;

#[OpenApi]
impl Hex {
    /// Generate a solid color image
    #[allow(clippy::unused_async)]
    #[oai(path = "/hex/generate/:color/:height/:width", method = "get")]
    async fn hex(
        &self,
        color: Path<String>,
        height: Path<u32>,
        width: Path<u32>,
    ) -> Result<ImageResponse> {
        if height.0 > 10_000 || width.0 > 10_000 {
            let error = Error::from_status(StatusCode::BAD_REQUEST);
            return Err(error);
        }

        let Ok(color) = color.0.parse::<Color>() else {
            let error = Error::from_status(StatusCode::BAD_REQUEST);
            return Err(error);
        };

        let rgba = Rgba([color.red, color.green, color.blue, color.alpha]);
        let image = ImageBuffer::from_pixel(width.0, height.0, rgba);
        let mut bytes = Vec::new();
        image
            .write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
            .map_err(InternalServerError)?;

        Ok(ImageResponse::Png(Binary(bytes)))
    }
}

Rocket:

use std::io::Cursor;

use color_processing::Color;
use image::{ImageBuffer, ImageFormat, Rgba};
use rocket::response::status::BadRequest;
use rocket_okapi::openapi;

use crate::host::API;

const MAX_SIZE: u32 = 10_000;

#[derive(Responder)]
#[response(status = 200, content_type = "image/png")]
struct Png(Vec<u8>);

#[openapi]
#[get("/hex/<color>/<height>/<width>")]
pub fn hex(_host: API, color: &str, height: u32, width: u32) -> Result<Png, BadRequest<String>> {
    if height > MAX_SIZE || width > MAX_SIZE {
        let size = height.max(width);
        let err = format!("{size} is larger than the maximum size ({MAX_SIZE})");
        return Err(BadRequest(err));
    }

    let color = color.parse::<Color>().map_err(BadRequest)?;
    let rgba = Rgba([color.red, color.green, color.blue, color.alpha]);
    let image = ImageBuffer::from_pixel(width, height, rgba);
    let mut bytes = Vec::new();
    image
        .write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
        .map_err(|error| {
            eprintln!("   >> Error: {error}");
            BadRequest(error.to_string())
        })?;

    Ok(Png(bytes))
}

EDIT: I found the rocket responder, but it doesn't implement OpenApiResponder/Inner.

ShayBox commented 6 months ago

I figured out how to do it, by searching GitHub I found Revolt has this code

use std::io::Cursor;

use color_processing::Color;
use image::{ImageBuffer, ImageFormat, Rgba};
use rocket::response::status::BadRequest;
use rocket_okapi::{
    gen::OpenApiGenerator,
    okapi::openapi3::Responses,
    openapi,
    response::OpenApiResponderInner,
    OpenApiError,
};

const MAX_SIZE: u32 = 2048;

#[derive(Responder)]
#[response(status = 200, content_type = "image/png")]
pub struct Image(Vec<u8>);

impl OpenApiResponderInner for Image {
    fn responses(_gen: &mut OpenApiGenerator) -> Result<Responses, OpenApiError> {
        use rocket_okapi::okapi::{
            openapi3::{MediaType, RefOr, Response, SchemaObject},
            schemars::{
                schema::{InstanceType, SingleOrVec},
                Map,
            },
        };

        let content = Map::from([(
            String::from("image/png"),
            MediaType {
                schema: Some(SchemaObject {
                    instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
                    format: Some(String::from("binary")),
                    ..Default::default()
                }),
                ..Default::default()
            },
        )]);

        let responses = Map::from([(
            String::from("200"),
            RefOr::Object(Response {
                description: String::from("Image"),
                content,
                ..Default::default()
            }),
        )]);

        Ok(Responses {
            responses,
            ..Default::default()
        })
    }
}

/// Generate a solid color image
///
/// # Errors
///
/// Will return `Err` if `height`/`width` are too big or `&str::parse`/`ImageBuffer::write_to` errors.
#[openapi(tag = "Misc")]
#[get("/hex/<color>/<height>/<width>")]
pub fn hex(color: &str, height: u32, width: u32) -> Result<Image, BadRequest<String>> {
    if height > MAX_SIZE || width > MAX_SIZE {
        let size = height.max(width);
        let err = format!("{size} is larger than the maximum size ({MAX_SIZE})");
        return Err(BadRequest(err));
    }

    let color = color.parse::<Color>().map_err(BadRequest)?;
    let rgba = Rgba([color.red, color.green, color.blue, color.alpha]);
    let image = ImageBuffer::from_pixel(width, height, rgba);
    let mut bytes = Vec::new();
    image
        .write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
        .map_err(|error| {
            eprintln!("   >> Error: {error}");
            BadRequest(error.to_string())
        })?;

    Ok(Image(bytes))
}