kyanite 0.5.1

A modular, expandable, gallery collector.
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

use crate::collectors::danbooru::DanbooruCollector;
use log::{debug, error, info};

use crate::collectors::e621::E621Collector;
use crate::collectors::gelbooru::GelbooruCollector;
use crate::collectors::konachan::KonachanCollector;
use crate::collectors::rule34::Rule34Collector;
use crate::collectors::xbooru::XBooruCollector;
use crate::collectors::yandere::YandereCollector;
use crate::item::KyaniteItem;
use crate::manifest::KyaniteManifest;
use crate::params::KyaniteParams;
use crate::stats::StatsContainer;
use crate::utility::KyaniteUtility;

#[async_trait::async_trait]
pub trait KyaniteCollector {
    fn id(&self) -> &'static str;
    fn name(&self) -> &'static str;
    fn api_base(&self) -> &'static str;
    fn site_base(&self) -> &'static str;
    fn tags_argument(&self) -> &'static str;
    fn page_argument(&self) -> &'static str;
    fn starting_marker(&self) -> &'static str {
        if self.api_base().contains('?') {
            "&"
        } else {
            "?"
        }
    }
    fn api_by_page(&self, tags: String, page: u64) -> String {
        let api = format!(
            "{}{}{}={}&{}={}",
            &self.api_base(),
            &self.starting_marker(),
            &self.tags_argument(),
            tags,
            &self.page_argument(),
            page
        );
        debug!("{} API: {}", &self.name(), &api);
        api
    }
    async fn collect(&self, tags: Vec<String>) -> anyhow::Result<Vec<KyaniteItem>>;
}

pub struct CollectorCore {
    stats: StatsContainer,
    params: KyaniteParams,
    manifest: KyaniteManifest,
    collectors: Vec<Box<dyn KyaniteCollector>>,
}

impl CollectorCore {
    pub fn new(params: KyaniteParams) -> anyhow::Result<Self> {
        let stats = StatsContainer::new();
        let collectors = vec![
            DanbooruCollector::boxed(),
            E621Collector::boxed(),
            GelbooruCollector::boxed(),
            KonachanCollector::boxed(),
            Rule34Collector::boxed(),
            XBooruCollector::boxed(),
            YandereCollector::boxed(),
        ];
        let manifest = KyaniteManifest::new()?;
        Ok(Self {
            stats,
            params,
            manifest,
            collectors,
        })
    }

    pub async fn collect(&mut self) -> anyhow::Result<Vec<KyaniteItem>> {
        info!(
            "Searching for {} on {}.",
            &self.params.tags.join(", "),
            &self.params.sources.join(", ")
        );
        let mut items = Vec::new();
        for collector in &self.collectors {
            for source in &self.params.sources {
                if source == collector.id() || source == collector.name() || source == "all" {
                    let collected = match collector.collect((&self.params.tags).to_owned()).await {
                        Ok(clctd) => clctd,
                        Err(why) => {
                            error!(
                                "Failed collecting items from {}: {:#?}",
                                collector.id(),
                                why
                            );
                            Vec::new()
                        }
                    };
                    for item in collected {
                        items.push(item);
                    }
                }
            }
        }
        let trimmed = KyaniteItem::trim(items);
        let skipped = KyaniteItem::skip(trimmed)?;
        let sorted = KyaniteItem::sort(skipped);
        Ok(sorted)
    }

    pub async fn download(&mut self, items: Option<Vec<KyaniteItem>>) -> anyhow::Result<()> {
        let items = match items {
            Some(items) => items,
            None => self.collect().await?,
        };
        let total = items.len();
        for mut item in items {
            let resp = item.save(&mut self.stats, &mut self.manifest).await?;
            info!(
                "{} [{}] [{}] [{}/{}]: {}",
                resp,
                self.stats.describe(),
                KyaniteUtility::human_size(self.stats.size, 3f64, "GiB"),
                self.stats.count(),
                &total,
                item.describe()
            );
        }
        Ok(())
    }
}