1use std::fs::File;
2use std::io::{self, ErrorKind, prelude::*, Write};
3use std::time::{Duration, SystemTime};
4
5use crate::{SimpleError, XResult};
6use crate::util_file;
7use crate::util_msg;
8use crate::util_size;
9
10pub const DEFAULT_BUF_SIZE: usize = 8 * 1024;
11
12pub fn stdout_or_file_write(file: Option<&str>, overwrite: bool) -> XResult<Box<dyn Write>> {
13 match file {
14 None => Ok(Box::new(io::stdout())),
15 Some(output) => {
16 if File::open(output).is_ok() && !overwrite {
17 return Err(SimpleError::new(format!("File exists: {}", output)).into());
18 }
19 Ok(Box::new(File::create(output).map_err(|e| {
20 SimpleError::new(format!("Create file: {}, failed: {}", output, e))
21 })?))
22 }
23 }
24}
25
26pub struct PrintStatusContext {
27 pub print_interval_time: Duration,
28 pub print_interval_bytes: i64,
29 pub init_print_time: SystemTime,
30 pub last_print_time: SystemTime,
31 pub total_written_bytes: i64,
32}
33
34impl PrintStatusContext {
35 pub fn new() -> Self {
36 Self::new_with(Duration::from_millis(100), 512 * 1024)
37 }
38
39 pub fn new_with(print_interval_time: Duration, print_interval_bytes: i64) -> Self {
40 Self {
41 print_interval_time,
42 print_interval_bytes,
43 init_print_time: SystemTime::now(),
44 last_print_time: SystemTime::now(),
45 total_written_bytes: 0,
46 }
47 }
48
49 pub fn check_print(&mut self, total: i64, written: i64) -> (bool, Duration) {
50 let now = SystemTime::now();
51 let total_cost = now.duration_since(self.init_print_time).unwrap_or_else(|_| Duration::from_millis(0));
52 let last_print_cost = now.duration_since(self.last_print_time).unwrap_or_else(|_| Duration::from_millis(0));
53 let should_update_status_line = || {
54 if total > written && (total - written < self.print_interval_bytes) {
55 return true;
56 }
57 if written > self.total_written_bytes && (written - self.total_written_bytes > self.print_interval_bytes) {
58 return true;
59 }
60 last_print_cost.as_millis() > self.print_interval_time.as_millis()
61 };
62 if should_update_status_line() {
63 self.last_print_time = now;
64 self.total_written_bytes = written;
65 (true, total_cost)
66 } else {
67 (false, total_cost)
68 }
69 }
70}
71
72impl Default for PrintStatusContext {
73 fn default() -> Self {
74 PrintStatusContext::new()
75 }
76}
77
78pub fn get_read_stdin_or_file(file: &str) -> XResult<Box<dyn Read>> {
79 if file.is_empty() {
80 Ok(Box::new(io::stdin()))
81 } else {
82 match File::open(util_file::resolve_file_path(file)) {
83 Ok(f) => Ok(Box::new(f)),
84 Err(err) => Err(SimpleError::new(format!("Open file {}, erorr: {}", file, err)).into()),
85 }
86 }
87}
88
89pub fn read_to_string(read: &mut dyn Read) -> XResult<String> {
90 let mut buffer = String::new();
91 read.read_to_string(&mut buffer)?;
92 Ok(buffer)
93}
94
95pub fn read_to_bytes(read: &mut dyn Read) -> XResult<Vec<u8>> {
96 let mut buffer = vec![];
97 read.read_to_end(&mut buffer)?;
98 Ok(buffer)
99}
100
101pub fn copy_io_default<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W, total: i64) -> io::Result<u64>
102 where R: io::Read, W: io::Write {
103 copy_io_with_head(reader, writer, total, "Downloading", &mut PrintStatusContext::default())
104}
105
106pub fn copy_io<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W, total: i64, print_status_context: &mut PrintStatusContext)
107 -> io::Result<u64>
108 where R: io::Read, W: io::Write {
109 copy_io_with_head(reader, writer, total, "Downloading", print_status_context)
110}
111
112pub fn copy_io_with_head<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W, total: i64, head: &str, print_status_context: &mut PrintStatusContext) -> io::Result<u64>
113 where R: io::Read, W: io::Write {
114 let written = copy_io_callback(reader, writer, total, print_status_context, &mut |total, written, _len, print_status_context| {
115 print_status_last_line(head, total, written as i64, print_status_context);
116 });
117 println!();
118 written
119}
120
121pub fn copy_io_callback<R: ?Sized, W: ?Sized, FCallback>(reader: &mut R, writer: &mut W, total: i64, print_status_context: &mut PrintStatusContext, callback: &mut FCallback) -> io::Result<u64>
122 where R: io::Read,
123 W: io::Write,
124 FCallback: Fn(i64, u64, usize, &mut PrintStatusContext) {
125 let mut written = 0u64;
126 let mut buf: [u8; DEFAULT_BUF_SIZE] = [0u8; DEFAULT_BUF_SIZE];
127 loop {
128 let len = match reader.read(&mut buf) {
129 Ok(0) => return Ok(written),
130 Ok(len) => len,
131 Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
132 Err(e) => return Err(e),
133 };
134 writer.write_all(&buf[..len])?;
135 written += len as u64;
136 callback(total, written, len, print_status_context);
137 }
138}
139
140pub fn print_status_last_line(head: &str, total: i64, written: i64, print_status_context: &mut PrintStatusContext) {
141 let mut download_speed = "-".to_string();
142 let (is_print, cost) = print_status_context.check_print(total, written);
143 if !is_print {
144 return;
145 }
146 let cost_as_secs = cost.as_secs();
147 if cost_as_secs > 0 {
148 download_speed = format!("{}/s", util_size::get_display_size(written / (cost_as_secs as i64)));
149 }
150 if total > 0 {
151 util_msg::print_lastline(&format!("{}, Total: {}, Finished: {}, Speed: {}",
152 head,
153 util_size::get_display_size(total),
154 util_size::get_display_size(written),
155 download_speed));
156 } else {
157 util_msg::print_lastline(&format!("{}, Finished: {}, Speed: {}",
158 head,
159 util_size::get_display_size(written),
160 download_speed));
161 }
162}