async
This commit is contained in:
+81
-76
@@ -1,18 +1,20 @@
|
||||
extern crate async_zip;
|
||||
extern crate encoding;
|
||||
extern crate zip;
|
||||
extern crate tokio;
|
||||
|
||||
use async_zip::read::fs::ZipFileReader;
|
||||
use async_zip::write::{EntryOptions, ZipFileWriter};
|
||||
use async_zip::Compression;
|
||||
use clap::{Parser, ValueEnum};
|
||||
use encoding::label::encoding_from_whatwg_label;
|
||||
use encoding::EncodingRef;
|
||||
use encoding::{DecoderTrap, EncoderTrap};
|
||||
use regex::Regex;
|
||||
use std::fs;
|
||||
use std::io::{Read, Write};
|
||||
use std::sync::mpsc;
|
||||
use std::thread;
|
||||
use zip::ZipWriter;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
|
||||
use tokio::{fs, task};
|
||||
|
||||
/// transcode txt files in zip archieve
|
||||
/// transcode txt files in zip archive
|
||||
#[derive(Parser, Debug)]
|
||||
#[clap(author, version, about, long_about = None)]
|
||||
#[clap(propagate_version = true)]
|
||||
@@ -29,10 +31,6 @@ struct Cli {
|
||||
#[clap(arg_enum, short, long, default_value = "zstd")]
|
||||
compression: OutputFileCompression,
|
||||
|
||||
/// output compression level
|
||||
#[clap(arg_enum, short='l', long, value_parser = clap::value_parser!(i32).range(1..=9), default_value = "5")]
|
||||
compression_level: i32,
|
||||
|
||||
/// filename filter (regex)
|
||||
#[clap(short, long, default_value = r#".*\.txt$"#)]
|
||||
regex: String,
|
||||
@@ -55,17 +53,23 @@ enum OutputFileCompression {
|
||||
Deflate,
|
||||
/// Compress the file using BZIP2
|
||||
Bzip2,
|
||||
/// Compress the file using LZMA
|
||||
Lzma,
|
||||
/// Compress the file using ZStandard
|
||||
Zstd,
|
||||
/// Compress the file using XZ
|
||||
Xz,
|
||||
}
|
||||
|
||||
impl From<OutputFileCompression> for zip::CompressionMethod {
|
||||
impl From<OutputFileCompression> for Compression {
|
||||
fn from(compression: OutputFileCompression) -> Self {
|
||||
match compression {
|
||||
OutputFileCompression::Store => Self::Stored,
|
||||
OutputFileCompression::Deflate => Self::Deflated,
|
||||
OutputFileCompression::Bzip2 => Self::Bzip2,
|
||||
OutputFileCompression::Deflate => Self::Deflate,
|
||||
OutputFileCompression::Bzip2 => Self::Bz,
|
||||
OutputFileCompression::Lzma => Self::Lzma,
|
||||
OutputFileCompression::Zstd => Self::Zstd,
|
||||
OutputFileCompression::Xz => Self::Xz,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -76,43 +80,55 @@ struct FileData {
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
fn reader_task(tx: mpsc::Sender<FileData>, input_filename: String, regex: Regex) {
|
||||
let zip_file = fs::File::open(input_filename).unwrap();
|
||||
let mut archive = zip::ZipArchive::new(zip_file).unwrap();
|
||||
async fn reader_task(tx: UnboundedSender<FileData>, input_filename: String, regex: Regex) {
|
||||
let archive = ZipFileReader::new(input_filename).await.unwrap();
|
||||
|
||||
let mut source_files: Vec<String> = archive
|
||||
.file_names()
|
||||
.filter(|name| regex.is_match(name))
|
||||
.map(|s| s.to_string())
|
||||
let mut source_files: Vec<(usize, String, u32)> = archive
|
||||
.entries()
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, entry)| !entry.dir())
|
||||
.filter(|(_, entry)| regex.is_match(entry.name()))
|
||||
.map(|(index, entry)| {
|
||||
(
|
||||
index,
|
||||
entry.name().to_string(),
|
||||
entry.uncompressed_size().unwrap(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
source_files.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.partial_cmp(name_b).unwrap());
|
||||
|
||||
println!("processing {} files...", source_files.len());
|
||||
|
||||
source_files.sort();
|
||||
|
||||
for name in source_files {
|
||||
let mut file = archive.by_name(name.as_str()).unwrap();
|
||||
let mut data = Vec::with_capacity(file.size().try_into().unwrap());
|
||||
file.read_to_end(&mut data).unwrap();
|
||||
drop(file);
|
||||
let mut count: usize = 0;
|
||||
for (index, name, uncompressed_size) in source_files {
|
||||
let mut entry_reader = archive.entry_reader(index).await.unwrap();
|
||||
let mut data = Vec::with_capacity(uncompressed_size.try_into().unwrap());
|
||||
entry_reader.read_to_end(&mut data).await.unwrap();
|
||||
drop(entry_reader);
|
||||
|
||||
tx.send(FileData { name, data }).unwrap();
|
||||
count += 1;
|
||||
}
|
||||
|
||||
println!("read done ✅");
|
||||
println!("read {count} files done ✅");
|
||||
}
|
||||
|
||||
fn transcoder_task(
|
||||
rx: mpsc::Receiver<FileData>,
|
||||
tx: mpsc::Sender<FileData>,
|
||||
async fn transcoder_task(
|
||||
mut rx: UnboundedReceiver<FileData>,
|
||||
tx: UnboundedSender<FileData>,
|
||||
encoding_from: EncodingRef,
|
||||
encoding_to: EncodingRef,
|
||||
) {
|
||||
while let Ok(FileData { name, data }) = rx.recv() {
|
||||
let text = encoding_from.decode(&data, DecoderTrap::Ignore).unwrap();
|
||||
let new_data = encoding_to
|
||||
.encode(text.as_str(), EncoderTrap::Ignore)
|
||||
.unwrap();
|
||||
while let Some(FileData { name, data }) = rx.recv().await {
|
||||
let new_data = task::spawn_blocking(move || {
|
||||
let text = encoding_from.decode(&data, DecoderTrap::Ignore).unwrap();
|
||||
encoding_to
|
||||
.encode(text.as_str(), EncoderTrap::Ignore)
|
||||
.unwrap()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
tx.send(FileData {
|
||||
name,
|
||||
data: new_data,
|
||||
@@ -122,59 +138,48 @@ fn transcoder_task(
|
||||
println!("transcode done ✅");
|
||||
}
|
||||
|
||||
fn writer_task(
|
||||
rx: mpsc::Receiver<FileData>,
|
||||
async fn writer_task(
|
||||
mut rx: UnboundedReceiver<FileData>,
|
||||
output_filename: String,
|
||||
compression: zip::CompressionMethod,
|
||||
compression_level: i32,
|
||||
compression: Compression,
|
||||
) {
|
||||
let options = zip::write::FileOptions::default()
|
||||
.compression_method(compression)
|
||||
.compression_level(Some(compression_level));
|
||||
let mut outfile = fs::File::create(output_filename)
|
||||
.await
|
||||
.expect("output file");
|
||||
let mut writer = ZipFileWriter::new(&mut outfile);
|
||||
|
||||
let mut outfile = fs::File::create(output_filename).expect("output file");
|
||||
let mut zip_writer = ZipWriter::new(&mut outfile);
|
||||
|
||||
while let Ok(FileData { name, data }) = rx.recv() {
|
||||
zip_writer.start_file(name, options).unwrap();
|
||||
zip_writer.write_all(&data).unwrap();
|
||||
while let Some(FileData { name, data }) = rx.recv().await {
|
||||
let opts = EntryOptions::new(name, compression);
|
||||
writer.write_entry_whole(opts, &data).await.unwrap();
|
||||
}
|
||||
zip_writer.finish().unwrap();
|
||||
writer.close().await.unwrap();
|
||||
println!("write done ✅");
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let args = Cli::parse();
|
||||
|
||||
let regex = Regex::new(&args.regex).expect("regex");
|
||||
let encoding_from = encoding_from_whatwg_label(&args.from).expect("input encoding");
|
||||
let encoding_to = encoding_from_whatwg_label(&args.to).expect("output encoding");
|
||||
let compression: zip::CompressionMethod = args.compression.into();
|
||||
let compression_level = args.compression_level;
|
||||
let compression: Compression = args.compression.into();
|
||||
let input_filename = args.src;
|
||||
let output_filename = args.dst;
|
||||
|
||||
let (reader_tx, reader_rx) = mpsc::channel::<FileData>();
|
||||
let (transcoder_tx, transcoder_rx) = mpsc::channel::<FileData>();
|
||||
let (reader_tx, reader_rx) = mpsc::unbounded_channel::<FileData>();
|
||||
let (transcoder_tx, transcoder_rx) = mpsc::unbounded_channel::<FileData>();
|
||||
|
||||
let handles = vec![
|
||||
thread::spawn(move || reader_task(reader_tx, input_filename, regex)),
|
||||
thread::spawn(move || {
|
||||
transcoder_task(reader_rx, transcoder_tx, encoding_from, encoding_to)
|
||||
}),
|
||||
thread::spawn(move || {
|
||||
writer_task(
|
||||
transcoder_rx,
|
||||
output_filename,
|
||||
compression,
|
||||
compression_level,
|
||||
)
|
||||
}),
|
||||
];
|
||||
|
||||
for handle in handles {
|
||||
handle.join().expect("thread paniced");
|
||||
}
|
||||
tokio::try_join!(
|
||||
tokio::spawn(reader_task(reader_tx, input_filename, regex)),
|
||||
tokio::spawn(transcoder_task(
|
||||
reader_rx,
|
||||
transcoder_tx,
|
||||
encoding_from,
|
||||
encoding_to
|
||||
)),
|
||||
tokio::spawn(writer_task(transcoder_rx, output_filename, compression))
|
||||
)?;
|
||||
|
||||
println!("all done ✅");
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user