notatin/
progress.rs

1/*
2 * Copyright 2023 Aon Cyber Solutions
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16use crate::err::Error;
17use crossterm::{cursor, QueueableCommand};
18use std::io;
19use std::io::{Stdout, Write};
20
21pub fn new(update_console: bool) -> Box<dyn UpdateProgressTrait> {
22    if update_console {
23        Box::new(UpdateConsole {
24            need_final_newline: false,
25            stdout: io::stdout(),
26        })
27    } else {
28        Box::new(UpdateNull {})
29    }
30}
31
32pub trait UpdateProgressTrait {
33    fn update_progress(&mut self, index: usize) -> Result<(), Error>;
34    fn update(&mut self, msg: &str) -> Result<(), Error>;
35    fn write(&mut self, msg: &str) -> Result<(), Error>;
36}
37
38struct UpdateConsole {
39    need_final_newline: bool,
40    stdout: Stdout,
41}
42
43impl UpdateProgressTrait for UpdateConsole {
44    fn update_progress(&mut self, index: usize) -> Result<(), Error> {
45        if index % 1000 == 0 {
46            self.stdout.write_all(".".as_bytes())?;
47            self.stdout.flush()?;
48        }
49        Ok(())
50    }
51
52    fn update(&mut self, msg: &str) -> Result<(), Error> {
53        self.stdout.queue(cursor::SavePosition)?;
54        self.stdout.write_all(msg.as_bytes())?;
55        self.stdout.queue(cursor::RestorePosition)?;
56        self.stdout.flush()?;
57        self.need_final_newline = true;
58        Ok(())
59    }
60
61    fn write(&mut self, msg: &str) -> Result<(), Error> {
62        self.stdout.write_all(msg.as_bytes())?;
63        self.stdout.flush()?;
64        Ok(())
65    }
66}
67
68impl Drop for UpdateConsole {
69    fn drop(&mut self) {
70        if self.need_final_newline {
71            self.stdout.write_all("\n".as_bytes()).unwrap_or_default();
72            self.stdout.flush().unwrap_or_default();
73        }
74    }
75}
76
77struct UpdateNull {}
78
79impl UpdateProgressTrait for UpdateNull {
80    fn update_progress(&mut self, _index: usize) -> Result<(), Error> {
81        Ok(())
82    }
83
84    fn update(&mut self, _msg: &str) -> Result<(), Error> {
85        Ok(())
86    }
87
88    fn write(&mut self, _msg: &str) -> Result<(), Error> {
89        Ok(())
90    }
91}