1pub mod errors;
5
6use serde::{Deserialize, Serialize};
7use serde_json::ser::to_string as to_json_string;
8use std::collections::hash_map::Entry;
9use std::collections::HashMap;
10use std::io::{self, Write};
11#[derive(Serialize, Deserialize, Debug, PartialEq, Hash, Eq, Ord, PartialOrd)]
13pub struct Task {
14 pub name: String,
16 pub progress: Status,
24}
25impl Task {
26 pub fn new(name: impl Into<String>) -> Self {
27 Task {
28 name: name.into(),
29 progress: Status::Pending { total: None },
30 }
31 }
32 pub fn total(mut self, new_total: usize) -> Self {
36 self.progress = Status::Pending {
37 total: Some(new_total),
38 };
39 self
40 }
41 pub fn name(mut self, new_name: impl Into<String>) -> Self {
43 self.name = new_name.into();
44 self
45 }
46}
47impl From<String> for Task {
48 fn from(name: String) -> Self {
49 Task::new(name)
50 }
51}
52#[derive(Serialize, Deserialize, Debug, PartialEq, Hash, Eq, Ord, PartialOrd)]
57#[serde(tag = "status")]
58pub enum Status {
59 #[serde(rename = "pending")]
79 Pending { total: Option<usize> },
80 #[serde(rename = "error")]
82 Error { message: String },
83 #[serde(rename = "finished")]
85 Finished,
86 #[serde(rename = "running")]
88 Running,
89 #[serde(rename = "in_progress")]
97 InProgress {
98 done: usize,
99 total: usize,
100 },
102}
103#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
113pub struct TaskManager {
114 pub tasks: HashMap<String, Task>,
115 pub task_list: Vec<String>,
116 pub task_counter: usize,
117 silent: bool,
118}
119
120impl TaskManager {
121 #[inline]
122 fn output(&self) {
123 if self.silent {
124 return;
125 }
126 println!(
127 "{}",
128 to_json_string(&self.tasks.values().collect::<Vec<_>>()).expect("Should never happen")
129 );
130 io::stdout().flush().unwrap();
131 }
132 pub fn new() -> Self {
133 Self {
134 tasks: HashMap::new(),
135 task_list: Vec::new(),
136 task_counter: 0,
137 silent: false,
138 }
139 }
140
141 pub fn add_task(&mut self, task: Task) -> Result<(), errors::InterprogError> {
142 let name = task.name.clone();
143 match self.tasks.entry(name.clone()) {
144 Entry::Occupied(_) => return Err(errors::InterprogError::TaskAlreadyExists),
145 Entry::Vacant(entry) => entry.insert(task),
146 };
147 self.task_list.push(name);
148 Ok(())
149 }
150
151 pub fn start_task(&mut self, task_name: impl AsRef<str>) -> Result<(), errors::InterprogError> {
152 let task = &mut self
153 .tasks
154 .get_mut(task_name.as_ref())
155 .ok_or(errors::InterprogError::NonexistentTask)?;
156 if let Status::Pending { total } = &task.progress {
157 match total {
158 Some(total) => {
159 task.progress = Status::InProgress {
160 done: 0,
161 total: *total,
162 };
164 }
165 None => task.progress = Status::Running,
166 }
167 } else {
168 return Err(errors::InterprogError::TaskAlreadyStarted);
169 }
170 self.output();
171 Ok(())
172 }
173 pub fn start(&mut self) -> Result<(), errors::InterprogError> {
174 let task_name: String = self
175 .task_list
176 .get(self.task_counter)
177 .ok_or(errors::InterprogError::NonexistentTask)?
178 .clone();
179 self.start_task(&task_name)
180 }
181
182 pub fn increment_task(
183 &mut self,
184 task_name: impl AsRef<str>,
185 by: usize,
186 ) -> Result<(), errors::InterprogError> {
187 let task = &mut self
188 .tasks
189 .get_mut(task_name.as_ref())
190 .ok_or(errors::InterprogError::NonexistentTask)?;
191 match &task.progress {
193 Status::Pending { total: Some(total) } => {
194 task.progress = Status::InProgress {
195 done: 1,
196 total: *total,
197 };
199 }
200 Status::InProgress {
201 done,
202 total,
203 } => {
205 if done >= total {
206 return Err(errors::InterprogError::TaskAlreadyFinished);
207 }
208 task.progress = Status::InProgress {
210 done: done + by,
211 total: *total,
212 };
214 }
215 Status::Running | Status::Pending { total: None } => {
216 return Err(errors::InterprogError::InvalidTaskType)
217 }
218 Status::Finished | Status::Error { message: _ } => {
219 return Err(errors::InterprogError::TaskAlreadyFinished)
220 }
221 }
222 self.output();
223 Ok(())
224 }
225 pub fn increment(&mut self, by: usize) -> Result<(), errors::InterprogError> {
226 let task_name: String = self
227 .task_list
228 .get(self.task_counter)
229 .ok_or(errors::InterprogError::NonexistentTask)?
230 .clone();
231 self.increment_task(&task_name, by)
232 }
233
234 pub fn finish_task(
235 &mut self,
236 task_name: impl AsRef<str>,
237 ) -> Result<(), errors::InterprogError> {
238 let task = &mut self
239 .tasks
240 .get_mut(task_name.as_ref())
241 .ok_or(errors::InterprogError::NonexistentTask)?;
242 task.progress = Status::Finished;
255 self.task_counter += 1;
256 self.output();
257 Ok(())
258 }
259 pub fn finish(&mut self) -> Result<(), errors::InterprogError> {
260 let task_name: String = self
261 .task_list
262 .get(self.task_counter)
263 .ok_or(errors::InterprogError::NonexistentTask)?
264 .clone();
265 self.finish_task(&task_name)
266 }
267
268 pub fn error_task(
269 &mut self,
270 task_name: impl AsRef<str>,
271 message: impl Into<String>,
272 ) -> Result<(), errors::InterprogError> {
273 let task = &mut self
274 .tasks
275 .get_mut(task_name.as_ref())
276 .ok_or(errors::InterprogError::NonexistentTask)?;
277 task.progress = Status::Error {
278 message: message.into(),
279 };
280 self.task_counter += 1;
281 self.output();
282 Ok(())
283 }
284 pub fn error(&mut self, message: impl Into<String>) -> Result<(), errors::InterprogError> {
285 let task_name: String = self
286 .task_list
287 .get(self.task_counter)
288 .ok_or(errors::InterprogError::NonexistentTask)?
289 .clone();
290 self.error_task(&task_name, message)
291 }
292}
293impl Default for TaskManager {
294 fn default() -> Self {
295 Self::new()
296 }
297}
298#[cfg(test)]
299mod tests {
300 use crate::{Task, TaskManager};
301
302 #[test]
303 fn it_works() {
304 let mut manager = TaskManager::new();
305 manager.add_task(Task::new("name")).unwrap();
306 manager.start().unwrap();
307 manager.finish().unwrap();
308 }
309 #[test]
310 fn real_example() {
311 let mut manager = TaskManager::new();
312 manager.add_task(Task::new("Log in")).unwrap();
313 manager.start().unwrap();
314 manager.finish().unwrap();
315 let classes = vec!["English", "History", "Science", "Math"];
316 for class in &classes {
317 manager
318 .add_task(Task::new(format!("Scraping {class}")).total(4))
319 .unwrap();
320 }
321 for _ in 0..4 {
322 for class in &classes {
323 manager
324 .increment_task(format!("Scraping {class}"), 1)
325 .unwrap();
326 }
327 }
328 }
329 #[test]
330 fn static_names() {
331 let mut manager = TaskManager::new();
332 manager.add_task(Task::new("Log in")).unwrap();
333 manager.start_task("Log in").unwrap();
334 manager.finish().unwrap();
335 }
336}