1use indexmap::IndexMap;
2use nu_engine::{ClosureEval, command_prelude::*};
3use nu_protocol::engine::Closure;
4
5#[derive(Clone)]
6pub struct Rename;
7
8impl Command for Rename {
9 fn name(&self) -> &str {
10 "rename"
11 }
12
13 fn signature(&self) -> Signature {
14 Signature::build("rename")
15 .input_output_types(vec![
16 (Type::record(), Type::record()),
17 (Type::table(), Type::table()),
18 ])
19 .named(
20 "column",
21 SyntaxShape::Record(vec![]),
22 "column name to be changed",
23 Some('c'),
24 )
25 .named(
26 "block",
27 SyntaxShape::Closure(Some(vec![SyntaxShape::Any])),
28 "A closure to apply changes on each column",
29 Some('b'),
30 )
31 .rest(
32 "rest",
33 SyntaxShape::String,
34 "The new names for the columns.",
35 )
36 .category(Category::Filters)
37 }
38
39 fn description(&self) -> &str {
40 "Creates a new table with columns renamed."
41 }
42
43 fn run(
44 &self,
45 engine_state: &EngineState,
46 stack: &mut Stack,
47 call: &Call,
48 input: PipelineData,
49 ) -> Result<PipelineData, ShellError> {
50 rename(engine_state, stack, call, input)
51 }
52
53 fn examples(&self) -> Vec<Example<'_>> {
54 vec![
55 Example {
56 description: "Rename a column",
57 example: "[[a, b]; [1, 2]] | rename my_column",
58 result: Some(Value::test_list(vec![Value::test_record(record! {
59 "my_column" => Value::test_int(1),
60 "b" => Value::test_int(2),
61 })])),
62 },
63 Example {
64 description: "Rename many columns",
65 example: "[[a, b, c]; [1, 2, 3]] | rename eggs ham bacon",
66 result: Some(Value::test_list(vec![Value::test_record(record! {
67 "eggs" => Value::test_int(1),
68 "ham" => Value::test_int(2),
69 "bacon" => Value::test_int(3),
70 })])),
71 },
72 Example {
73 description: "Rename a specific column",
74 example: "[[a, b, c]; [1, 2, 3]] | rename --column { a: ham }",
75 result: Some(Value::test_list(vec![Value::test_record(record! {
76 "ham" => Value::test_int(1),
77 "b" => Value::test_int(2),
78 "c" => Value::test_int(3),
79 })])),
80 },
81 Example {
82 description: "Rename the fields of a record",
83 example: "{a: 1 b: 2} | rename x y",
84 result: Some(Value::test_record(record! {
85 "x" => Value::test_int(1),
86 "y" => Value::test_int(2),
87 })),
88 },
89 Example {
90 description: "Rename fields based on a given closure",
91 example: "{abc: 1, bbc: 2} | rename --block {str replace --all 'b' 'z'}",
92 result: Some(Value::test_record(record! {
93 "azc" => Value::test_int(1),
94 "zzc" => Value::test_int(2),
95 })),
96 },
97 ]
98 }
99}
100
101fn rename(
102 engine_state: &EngineState,
103 stack: &mut Stack,
104 call: &Call,
105 input: PipelineData,
106) -> Result<PipelineData, ShellError> {
107 let head = call.head;
108 let columns: Vec<String> = call.rest(engine_state, stack, 0)?;
109 let specified_column: Option<Record> = call.get_flag(engine_state, stack, "column")?;
110 let specified_column: Option<IndexMap<String, String>> = match specified_column {
112 Some(query) => {
113 let mut columns = IndexMap::new();
114 for (col, val) in query {
115 let val_span = val.span();
116 match val {
117 Value::String { val, .. } => {
118 columns.insert(col, val);
119 }
120 _ => {
121 return Err(ShellError::TypeMismatch {
122 err_message: "new column name must be a string".to_owned(),
123 span: val_span,
124 });
125 }
126 }
127 }
128 if columns.is_empty() {
129 return Err(ShellError::TypeMismatch {
130 err_message: "The column info cannot be empty".to_owned(),
131 span: call.head,
132 });
133 }
134 Some(columns)
135 }
136 None => None,
137 };
138 let closure: Option<Closure> = call.get_flag(engine_state, stack, "block")?;
139
140 let mut closure = closure.map(|closure| ClosureEval::new(engine_state, stack, closure));
141
142 let metadata = input.metadata();
143 input
144 .map(
145 move |item| {
146 let span = item.span();
147 match item {
148 Value::Record { val: record, .. } => {
149 let record = if let Some(closure) = &mut closure {
150 record
151 .into_owned()
152 .into_iter()
153 .map(|(col, val)| {
154 let col = Value::string(col, span);
155 let data = closure.run_with_value(col)?;
156 let col = data.collect_string_strict(span)?.0;
157 Ok((col, val))
158 })
159 .collect::<Result<Record, _>>()
160 } else {
161 match &specified_column {
162 Some(columns) => {
163 let mut renamed = 0;
166 let record = record
167 .into_owned()
168 .into_iter()
169 .map(|(col, val)| {
170 let col = if let Some(col) = columns.get(&col) {
171 renamed += 1;
172 col.clone()
173 } else {
174 col
175 };
176
177 (col, val)
178 })
179 .collect::<Record>();
180
181 let missing_column = if renamed < columns.len() {
182 columns.iter().find_map(|(col, new_col)| {
183 (!record.contains(new_col)).then_some(col)
184 })
185 } else {
186 None
187 };
188
189 if let Some(missing) = missing_column {
190 Err(ShellError::UnsupportedInput {
191 msg: format!(
192 "The column '{missing}' does not exist in the input"
193 ),
194 input: "value originated from here".into(),
195 msg_span: head,
196 input_span: span,
197 })
198 } else {
199 Ok(record)
200 }
201 }
202 None => Ok(record
203 .into_owned()
204 .into_iter()
205 .enumerate()
206 .map(|(i, (col, val))| {
207 (columns.get(i).cloned().unwrap_or(col), val)
208 })
209 .collect()),
210 }
211 };
212
213 match record {
214 Ok(record) => Value::record(record, span),
215 Err(err) => Value::error(err, span),
216 }
217 }
218 Value::Error { .. } => item,
220 other => Value::error(
221 ShellError::OnlySupportsThisInputType {
222 exp_input_type: "record".into(),
223 wrong_type: other.get_type().to_string(),
224 dst_span: head,
225 src_span: other.span(),
226 },
227 head,
228 ),
229 }
230 },
231 engine_state.signals(),
232 )
233 .map(|data| data.set_metadata(metadata))
234}
235
236#[cfg(test)]
237mod test {
238 use super::*;
239
240 #[test]
241 fn test_examples() {
242 use crate::test_examples;
243
244 test_examples(Rename {})
245 }
246}