shell_rs/
head.rs

1// Copyright (c) 2021 Xu Shaohua <shaohua@biofan.org>. All rights reserved.
2// Use of this source is governed by Apache-2.0 License that can be found
3// in the LICENSE file.
4
5use std::fs::File;
6use std::io::{BufRead, BufReader, Read};
7use std::path::Path;
8
9use crate::error::Error;
10
11pub const DEFAULT_LINES: usize = 10;
12
13#[derive(Debug, Default, Clone, Copy, PartialEq)]
14pub struct Options {
15    /// Get the first n bytes of each file
16    pub bytes: Option<usize>,
17
18    /// Get the first n lines instead of the first `10` lines.
19    /// If both `bytes` and `lines` are specified, only `lines` is used.
20    pub lines: Option<usize>,
21}
22
23impl Options {
24    pub fn with_bytes(bytes: usize) -> Self {
25        Self {
26            bytes: Some(bytes),
27            lines: None,
28        }
29    }
30
31    pub fn with_lines(lines: usize) -> Self {
32        Self {
33            lines: Some(lines),
34            bytes: None,
35        }
36    }
37}
38
39/// Output the first part of files.
40pub fn head<P: AsRef<Path>>(file: P, options: &Options) -> Result<String, Error> {
41    let mut fd = File::open(file)?;
42
43    if options.lines.is_none() && options.bytes.is_some() {
44        let mut bytes = options.bytes.unwrap();
45        const BUF_SIZE: usize = 4 * 1024;
46        let mut buf = [0; BUF_SIZE];
47        let mut n_read = 1;
48        let mut result = Vec::new();
49        while bytes > 0 && n_read > 0 {
50            let chunk_size = usize::min(BUF_SIZE, bytes);
51            n_read = fd.read(&mut buf[0..chunk_size])?;
52            bytes -= n_read;
53            result.extend_from_slice(&buf[0..n_read]);
54        }
55        let result = String::from_utf8(result)?;
56        return Ok(result);
57    }
58
59    let mut buf = BufReader::new(fd);
60    let mut result = String::new();
61    let mut lines = options.lines.unwrap_or(DEFAULT_LINES);
62    let mut n_read = 1;
63    let mut line = String::new();
64    while lines > 0 && n_read > 0 {
65        n_read = buf.read_line(&mut line)?;
66        result += &line;
67        lines -= 1;
68    }
69
70    Ok(result)
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn test_head() {
79        let input_file = "tests/Rust_Wikipedia.html";
80
81        assert!(head(input_file, &Options::default()).is_ok());
82        assert!(head(input_file, &Options::with_lines(20)).is_ok());
83        assert!(head(input_file, &Options::with_bytes(1024)).is_ok());
84    }
85}
shell_rs/head.rs

shell_rs/
head.rs
