feat(scan): support context, before/after flags

connor15mcc · HerringtonDarkholme · commit 36ad9622debc · 2024-10-22T15:58:45.000-04:00
fix ast-grep#1549 Add support for `--before`, `--after` | `--context` CLI flags for `sg scan`, matching the existing flags for `sg fix` When developing complex rules (often those that necessitate a yaml rule file to begin with), this context can be instrumental in evaluating the quality of a fix.
diff --git a/crates/cli/src/run.rs b/crates/cli/src/run.rs
@@ -116,39 +116,6 @@ pub struct RunArg {
   /// and to disable heading when piping to another program or redirected to files.
   #[clap(long, default_value = "auto", value_name = "WHEN")]
   heading: Heading,
-
-  // context related options
-  /// Show NUM lines after each match.
-  ///
-  /// It conflicts with both the -C/--context flag.
-  #[clap(
-    short = 'A',
-    long,
-    default_value = "0",
-    conflicts_with = "context",
-    value_name = "NUM"
-  )]
-  after: u16,
-
-  /// Show NUM lines before each match.
-  ///
-  /// It conflicts with both the -C/--context flag.
-  #[clap(
-    short = 'B',
-    long,
-    default_value = "0",
-    conflicts_with = "context",
-    value_name = "NUM"
-  )]
-  before: u16,
-
-  /// Show NUM lines around each match.
-  ///
-  /// This is equivalent to providing both the
-  /// -B/--before and -A/--after flags with the same value.
-  /// It conflicts with both the -B/--before and -A/--after flags.
-  #[clap(short = 'C', long, default_value = "0", value_name = "NUM")]
-  context: u16,
 }
 
 impl RunArg {
@@ -170,10 +137,10 @@ 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.context != 0 {
-    (arg.context, arg.context)
+  let context = if arg.output.context != 0 {
+    (arg.output.context, arg.output.context)
   } else {
-    (arg.before, arg.after)
+    (arg.output.before, arg.output.after)
   };
   if let Some(json) = arg.output.json {
     let printer = JSONPrinter::stdout(json).context(context);
@@ -405,10 +372,10 @@ mod test {
         json: None,
         update_all: false,
         tracing: Default::default(),
+        before: 0,
+        after: 0,
+        context: 0,
       },
-      before: 0,
-      after: 0,
-      context: 0,
     }
   }
 
diff --git a/crates/cli/src/scan.rs b/crates/cli/src/scan.rs
@@ -74,6 +74,11 @@ pub struct ScanArg {
 
 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)
+  };
   if let Some(_format) = &arg.format {
     let printer = CloudPrinter::stdout();
     return run_scan(arg, printer);
@@ -82,7 +87,9 @@ pub fn run_with_config(arg: ScanArg) -> Result<()> {
     let printer = JSONPrinter::stdout(json);
     return run_scan(arg, printer);
   }
-  let printer = ColoredPrinter::stdout(arg.output.color).style(arg.report_style);
+  let printer = ColoredPrinter::stdout(arg.output.color)
+    .style(arg.report_style)
+    .context(context);
   let interactive = arg.output.needs_interactive();
   if interactive {
     let from_stdin = arg.input.stdin;
@@ -383,6 +390,9 @@ rule:
         update_all: false,
         color: ColorArg::Never,
         tracing: Default::default(),
+        before: 0,
+        after: 0,
+        context: 0,
       },
       format: None,
     }
diff --git a/crates/cli/src/utils/args.rs b/crates/cli/src/utils/args.rs
@@ -153,6 +153,39 @@ 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,
+
+  // context related options
+  /// Show NUM lines after each match.
+  ///
+  /// It conflicts with both the -C/--context flag.
+  #[clap(
+    short = 'A',
+    long,
+    default_value = "0",
+    conflicts_with = "context",
+    value_name = "NUM"
+  )]
+  pub after: u16,
+
+  /// Show NUM lines before each match.
+  ///
+  /// It conflicts with both the -C/--context flag.
+  #[clap(
+    short = 'B',
+    long,
+    default_value = "0",
+    conflicts_with = "context",
+    value_name = "NUM"
+  )]
+  pub before: u16,
+
+  /// Show NUM lines around each match.
+  ///
+  /// This is equivalent to providing both the
+  /// -B/--before and -A/--after flags with the same value.
+  /// It conflicts with both the -B/--before and -A/--after flags.
+  #[clap(short = 'C', long, default_value = "0", value_name = "NUM")]
+  pub context: u16,
 }
 
 impl OutputArgs {
