csv_pipeline/
headers.rs

1use crate::Row;
2use csv::StringRecordIter;
3use std::collections::BTreeMap;
4use std::fmt;
5
6/// The headers of a CSV file
7#[derive(Debug, Clone, PartialEq)]
8pub struct Headers {
9	indexes: BTreeMap<String, usize>,
10	row: Row,
11}
12pub enum RenameError {
13	DuplicateColumn(usize),
14	MissingColumn,
15}
16impl fmt::Display for RenameError {
17	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
18		match self {
19			RenameError::DuplicateColumn(index) => write!(f, "Duplicate column at index {}", index),
20			RenameError::MissingColumn => write!(f, "Missing column"),
21		}
22	}
23}
24
25impl Headers {
26	pub fn new() -> Self {
27		Headers {
28			indexes: BTreeMap::new(),
29			row: Row::new(),
30		}
31	}
32
33	/// Returns `Error::MissingColumn` if `from` is non-existant or `Error::DuplicateColumn` the new name already exists
34	pub fn rename(&mut self, from: &str, to: &str) -> Result<(), RenameError> {
35		if let Some(index) = self.get_index(to) {
36			return Err(RenameError::DuplicateColumn(index));
37		}
38		let index = match self.indexes.remove(from) {
39			Some(index) => index,
40			None => return Err(RenameError::MissingColumn),
41		};
42		self.indexes.insert(to.to_string(), index);
43		let mut row_vec: Vec<_> = self.row.into_iter().collect();
44		row_vec[index] = to;
45		self.row = row_vec.into_iter().collect();
46		Ok(())
47	}
48
49	/// Returns false if the field already exists
50	pub fn push_field(&mut self, name: &str) -> bool {
51		if self.indexes.contains_key(name) {
52			return false;
53		}
54
55		self.row.push_field(name);
56		self.indexes.insert(name.to_string(), self.row.len() - 1);
57
58		true
59	}
60
61	pub fn contains(&self, name: &str) -> bool {
62		self.indexes.contains_key(name)
63	}
64
65	pub fn get_field<'a>(&self, row: &'a Row, name: &str) -> Option<&'a str> {
66		self.indexes.get(name).and_then(|index| row.get(*index))
67	}
68
69	pub fn get_index(&self, name: &str) -> Option<usize> {
70		self.indexes.get(name).copied()
71	}
72
73	pub fn get_row(&self) -> &Row {
74		&self.row
75	}
76
77	/// If a column is duplicated, errors with the column name
78	pub fn from_row(row: Row) -> Result<Self, String> {
79		let mut header = Headers::new();
80		for field in &row {
81			let added = header.push_field(field);
82			if !added {
83				return Err(field.to_string());
84			}
85		}
86		Ok(header)
87	}
88}
89
90impl<'a> IntoIterator for &'a Headers {
91	type Item = &'a str;
92	type IntoIter = StringRecordIter<'a>;
93
94	fn into_iter(self) -> StringRecordIter<'a> {
95		self.row.into_iter()
96	}
97}
98impl From<Headers> for Row {
99	fn from(headers: Headers) -> Row {
100		headers.row
101	}
102}