rust值日排班工具 - 果核剥壳

公司里每个月都需要做一份卫生值日安排表，一直都是人工编辑安排，需要根据人数、区域和工作日进行汇总分析，人工编辑难免出错，因为学过一段rust于是想用rust做一个工具自动生成

use std::{env, io};
use std::collections::{HashMap, HashSet};
use std::io::BufRead;

use chrono::{Datelike, Local};
use rand::{Rng, thread_rng};
use rand::prelude::SliceRandom;
use rust_xlsxwriter::{Format, FormatAlign, FormatBorder, Workbook, XlsxError};

fn main() -> Result<(), XlsxError> {
let stdin = io::stdin();
// 参与人员
let mut staff:Vec<&str> = vec![
"赵云", "马超", "黄忠", "诸葛亮",
"刘备", "关羽", "张飞", "曹操",
"张辽", "孙权", "周瑜", "黄盖",
"华佗", "貂蝉", "大乔", "小乔"
];
println!("**************************************");
println!("**************************************");
println!("****** 欢迎使用员工值日生成系统 ******");
println!("**************************************");
println!("**************************************");
println!();
println!("描述：随机分配员工顺序，并按照每人值日总次数分配，同时兼顾左右侧分配平衡");
println!();
println!("tips：按回车跳过可取默认值");
println!("tips：输出文件到当前目录下");
println!();
println!();
println!("1/3 请输入参与人员（不能少于6个且不能重复），使用空格键分隔，默认：赵云 马超 黄忠 诸葛亮 刘备 关羽 张飞");
let mut staff_line = String::new();
loop {
stdin.lock().read_line(&mut staff_line).unwrap();
staff_line = staff_line.trim().to_string();
// 输入项超过6个
if (staff_line.split_whitespace().collect::<HashSet<&str>>().len() == staff_line.split_whitespace().count()
&& staff_line.split_whitespace().count() >= 6 )
|| staff_line.is_empty() {
break;
} else {
staff_line.clear();
println!("参与人员少于6个或有重复，请重新输入");
}
}

if !staff_line.is_empty() {
staff = staff_line.trim()
.split_whitespace().collect();
}

// 年月
let now = Local::now();
let mut year = now.year();
let mut month = now.month();
println!("2/3 请输入年份，默认当年：{}", year);
let mut year_line = String::new();
loop {
stdin.lock().read_line(&mut year_line).unwrap();
if !year_line.trim().is_empty() {
if let Ok(y) = year_line.trim().parse::<i32>() {
year = y;
} else {
year_line = String::new();
println!("年份解析出错，请重新输入或回车跳过");
continue;
}
}
break;
}

println!("3/3 请输入月份，默认当月： {}", month);
let mut month_line = String::new();
loop {
stdin.lock().read_line(&mut month_line).unwrap();
if !month_line.trim().is_empty() {
if let Ok(m) = month_line.trim().parse::<u32>(){
month = m;
} else {
month_line = String::new();
println!("月份解析出错，请重新输入或回车跳过");
continue;
}
}
break;
}

let mut total_group = vec![];
let mut staff_count: HashMap<&str, usize> = HashMap::new();
let mut left_count: HashMap<&str, usize> = HashMap::new();
for s in &staff {
staff_count.insert(*s, 0);
left_count.insert(*s, 0);
}

staff.shuffle(&mut thread_rng());
for _ in 1..30 {
let arr: Vec<_> = staff.chunks_exact(3).collect();
let arr: Vec<_> = arr.iter().map(|&e| e).take(arr.len() - arr.len()%2).collect();
for x in 0..arr.len(){
total_group.push(arr[x].iter().map(|s| s.to_string()).collect::<Vec<String>>());
for &i in arr[x] {
let c = staff_count.get(i).unwrap_or(&0);
staff_count.insert(i, c + 1);
// 在左侧出现的次数
if x%2 == 0 {
let c = left_count.get(i).unwrap_or(&0);
left_count.insert(i, c + 1);
}
}
}
// 按照出现次数排序
staff.sort_by(|&a,&b| {
// 总次数排序
let mut ac = staff_count.get(&a).unwrap();
let mut bc = staff_count.get(&b).unwrap();
// 总次数相同 则按左侧出现次数排序
if ac == bc {
ac = left_count.get(&a).unwrap();
bc = left_count.get(&b).unwrap();
}

// 两侧不同则排序，否则随机排序
if ac != bc {
ac.cmp(bc)
} else {
if thread_rng().gen_bool(0.5) {
a.cmp(&b)
} else {
b.cmp(&a)
}
}
});

if total_group.len()>=60 {
break;
}
}

let days_in_month = num_days_month(year, month);
// 居中格式
let title_format = Format::new()
.set_border(FormatBorder::Thin)
.set_font_size(20)
.set_font_name("微软雅黑")
.set_align(FormatAlign::Center)
.set_align(FormatAlign::VerticalCenter)
;
// 标题字段
let field_formst = Format::new()
.set_font_name("微软雅黑")
.set_border(FormatBorder::Thin)
.set_font_size(12)
.set_bold()
.set_align(FormatAlign::Center)
.set_align(FormatAlign::VerticalCenter)
;

// 居中格式
let border_format = Format::new()
.set_font_name("微软雅黑")
.set_border(FormatBorder::Thin)
.set_font_size(12)
.set_align(FormatAlign::Center)
.set_align(FormatAlign::VerticalCenter)
;
let end_format= Format::new()
.set_font_name("微软雅黑")
.set_border(FormatBorder::Thin)
.set_font_size(12)
.set_align(FormatAlign::VerticalCenter)
.set_text_wrap()
;

let mut wb = Workbook::new();
let ws = wb.add_worksheet();
// 标题
let title = format!("{}年{}月值日安排", year, month);
ws.merge_range(0, 0, 0, 4, title.as_str(), &title_format)?;
ws.set_row_height(0, 38)?;
ws.set_column_width(0, 11)?;
ws.set_column_width(1, 10)?;
ws.set_column_width(2, 28)?;
ws.set_column_width(3, 28)?;
// 添加字段
ws.write_with_format(1, 0, "日期", &field_formst)?;
ws.write_with_format(1, 1, "星期", &field_formst)?;
ws.write_with_format(1, 2, "左大厅（包含卫生间）", &field_formst)?;
ws.write_with_format(1, 3, "右大厅（包含吸烟区）", &field_formst)?;
ws.write_with_format(1, 4, "绿植", &field_formst)?;
// 实际计数
let mut fac_i = 0;
// 周次
let mut week_i = 0;
// 行数
let mut data_row = 2;
// 月内星期开始
let mut w_start = 2;
// 月内星期结束
let mut w_end;
for day in 1..=days_in_month {
let date = Local::now()
.with_year(year).unwrap()
.with_month(month).unwrap()
.with_day(day).unwrap();
// 设置行高
ws.set_row_height(data_row, 20)?;

// 日期
let day_str = format!("{}月{}日", month, day);
// 星期
let weekday = to_week_name(date.weekday());
// 左大厅
let left_staff = total_group[fac_i*2].join("、");
// 右大厅
let right_staff = total_group[fac_i*2+1].join("、");

// 日期
ws.write_with_format(data_row, 0, day_str, &field_formst)?;
// 星期
ws.write_with_format(data_row, 1, weekday, &field_formst)?;

if vec![chrono::Weekday::Sat, chrono::Weekday::Sun].contains(&date.weekday()) {
ws.merge_range(data_row, 2, data_row, 4, "",&border_format)?;
} else {
// 左大厅
ws.write_with_format(data_row, 2, left_staff, &border_format)?;
// 右大厅
ws.write_with_format(data_row, 3, right_staff, &border_format)?;
fac_i += 1;
}

// 月初或是周一 设置开始行
if day == 1 || date.weekday() == chrono::Weekday::Mon {
w_start = data_row;
}

// 月末或周五 设置结束行
if (day == days_in_month && !vec![chrono::Weekday::Sat, chrono::Weekday::Sun].contains(&date.weekday()) ) || date.weekday() == chrono::Weekday::Fri {
w_end = data_row;
if w_start == w_end {
// 最后一天单条
ws.write_with_format(w_start, 4,staff[week_i], &border_format)?;
} else {
// 绿植
ws.merge_range(w_start, 4, w_end, 4, staff[week_i], &border_format)?;
}

week_i += 1;
}

data_row += 1;
}
let counting = staff_count.iter()
.map(|(k, v)| format!("{}（{}次）", *k, v))
.collect::<Vec<String>>()
.join(" ");
let ending = format!("值日汇总：{}", counting);

ws.merge_range(data_row,0, data_row, 4, ending.as_str(), &end_format)?;
ws.set_row_height(data_row, 60)?;
let mut cur_dir = env::current_dir().unwrap();
cur_dir.push(format!("{title}.xlsx"));
wb.save(cur_dir.to_str().unwrap())?;
Ok(())

}

/// 星期转汉字
fn to_week_name(weekday: chrono::Weekday) -> &'static str {
let weekday = match weekday {
chrono::Weekday::Mon => "星期一",
chrono::Weekday::Tue => "星期二",
chrono::Weekday::Wed => "星期三",
chrono::Weekday::Thu => "星期四",
chrono::Weekday::Fri => "星期五",
chrono::Weekday::Sat => "星期六",
chrono::Weekday::Sun => "星期日",
};
weekday
}

/// 获取月内的最大天数
fn num_days_month(year: i32, month: u32) -> u32 {
let days: u32 = match month {
2 if year % 4 == 0 && year % 100 != 0 || year % 400 == 0 => 29,
2 => 28,
4 | 6 | 9 | 11 => 30,
_ => 31,
};
days
}