From 9867b2eaf656dfa4f25409e5b29ecfe738a138cf Mon Sep 17 00:00:00 2001 From: lifegpc Date: Fri, 10 Mar 2023 14:17:35 +0000 Subject: [PATCH] Add ai_type to json data file. Remove unneeded unwrap when writing data to json. --- src/data/data.rs | 32 ++++++++++++++++++++++++++++++++ src/data/json.rs | 39 ++++++++++++++++++++++----------------- src/ext/json.rs | 40 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+), 17 deletions(-) diff --git a/src/data/data.rs b/src/data/data.rs index 0b55966..1d44281 100644 --- a/src/data/data.rs +++ b/src/data/data.rs @@ -5,10 +5,27 @@ use crate::opt::author_name_filter::AuthorFiler; use crate::opthelper::get_helper; use crate::pixiv_link::PixivID; use crate::pixiv_link::ToPixivID; +use int_enum::IntEnum; use json::JsonValue; use std::convert::TryInto; use xml::unescape; +/// The type of the artwork +#[repr(u8)] +#[derive(Clone, Copy, Debug, Eq, PartialEq, IntEnum)] +pub enum PixivAiType { + /// Other types + NonAI = 1, + /// Artwork is generated by AI. + AI = 2, +} + +impl PixivAiType { + pub fn is_ai(&self) -> bool { + *self == Self::AI + } +} + /// Pixiv's basic data pub struct PixivData { /// ID @@ -20,6 +37,8 @@ pub struct PixivData { pub description: Option, /// Tags (Original, translated) pub tags: Option)>>, + /// AI type + pub ai_type: Option, } impl PixivData { @@ -34,6 +53,7 @@ impl PixivData { author: None, description: None, tags: None, + ai_type: None, }) } @@ -100,6 +120,18 @@ impl PixivData { } } self.tags.replace(tags); + if self.ai_type.is_none() { + let ai_type = value["aiType"].as_u8(); + match ai_type { + Some(ai_type) => match PixivAiType::from_int(ai_type) { + Ok(ai_type) => { + self.ai_type.replace(ai_type); + } + Err(_) => {} + }, + None => {} + } + } } } diff --git a/src/data/json.rs b/src/data/json.rs index b06345d..fe350de 100644 --- a/src/data/json.rs +++ b/src/data/json.rs @@ -1,10 +1,11 @@ use super::fanbox::FanboxData; use crate::data::data::PixivData; -use crate::ext::json::ToJson; +use crate::ext::json::{ToJson, ToJson2}; use crate::gettext; use crate::parser::description::parse_description; use crate::pixiv_link::PixivID; use crate::pixiv_link::ToPixivID; +use int_enum::IntEnum; use json::JsonValue; use std::collections::HashMap; use std::convert::From; @@ -34,14 +35,9 @@ impl JSONDataFile { }) } - pub fn add(&mut self, key: &str, value: T) -> Result<(), ()> { - let v = value.to_json(); - if v.is_some() { - self.maps.insert(String::from(key), v.unwrap()); - Ok(()) - } else { - Err(()) - } + pub fn add(&mut self, key: &str, value: T) { + let v = value.to_json2(); + self.maps.insert(String::from(key), v); } pub fn save + ?Sized>(&self, path: &S) -> bool { @@ -92,17 +88,16 @@ impl From<&PixivData> for JSONDataFile { maps: HashMap::new(), }; if p.title.is_some() { - f.add("title", p.title.as_ref().unwrap()).unwrap(); + f.add("title", p.title.as_ref().unwrap()); } if p.author.is_some() { - f.add("author", p.author.as_ref().unwrap()).unwrap(); + f.add("author", p.author.as_ref().unwrap()); } if p.description.is_some() { - f.add("description", p.description.as_ref().unwrap()) - .unwrap(); + f.add("description", p.description.as_ref().unwrap()); let pd = parse_description(p.description.as_ref().unwrap()); if pd.is_some() { - f.add("parsed_description", pd.unwrap()).unwrap(); + f.add("parsed_description", pd.unwrap()); } } match p.tags.as_ref() { @@ -118,10 +113,20 @@ impl From<&PixivData> for JSONDataFile { ]) .unwrap(); } - f.add("tags", t).unwrap(); + f.add("tags", t); } None => {} } + match &p.ai_type { + Some(ai_type) => { + f.add("is_ai", ai_type.is_ai()); + f.add("ai_type", ai_type.int_value()); + } + None => { + f.add("is_ai", JsonValue::Null); + f.add("ai_type", JsonValue::Null); + } + } f } } @@ -132,7 +137,7 @@ impl From for JSONDataFile { id: d.id.clone(), maps: HashMap::new(), }; - f.add("raw", d.raw).unwrap(); + f.add("raw", d.raw); f } } @@ -143,7 +148,7 @@ impl From<&FanboxData> for JSONDataFile { id: d.id.clone(), maps: HashMap::new(), }; - f.add("raw", d.raw.clone()).unwrap(); + f.add("raw", d.raw.clone()); f } } diff --git a/src/ext/json.rs b/src/ext/json.rs index 2e23fe1..725fc4b 100644 --- a/src/ext/json.rs +++ b/src/ext/json.rs @@ -92,6 +92,46 @@ impl ToJson2 for RwLockWriteGuard<'_, T> { } } +impl ToJson for bool { + #[inline] + fn to_json(&self) -> Option { + Some(JsonValue::from(self.clone())) + } +} + +impl ToJson2 for bool { + #[inline] + fn to_json2(&self) -> JsonValue { + JsonValue::from(self.clone()) + } +} + +macro_rules! impl_tojson_for_int { + ($int:ty) => { + impl ToJson for $int { + #[inline] + fn to_json(&self) -> Option { + Some(JsonValue::from(self.clone())) + } + } + impl ToJson2 for $int { + #[inline] + fn to_json2(&self) -> JsonValue { + JsonValue::from(self.clone()) + } + } + }; +} + +impl_tojson_for_int!(i8); +impl_tojson_for_int!(i16); +impl_tojson_for_int!(i32); +impl_tojson_for_int!(i64); +impl_tojson_for_int!(u8); +impl_tojson_for_int!(u16); +impl_tojson_for_int!(u32); +impl_tojson_for_int!(u64); + pub trait FromJson where Self: Sized,