nu_command/strings/str_/
length.rs

1use crate::{grapheme_flags, grapheme_flags_const};
2use nu_cmd_base::input_handler::{CmdArgument, operate};
3use nu_engine::command_prelude::*;
4
5use unicode_segmentation::UnicodeSegmentation;
6
7struct Arguments {
8    cell_paths: Option<Vec<CellPath>>,
9    graphemes: bool,
10    chars: bool,
11}
12
13impl CmdArgument for Arguments {
14    fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
15        self.cell_paths.take()
16    }
17}
18
19#[derive(Clone)]
20pub struct StrLength;
21
22impl Command for StrLength {
23    fn name(&self) -> &str {
24        "str length"
25    }
26
27    fn signature(&self) -> Signature {
28        Signature::build("str length")
29            .input_output_types(vec![
30                (Type::String, Type::Int),
31                (Type::List(Box::new(Type::String)), Type::List(Box::new(Type::Int))),
32                (Type::table(), Type::table()),
33                (Type::record(), Type::record()),
34            ])
35            .allow_variants_without_examples(true)
36            .switch(
37                "grapheme-clusters",
38                "count length in grapheme clusters (all visible chars have length 1)",
39                Some('g'),
40            )
41            .switch(
42                "utf-8-bytes",
43                "count length in UTF-8 bytes (default; all non-ASCII chars have length 2+)",
44                Some('b'),
45            )
46            .switch(
47                "chars",
48                "count length in chars",
49                Some('c'),
50            )
51            .rest(
52                "rest",
53                SyntaxShape::CellPath,
54                "For a data structure input, replace strings at the given cell paths with their length.",
55            )
56            .category(Category::Strings)
57    }
58
59    fn description(&self) -> &str {
60        "Output the length of any strings in the pipeline."
61    }
62
63    fn search_terms(&self) -> Vec<&str> {
64        vec!["size", "count"]
65    }
66
67    fn is_const(&self) -> bool {
68        true
69    }
70
71    fn run(
72        &self,
73        engine_state: &EngineState,
74        stack: &mut Stack,
75        call: &Call,
76        input: PipelineData,
77    ) -> Result<PipelineData, ShellError> {
78        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 0)?;
79        let chars = call.has_flag(engine_state, stack, "chars")?;
80        run(
81            cell_paths,
82            engine_state,
83            call,
84            input,
85            grapheme_flags(engine_state, stack, call)?,
86            chars,
87        )
88    }
89
90    fn run_const(
91        &self,
92        working_set: &StateWorkingSet,
93        call: &Call,
94        input: PipelineData,
95    ) -> Result<PipelineData, ShellError> {
96        let cell_paths: Vec<CellPath> = call.rest_const(working_set, 0)?;
97        let chars = call.has_flag_const(working_set, "chars")?;
98        run(
99            cell_paths,
100            working_set.permanent(),
101            call,
102            input,
103            grapheme_flags_const(working_set, call)?,
104            chars,
105        )
106    }
107
108    fn examples(&self) -> Vec<Example<'_>> {
109        vec![
110            Example {
111                description: "Return the lengths of a string in bytes",
112                example: "'hello' | str length",
113                result: Some(Value::test_int(5)),
114            },
115            Example {
116                description: "Count length of a string in grapheme clusters",
117                example: "'🇯🇵ほげ ふが ぴよ' | str length  --grapheme-clusters",
118                result: Some(Value::test_int(9)),
119            },
120            Example {
121                description: "Return the lengths of multiple strings in bytes",
122                example: "['hi' 'there'] | str length",
123                result: Some(Value::list(
124                    vec![Value::test_int(2), Value::test_int(5)],
125                    Span::test_data(),
126                )),
127            },
128            Example {
129                description: "Return the lengths of a string in chars",
130                example: "'hällo' | str length --chars",
131                result: Some(Value::test_int(5)),
132            },
133        ]
134    }
135}
136
137fn run(
138    cell_paths: Vec<CellPath>,
139    engine_state: &EngineState,
140    call: &Call,
141    input: PipelineData,
142    graphemes: bool,
143    chars: bool,
144) -> Result<PipelineData, ShellError> {
145    let args = Arguments {
146        cell_paths: (!cell_paths.is_empty()).then_some(cell_paths),
147        graphemes,
148        chars,
149    };
150    operate(action, args, input, call.head, engine_state.signals())
151}
152
153fn action(input: &Value, arg: &Arguments, head: Span) -> Value {
154    match input {
155        Value::String { val, .. } => Value::int(
156            if arg.graphemes {
157                val.graphemes(true).count()
158            } else if arg.chars {
159                val.chars().count()
160            } else {
161                val.len()
162            } as i64,
163            head,
164        ),
165        Value::Error { .. } => input.clone(),
166        _ => Value::error(
167            ShellError::OnlySupportsThisInputType {
168                exp_input_type: "string".into(),
169                wrong_type: input.get_type().to_string(),
170                dst_span: head,
171                src_span: input.span(),
172            },
173            head,
174        ),
175    }
176}
177
178#[cfg(test)]
179mod test {
180    use super::*;
181
182    #[test]
183    fn use_utf8_bytes() {
184        let word = Value::string(String::from("🇯🇵ほげ ふが ぴよ"), Span::test_data());
185
186        let options = Arguments {
187            cell_paths: None,
188            graphemes: false,
189            chars: false,
190        };
191
192        let actual = action(&word, &options, Span::test_data());
193        assert_eq!(actual, Value::test_int(28));
194    }
195
196    #[test]
197    fn test_examples() {
198        use crate::test_examples;
199
200        test_examples(StrLength {})
201    }
202}