1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

extern crate unicode_segmentation;

use crate::prelude::*;
use indexmap::indexmap;
use nu_engine::WholeStreamCommand;
use nu_errors::ShellError;
use nu_protocol::{Signature, TaggedDictBuilder, UntaggedValue, Value};
use unicode_segmentation::UnicodeSegmentation;

pub struct Size;

impl WholeStreamCommand for Size {
    fn name(&self) -> &str {
        "size"
    }

    fn signature(&self) -> Signature {
        Signature::build("size")
    }

    fn usage(&self) -> &str {
        "Gather word count statistics on the text."
    }

    fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
        size(args)
    }

    fn examples(&self) -> Vec<Example> {
        vec![
            Example {
                description: "Count the number of words in a string",
                example: r#"echo "There are seven words in this sentence" | size"#,
                result: Some(vec![UntaggedValue::row(indexmap! {
                        "lines".to_string() => UntaggedValue::int(0).into(),
                        "words".to_string() => UntaggedValue::int(7).into(),
                        "chars".to_string() => UntaggedValue::int(38).into(),
                        "bytes".to_string() => UntaggedValue::int(38).into(),
                })
                .into()]),
            },
            Example {
                description: "Counts Unicode characters correctly in a string",
                example: r#"echo "Amélie Amelie" | size"#,
                result: Some(vec![UntaggedValue::row(indexmap! {
                        "lines".to_string() => UntaggedValue::int(0).into(),
                        "words".to_string() => UntaggedValue::int(2).into(),
                        "chars".to_string() => UntaggedValue::int(13).into(),
                        "bytes".to_string() => UntaggedValue::int(15).into(),
                })
                .into()]),
            },
        ]
    }
}

fn size(args: CommandArgs) -> Result<OutputStream, ShellError> {
    let input = args.input;
    let tag = args.call_info.name_tag;
    let name_span = tag.span;

    Ok(input
        .map(move |v| {
            if let Ok(s) = v.as_string() {
                Ok(count(&s, &v.tag))
            } else {
                Err(ShellError::labeled_error_with_secondary(
                    "Expected a string from pipeline",
                    "requires string input",
                    name_span,
                    "value originates from here",
                    v.tag.span,
                ))
            }
        })
        .into_input_stream())
}

fn count(contents: &str, tag: impl Into<Tag>) -> Value {
    let mut lines: i64 = 0;
    let mut words: i64 = 0;
    let mut chars: i64 = 0;
    let bytes = contents.len() as i64;
    let mut end_of_word = true;

    for c in UnicodeSegmentation::graphemes(contents, true) {
        chars += 1;

        match c {
            "\n" => {
                lines += 1;
                end_of_word = true;
            }
            " " => end_of_word = true,
            _ => {
                if end_of_word {
                    words += 1;
                }
                end_of_word = false;
            }
        }
    }

    let mut dict = TaggedDictBuilder::new(tag);
    //TODO: add back in name when we have it in the tag
    //dict.insert("name", value::string(name));
    dict.insert_untagged("lines", UntaggedValue::int(lines));
    dict.insert_untagged("words", UntaggedValue::int(words));
    dict.insert_untagged("chars", UntaggedValue::int(chars));
    dict.insert_untagged("bytes", UntaggedValue::int(bytes));

    dict.into_value()
}

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

    #[test]
    fn examples_work_as_expected() -> Result<(), ShellError> {
        use crate::examples::test as test_examples;

        test_examples(Size {})
    }
}