From 405d828b96496a6be1db62a749a7ec4f1c19cdc6 Mon Sep 17 00:00:00 2001
From: Kasper <kasperkh.kh@gmail.com>
Date: Mon, 9 Jan 2023 08:45:57 +0100
Subject: [PATCH] Add map_col method

---
 src/headers.rs            |  4 ++++
 src/lib.rs                | 14 +++++++-----
 src/pipeline.rs           | 45 +++++++++++++++++++++++++++++++++++----
 src/pipeline_iterators.rs | 43 +++++++++++++++++++++++++++++++++----
 4 files changed, 93 insertions(+), 13 deletions(-)

diff --git a/src/headers.rs b/src/headers.rs
index 0e6066a..976ec09 100644
--- a/src/headers.rs
+++ b/src/headers.rs
@@ -27,6 +27,10 @@ impl Headers {
 		self.indexes.get(name).and_then(|index| row.get(*index))
 	}
 
+	pub fn get_index(&self, name: &str) -> Option<usize> {
+		self.indexes.get(name).copied()
+	}
+
 	pub fn get_row(&self) -> &Row {
 		&self.row
 	}
diff --git a/src/lib.rs b/src/lib.rs
index 3dcd240..3138a08 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -13,6 +13,7 @@ pub enum Error {
 	Example,
 	Csv(csv::Error),
 	Io(std::io::Error),
+	MissingColumn(String),
 }
 impl From<csv::Error> for Error {
 	fn from(error: csv::Error) -> Error {
@@ -30,10 +31,13 @@ impl From<std::io::Error> for Error {
 fn test_pipeline() {
 	let mut csv_str = String::new();
 	let mut pipeline = PipelineBuilder::from_path("test/Countries.csv")
-		.add_col("Language", |_headers, row| match row.get(1) {
-			Some("Norway") => Ok("Norwegian".to_string()),
-			_ => Ok("Unknown".to_string()),
+		.add_col("Language", |headers, row| {
+			match headers.get_field(row, "Country") {
+				Some("Norway") => Ok("Norwegian".to_string()),
+				_ => Ok("Unknown".to_string()),
+			}
 		})
+		.map_col("Country", |id_str| Ok(id_str.to_uppercase()))
 		.flush(target::StringTarget::new(&mut csv_str))
 		.build();
 
@@ -45,7 +49,7 @@ fn test_pipeline() {
 	assert_eq!(
 		csv_str,
 		"ID,Country,Language\n\
-			1,Norway,Norwegian\n\
-			2,Tuvalu,Unknown\n"
+			1,NORWAY,Norwegian\n\
+			2,TUVALU,Unknown\n"
 	);
 }
diff --git a/src/pipeline.rs b/src/pipeline.rs
index 387cfc6..5611984 100644
--- a/src/pipeline.rs
+++ b/src/pipeline.rs
@@ -1,5 +1,5 @@
 use super::headers::Headers;
-use crate::pipeline_iterators::{AddCol, Flush, Map};
+use crate::pipeline_iterators::{AddCol, Flush, MapCol, MapRow};
 use crate::target::Target;
 use crate::{Error, Row, RowResult};
 use csv::{Reader, ReaderBuilder, StringRecordsIntoIter};
@@ -62,18 +62,55 @@ impl<'a> PipelineBuilder<'a> {
 		self
 	}
 
-	pub fn map<F>(mut self, get_value: F) -> Self
+	/// Maps each row
+	///
+	/// ## Example
+	///
+	/// ```
+	/// use csv_pipeline::PipelineBuilder;
+	///
+	/// PipelineBuilder::from_path("test/Countries.csv")
+	///   .map(|headers, row| {
+	///     Ok(row.into_iter().map(|field| field.to_uppercase()).collect())
+	///   });
+	/// ```
+	pub fn map<F>(mut self, get_row: F) -> Self
 	where
 		F: FnMut(&Headers, Row) -> Result<Row, Error> + 'a,
 	{
-		self.iterator = Box::new(Map {
+		self.iterator = Box::new(MapRow {
 			iterator: self.iterator,
-			f: get_value,
+			f: get_row,
 			headers: self.headers.clone(),
 		});
 		self
 	}
 
+	/// Maps each field of a column
+	///
+	/// ## Example
+	///
+	/// ```
+	/// use csv_pipeline::PipelineBuilder;
+	///
+	/// PipelineBuilder::from_path("test/Countries.csv")
+	///   .map_col("Country", |field| {
+	/// 	  Ok(field.to_uppercase())
+	///   });
+	/// ```
+	pub fn map_col<F>(mut self, col: &str, get_value: F) -> Self
+	where
+		F: FnMut(&str) -> Result<String, Error> + 'a,
+	{
+		self.iterator = Box::new(MapCol {
+			iterator: self.iterator,
+			f: get_value,
+			name: col.to_string(),
+			index: self.headers.get_index(col),
+		});
+		self
+	}
+
 	pub fn flush(mut self, target: impl Target + 'a) -> Self {
 		let flush = Flush::new(self.iterator, target, self.headers.clone());
 		self.iterator = Box::new(flush);
diff --git a/src/pipeline_iterators.rs b/src/pipeline_iterators.rs
index a887aa6..97a4aac 100644
--- a/src/pipeline_iterators.rs
+++ b/src/pipeline_iterators.rs
@@ -15,13 +15,12 @@ where
 	type Item = RowResult;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		let row = match self.iterator.next()? {
+		let mut row = match self.iterator.next()? {
 			Ok(row) => row,
 			Err(e) => return Some(Err(e)),
 		};
 		match (self.f)(&self.headers, &row) {
 			Ok(value) => {
-				let mut row = row;
 				row.push_field(&value);
 				Some(Ok(row))
 			}
@@ -30,12 +29,12 @@ where
 	}
 }
 
-pub struct Map<I, F: FnMut(&Headers, Row) -> Result<Row, Error>> {
+pub struct MapRow<I, F: FnMut(&Headers, Row) -> Result<Row, Error>> {
 	pub iterator: I,
 	pub f: F,
 	pub headers: Headers,
 }
-impl<I, F> Iterator for Map<I, F>
+impl<I, F> Iterator for MapRow<I, F>
 where
 	I: Iterator<Item = RowResult>,
 	F: FnMut(&Headers, Row) -> Result<Row, Error>,
@@ -54,6 +53,42 @@ where
 	}
 }
 
+pub struct MapCol<I, F: FnMut(&str) -> Result<String, Error>> {
+	pub iterator: I,
+	pub f: F,
+	pub name: String,
+	pub index: Option<usize>,
+}
+impl<I, F> Iterator for MapCol<I, F>
+where
+	I: Iterator<Item = RowResult>,
+	F: FnMut(&str) -> Result<String, Error>,
+{
+	type Item = RowResult;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let row = match self.iterator.next()? {
+			Ok(row) => row,
+			Err(e) => return Some(Err(e)),
+		};
+		let mut row_vec: Vec<_> = row.into_iter().collect();
+		let index = match self.index {
+			Some(index) => index,
+			None => return Some(Err(Error::MissingColumn(self.name.clone()))),
+		};
+		let field = match row_vec.get_mut(index) {
+			Some(field) => field,
+			None => return Some(Err(Error::MissingColumn(self.name.clone()))),
+		};
+		let new_value = match (self.f)(field) {
+			Ok(value) => value,
+			Err(e) => return Some(Err(e)),
+		};
+		*field = &new_value;
+		Some(Ok(row_vec.into()))
+	}
+}
+
 pub struct Flush<I, T> {
 	pub iterator: I,
 	pub target: T,
-- 
2.51.2