nu_command/strings/str_/
distance.rs

1use nu_cmd_base::input_handler::{CmdArgument, operate};
2use nu_engine::command_prelude::*;
3use nu_protocol::{engine::StateWorkingSet, levenshtein_distance};
4
5#[derive(Clone)]
6pub struct StrDistance;
7
8struct Arguments {
9    compare_string: String,
10    cell_paths: Option<Vec<CellPath>>,
11}
12
13impl CmdArgument for Arguments {
14    fn take_cell_paths(&mut self) -> Option<Vec<CellPath>> {
15        self.cell_paths.take()
16    }
17}
18
19impl Command for StrDistance {
20    fn name(&self) -> &str {
21        "str distance"
22    }
23
24    fn signature(&self) -> Signature {
25        Signature::build("str distance")
26            .input_output_types(vec![
27                (Type::String, Type::Int),
28                (Type::table(), Type::table()),
29                (Type::record(), Type::record()),
30            ])
31            .required(
32                "compare-string",
33                SyntaxShape::String,
34                "The first string to compare.",
35            )
36            .rest(
37                "rest",
38                SyntaxShape::CellPath,
39                "For a data structure input, check strings at the given cell paths, and replace with result.",
40            )
41            .category(Category::Strings)
42    }
43
44    fn description(&self) -> &str {
45        "Compare two strings and return the edit distance/Levenshtein distance."
46    }
47
48    fn search_terms(&self) -> Vec<&str> {
49        vec!["edit", "levenshtein"]
50    }
51
52    fn is_const(&self) -> bool {
53        true
54    }
55
56    fn run(
57        &self,
58        engine_state: &EngineState,
59        stack: &mut Stack,
60        call: &Call,
61        input: PipelineData,
62    ) -> Result<PipelineData, ShellError> {
63        let compare_string: String = call.req(engine_state, stack, 0)?;
64        let cell_paths: Vec<CellPath> = call.rest(engine_state, stack, 1)?;
65        let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
66        let args = Arguments {
67            compare_string,
68            cell_paths,
69        };
70        operate(action, args, input, call.head, engine_state.signals())
71    }
72
73    fn run_const(
74        &self,
75        working_set: &StateWorkingSet,
76        call: &Call,
77        input: PipelineData,
78    ) -> Result<PipelineData, ShellError> {
79        let compare_string: String = call.req_const(working_set, 0)?;
80        let cell_paths: Vec<CellPath> = call.rest_const(working_set, 1)?;
81        let cell_paths = (!cell_paths.is_empty()).then_some(cell_paths);
82        let args = Arguments {
83            compare_string,
84            cell_paths,
85        };
86        operate(
87            action,
88            args,
89            input,
90            call.head,
91            working_set.permanent().signals(),
92        )
93    }
94
95    fn examples(&self) -> Vec<Example<'_>> {
96        vec![
97            Example {
98                description: "get the edit distance between two strings",
99                example: "'nushell' | str distance 'nutshell'",
100                result: Some(Value::test_int(1)),
101            },
102            Example {
103                description: "Compute edit distance between strings in table and another string, using cell paths",
104                example: "[{a: 'nutshell' b: 'numetal'}] | str distance 'nushell' 'a' 'b'",
105                result: Some(Value::test_list(vec![Value::test_record(record! {
106                    "a" => Value::test_int(1),
107                    "b" => Value::test_int(4),
108                })])),
109            },
110            Example {
111                description: "Compute edit distance between strings in record and another string, using cell paths",
112                example: "{a: 'nutshell' b: 'numetal'} | str distance 'nushell' a b",
113                result: Some(Value::test_record(record! {
114                    "a" => Value::test_int(1),
115                    "b" => Value::test_int(4),
116                })),
117            },
118        ]
119    }
120}
121
122fn action(input: &Value, args: &Arguments, head: Span) -> Value {
123    let compare_string = &args.compare_string;
124    match input {
125        Value::String { val, .. } => {
126            let distance = levenshtein_distance(val, compare_string);
127            Value::int(distance as i64, head)
128        }
129        Value::Error { .. } => input.clone(),
130        _ => Value::error(
131            ShellError::OnlySupportsThisInputType {
132                exp_input_type: "string".into(),
133                wrong_type: input.get_type().to_string(),
134                dst_span: head,
135                src_span: input.span(),
136            },
137            head,
138        ),
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn test_examples() {
148        use crate::test_examples;
149
150        test_examples(StrDistance {})
151    }
152}