shell-rs 0.2.6

Rust reimplementation of common coreutils APIs
Documentation
// Copyright (c) 2021 Xu Shaohua <shaohua@biofan.org>. All rights reserved.
// Use of this source is governed by Apache-2.0 License that can be found
// in the LICENSE file.

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

use crate::error::Error;

pub const DEFAULT_LINES: usize = 10;

#[derive(Debug, Default, Clone, Copy, PartialEq)]
pub struct Options {
    /// Get the first n bytes of each file
    pub bytes: Option<usize>,

    /// Get the first n lines instead of the first `10` lines.
    /// If both `bytes` and `lines` are specified, only `lines` is used.
    pub lines: Option<usize>,
}

impl Options {
    pub fn with_bytes(bytes: usize) -> Self {
        Self {
            bytes: Some(bytes),
            lines: None,
        }
    }

    pub fn with_lines(lines: usize) -> Self {
        Self {
            lines: Some(lines),
            bytes: None,
        }
    }
}

/// Output the first part of files.
pub fn head<P: AsRef<Path>>(file: P, options: &Options) -> Result<String, Error> {
    let mut fd = File::open(file)?;

    if options.lines.is_none() && options.bytes.is_some() {
        let mut bytes = options.bytes.unwrap();
        const BUF_SIZE: usize = 4 * 1024;
        let mut buf = [0; BUF_SIZE];
        let mut n_read = 1;
        let mut result = Vec::new();
        while bytes > 0 && n_read > 0 {
            let chunk_size = usize::min(BUF_SIZE, bytes);
            n_read = fd.read(&mut buf[0..chunk_size])?;
            bytes -= n_read;
            result.extend_from_slice(&buf[0..n_read]);
        }
        let result = String::from_utf8(result)?;
        return Ok(result);
    }

    let mut buf = BufReader::new(fd);
    let mut result = String::new();
    let mut lines = options.lines.unwrap_or(DEFAULT_LINES);
    let mut n_read = 1;
    let mut line = String::new();
    while lines > 0 && n_read > 0 {
        n_read = buf.read_line(&mut line)?;
        result += &line;
        lines -= 1;
    }

    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_head() {
        let input_file = "tests/Rust_Wikipedia.html";

        assert!(head(input_file, &Options::default()).is_ok());
        assert!(head(input_file, &Options::with_lines(20)).is_ok());
        assert!(head(input_file, &Options::with_bytes(1024)).is_ok());
    }
}