1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

#[deny(missing_docs)]
extern crate base64;
extern crate reqwest;
extern crate serde;
extern crate snafu;

use serde::Serialize;
use snafu::{ResultExt, Snafu};

/// Various types used by trace.moe API.
pub mod types;

pub use types::Doc;

use types::{Me, SearchRequest, SearchResponse};

/// A wrapper around [`reqwest::Client`] to talk to trace.moe API.
pub struct Client {
    /// Defaults to `"https://trace.moe/api"`. Primarily useful for testing.
    pub base_uri: String,
    /// Access token that registered developers have.
    pub token: Option<String>,
    client: reqwest::Client,
}

impl Client {
    /// Create a [`Client`] with default settings.
    pub fn new() -> Self {
        Self {
            base_uri: "https://trace.moe/api".into(),
            client: reqwest::Client::new(),
            token: None,
        }
    }

    /// Create a [`Client`] with an API token.
    pub fn with_token(token: String) -> Self {
        Self {
            token: Some(token),
            ..Self::new()
        }
    }

    /// Search for image.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use moe::Client;
    /// # let client = Client::new();
    /// let image = std::fs::read("image.jpg").unwrap();
    /// let response = client.search(image).unwrap();
    /// assert_eq!("Hataraku Saibou", response.docs[0].title_romaji);
    /// ```
    ///
    /// # Errors
    ///
    /// - HTTP 400 if your search image is empty
    /// - HTTP 403 if are using an invalid token
    /// - HTTP 413 if the image is >1MB
    /// - HTTP 429 if are using requesting too fast
    /// - HTTP 500 or HTTP 503 if something went wrong in backend
    pub fn search(&self, image: Vec<u8>) -> Result<SearchResponse> {
        let body = SearchRequest::new(image);
        let mut response = self.request(reqwest::Method::POST, "search", &body)?;
        match response.status().as_u16() {
            400 => Err(Error::ImageEmpty),
            403 => Err(Error::InvalidToken),
            413 => Err(Error::ImageTooLarge),
            429 => Err(Error::RateLimit {
                message: response.text().context(ResponseEmpty)?,
            }),
            500 | 503 => Err(Error::InternalServerError),
            _ => response.json().context(JsonFailed),
        }
    }

    /// Check the search quota and limit for your account (or IP address).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use moe::Client;
    /// # let client = Client::new();
    /// let me = client.me().unwrap();
    /// assert_eq!("176.38.191.44", me.email);
    /// ```
    pub fn me(&self) -> Result<Me> {
        self.request(reqwest::Method::GET, "me", "")?
            .json()
            .context(JsonFailed)
    }

    fn request<T>(&self, method: reqwest::Method, path: &str, body: &T) -> Result<reqwest::Response>
    where
        T: Serialize + ?Sized,
    {
        let url = format!("{}/{}", self.base_uri, path);
        let mut request = self.client.request(method, &url);
        if let Some(token) = &self.token {
            request = request.query(&[("token", token)])
        }

        request.json(body).send().context(RequestFailed)
    }
}

/// An error that can happen while processing a request.
#[derive(Debug, Snafu)]
pub enum Error {
    /// The request failed to complete.
    RequestFailed { source: reqwest::Error },
    /// Parsing response failed (e.g. some fields are missing).
    JsonFailed { source: reqwest::Error },
    /// Response is empty when it's not expected.
    ResponseEmpty { source: reqwest::Error },

    /// An empty image was provided.
    ImageEmpty,
    /// Provided token is invalid, expired, etc.
    InvalidToken,
    /// Provided image is >1MB.
    ImageTooLarge,
    /// Limit or quota was reached.
    RateLimit { message: String },
    /// Something went wrong on the trace.moe side of things.
    InternalServerError,
}

/// A convenience type to wrap [`Error`].
pub type Result<T, E = Error> = std::result::Result<T, E>;