home / grad_rs / src / csv_reader.rs

csv_reader.rs



use std::fs::File;
use std::io::{BufRead, BufReader, Read, Seek};

pub struct CsvReader {
    separator: char,
    reader: BufReader<File>,
}

impl CsvReader {
    pub fn new(path: &str, separator: char) -> Self {
        let file = File::open(path).unwrap();
        CsvReader {
            separator,
            reader: BufReader::new(file),
        }
    }

    fn get_cell(row: &str, col: usize, separator: char) -> Option<String> {
        let mut count: usize = 0;
        let mut start = 0;
        let mut end = row.len();

        for (i, c) in row.char_indices() {
            if c == separator {
                if count == col {
                    end = i;
                    break;
                }
                count += 1;
                start = i + 1;
            }
        }
        if count == col {
            Some(row[start..end].to_owned())
        } else {
            None
        }
    }

    pub fn vertical_iter(&mut self, col: usize) -> impl Iterator<Item = String> + '_ {
        self.rewind();
        let separator = self.separator;
        self.reader
            .by_ref()
            .lines()
            .filter_map(move |line| line.ok())
            .filter_map(move |line| Self::get_cell(&line, col, separator))
    }

    pub fn get_col_count(&mut self) -> usize {
        self.rewind();
        self.reader
            .by_ref()
            .lines()
            .next()
            .and_then(Result::ok)
            .map(|line| line.chars().filter(|&c| c == self.separator).count() + 1)
            .unwrap_or(0)
    }

    pub fn rewind(&mut self) {
        self.reader.rewind().unwrap();
    }
}