From fbdb1251413ddbabd8b6b03ba0b9cddb347c6030 Mon Sep 17 00:00:00 2001
From: Yuheng Su <gipsyh.icu@gmail.com>
Date: Fri, 6 May 2022 23:53:02 +0800
Subject: [PATCH] Add split number flag in `split row` (#5434)

Signed-off-by: Yuheng Su <gipsyh.icu@gmail.com>
---
 crates/nu-command/src/strings/split/row.rs | 48 ++++++++++++++++------
 1 file changed, 36 insertions(+), 12 deletions(-)

diff --git a/crates/nu-command/src/strings/split/row.rs b/crates/nu-command/src/strings/split/row.rs
index ded65e29a..4961456aa 100644
--- a/crates/nu-command/src/strings/split/row.rs
+++ b/crates/nu-command/src/strings/split/row.rs
@@ -20,6 +20,12 @@ impl Command for SubCommand {
                 SyntaxShape::String,
                 "the character that denotes what separates rows",
             )
+            .named(
+                "number",
+                SyntaxShape::Int,
+                "Split into maximum number of items",
+                Some('n'),
+            )
             .category(Category::Strings)
     }
 
@@ -75,26 +81,44 @@ fn split_row(
 ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
     let name_span = call.head;
     let separator: Spanned<String> = call.req(engine_state, stack, 0)?;
-
+    let max_split: Option<usize> = call.get_flag(engine_state, stack, "number")?;
     input.flat_map(
-        move |x| split_row_helper(&x, &separator, name_span),
+        move |x| split_row_helper(&x, &separator, max_split, name_span),
         engine_state.ctrlc.clone(),
     )
 }
 
-fn split_row_helper(v: &Value, separator: &Spanned<String>, name: Span) -> Vec<Value> {
+fn split_row_helper(
+    v: &Value,
+    separator: &Spanned<String>,
+    max_split: Option<usize>,
+    name: Span,
+) -> Vec<Value> {
     match v.span() {
         Ok(v_span) => {
             if let Ok(s) = v.as_string() {
-                s.split(&separator.item)
-                    .filter_map(|s| {
-                        if s.trim() != "" {
-                            Some(Value::string(s, v_span))
-                        } else {
-                            None
-                        }
-                    })
-                    .collect()
+                match max_split {
+                    Some(max_split) => s
+                        .splitn(max_split, &separator.item)
+                        .filter_map(|s| {
+                            if s.trim() != "" {
+                                Some(Value::string(s, v_span))
+                            } else {
+                                None
+                            }
+                        })
+                        .collect(),
+                    None => s
+                        .split(&separator.item)
+                        .filter_map(|s| {
+                            if s.trim() != "" {
+                                Some(Value::string(s, v_span))
+                            } else {
+                                None
+                            }
+                        })
+                        .collect(),
+                }
             } else {
                 vec![Value::Error {
                     error: ShellError::PipelineMismatch("string".into(), name, v_span),