tidy_browser/
binary_cookies.rs

1use std::{fmt::Display, fs::File, io::IoSlice, path::Path};
2
3use binary_cookies::{
4    cookie::Cookie,
5    sync::{self, DecodeBinaryCookie},
6};
7use rayon::{iter::ParallelIterator, prelude::ParallelBridge};
8use snafu::ResultExt;
9
10use crate::{
11    args::Format,
12    error::{BinaryCookiesSnafu, IoSnafu, JsonSnafu, Result},
13    utils,
14};
15
16#[derive(Clone, Copy)]
17#[derive(Debug)]
18#[derive(Default)]
19#[derive(PartialEq, Eq, PartialOrd, Ord)]
20pub struct BinaryCookiesWriter;
21
22impl BinaryCookiesWriter {
23    pub fn write_data<A, B, S>(path: A, out: B, sep: S, format: Format) -> Result<()>
24    where
25        A: AsRef<Path>,
26        B: AsRef<Path>,
27        S: Display + Send + Clone + 'static + Sync,
28    {
29        let path = path.as_ref();
30        let out_path = out.as_ref();
31
32        let file = File::open(path).with_context(|_| IoSnafu { path: path.to_owned() })?;
33        let mut out = File::options()
34            .write(true)
35            .create(true)
36            .truncate(true)
37            .open(out_path)
38            .with_context(|_| IoSnafu { path: out_path.to_owned() })?;
39
40        let a = file
41            .decode()
42            .context(BinaryCookiesSnafu)?;
43        let (pages_handle, _meta_decoder) = a.into_handles();
44
45        let tmp = pages_handle
46            .decoders()
47            .par_bridge()
48            .filter_map(|mut v| v.decode().ok())
49            .map(sync::CookieHandle::into_decoders);
50        match format {
51            Format::Csv => {
52                let csvs: Vec<_> = tmp
53                    .flat_map(|v| {
54                        v.par_bridge().filter_map(|mut v| {
55                            v.decode()
56                                .map(|v| v.to_csv(sep.clone()))
57                                .ok()
58                        })
59                    })
60                    .collect();
61
62                let mut slices = Vec::with_capacity(2 + csvs.len() * 2);
63
64                let header = Cookie::csv_header(sep.clone());
65                slices.push(IoSlice::new(header.as_bytes()));
66                slices.push(IoSlice::new(b"\n"));
67
68                for cookie in &csvs {
69                    slices.push(IoSlice::new(cookie.as_bytes()));
70                    slices.push(IoSlice::new(b"\n"));
71                }
72                utils::write_all_vectored(&mut out, &mut slices)
73                    .with_context(|_| IoSnafu { path: out_path.to_owned() })
74            },
75            Format::Json => {
76                let cookies = tmp
77                    .map(|v| {
78                        v.par_bridge()
79                            .filter_map(|mut v| v.decode().ok())
80                            .collect::<Vec<_>>()
81                    })
82                    .collect::<Vec<_>>();
83                serde_json::to_writer(out, &cookies).context(JsonSnafu)
84            },
85            Format::JsonLines => {
86                let cookies = tmp
87                    .flat_map(|v| {
88                        v.par_bridge()
89                            .filter_map(|mut v| v.decode().ok())
90                            .collect::<Vec<_>>()
91                    })
92                    .collect::<Vec<_>>();
93                serde_jsonlines::write_json_lines(out_path, cookies)
94                    .with_context(|_| IoSnafu { path: out_path.to_owned() })
95            },
96        }
97    }
98}
tidy_browser/binary_cookies.rs

tidy_browser/
binary_cookies.rs
