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