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}