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}