refactor: move context related options to single args

HerringtonDarkholme · HerringtonDarkholme · commit 6b502763482d · 2024-10-23T00:01:56.000-04:00
diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs
@@ -218,6 +218,9 @@ mod test_cli {
     ok("scan --globs '*.js' --globs '*.ts'");
     ok("scan -j 12");
     ok("scan --threads 12");
+    ok("scan -A 12");
+    ok("scan --after 12");
+    ok("scan --context 1");
     error("scan -i --json dir"); // conflict
     error("scan --report-style rich --json dir"); // conflict
     error("scan -r test.yml --inline-rules '{}'"); // conflict
diff --git a/crates/cli/src/run.rs b/crates/cli/src/run.rs
@@ -11,7 +11,7 @@ use crate::config::register_custom_language;
 use crate::lang::SgLang;
 use crate::print::{ColoredPrinter, Diff, Heading, InteractivePrinter, JSONPrinter, Printer};
 use crate::utils::ErrorContext as EC;
-use crate::utils::{filter_file_pattern, InputArgs, MatchUnit, OutputArgs};
+use crate::utils::{filter_file_pattern, ContextArgs, InputArgs, MatchUnit, OutputArgs};
 use crate::utils::{DebugFormat, FileTrace, RunTrace};
 use crate::utils::{Items, PathWorker, StdInWorker, Worker};
 
@@ -108,6 +108,10 @@ pub struct RunArg {
   #[clap(flatten)]
   output: OutputArgs,
 
+  /// context related options
+  #[clap(flatten)]
+  context: ContextArgs,
+
   /// Controls whether to print the file name as heading.
   ///
   /// If heading is used, the file name will be printed as heading before all matches of that file.
@@ -137,11 +141,7 @@ impl RunArg {
 // Every run will include Search or Replace
 // Search or Replace by arguments `pattern` and `rewrite` passed from CLI
 pub fn run_with_pattern(arg: RunArg) -> Result<()> {
-  let context = if arg.output.context != 0 {
-    (arg.output.context, arg.output.context)
-  } else {
-    (arg.output.before, arg.output.after)
-  };
+  let context = arg.context.get();
   if let Some(json) = arg.output.json {
     let printer = JSONPrinter::stdout(json).context(context);
     return run_pattern_with_printer(arg, printer);
@@ -372,6 +372,8 @@ mod test {
         json: None,
         update_all: false,
         tracing: Default::default(),
+      },
+      context: ContextArgs {
         before: 0,
         after: 0,
         context: 0,
diff --git a/crates/cli/src/scan.rs b/crates/cli/src/scan.rs
@@ -18,7 +18,7 @@ use crate::print::{
 };
 use crate::utils::ErrorContext as EC;
 use crate::utils::RuleOverwrite;
-use crate::utils::{filter_file_interactive, InputArgs, OutputArgs, SeverityArg};
+use crate::utils::{filter_file_interactive, ContextArgs, InputArgs, OutputArgs, SeverityArg};
 use crate::utils::{FileTrace, RuleTrace, ScanTrace};
 use crate::utils::{Items, PathWorker, StdInWorker, Worker};
 
@@ -70,15 +70,14 @@ pub struct ScanArg {
   /// output related options
   #[clap(flatten)]
   output: OutputArgs,
+  /// context related options
+  #[clap(flatten)]
+  context: ContextArgs,
 }
 
 pub fn run_with_config(arg: ScanArg) -> Result<()> {
   register_custom_language(arg.config.clone())?;
-  let context = if arg.output.context != 0 {
-    (arg.output.context, arg.output.context)
-  } else {
-    (arg.output.before, arg.output.after)
-  };
+  let context = arg.context.get();
   if let Some(_format) = &arg.format {
     let printer = CloudPrinter::stdout();
     return run_scan(arg, printer);
@@ -390,6 +389,8 @@ rule:
         update_all: false,
         color: ColorArg::Never,
         tracing: Default::default(),
+      },
+      context: ContextArgs {
         before: 0,
         after: 0,
         context: 0,
diff --git a/crates/cli/src/utils/args.rs b/crates/cli/src/utils/args.rs
@@ -153,8 +153,18 @@ pub struct OutputArgs {
   /// tracing information outputs to stderr and does not affect the result of the search.
   #[clap(long, default_value = "nothing", value_name = "LEVEL")]
   pub tracing: Tracing,
+}
+
+impl OutputArgs {
+  // either explicit interactive or implicit update_all
+  pub fn needs_interactive(&self) -> bool {
+    self.interactive || self.update_all
+  }
+}
 
-  // context related options
+/// context related options
+#[derive(Args)]
+pub struct ContextArgs {
   /// Show NUM lines after each match.
   ///
   /// It conflicts with both the -C/--context flag.
@@ -188,10 +198,13 @@ pub struct OutputArgs {
   pub context: u16,
 }
 
-impl OutputArgs {
-  // either explicit interactive or implicit update_all
-  pub fn needs_interactive(&self) -> bool {
-    self.interactive || self.update_all
+impl ContextArgs {
+  pub fn get(&self) -> (u16, u16) {
+    if self.context > 0 {
+      (self.context, self.context)
+    } else {
+      (self.before, self.after)
+    }
   }
 }
 
diff --git a/crates/cli/src/utils/mod.rs b/crates/cli/src/utils/mod.rs
@@ -5,7 +5,7 @@ mod rule_overwrite;
 mod tracing;
 mod worker;
 
-pub use args::{InputArgs, OutputArgs, SeverityArg};
+pub use args::{ContextArgs, InputArgs, OutputArgs, SeverityArg};
 pub use debug_query::DebugFormat;
 pub use error_context::{exit_with_error, ErrorContext};
 pub use rule_overwrite::RuleOverwrite;
