From 6f67097f2fcd58f0ed95593a53ab74bab69dd69d Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Tue, 26 Sep 2023 11:47:28 +0530
Subject: [PATCH] Add `cell` field to JSON output format

---
 crates/ruff_linter/src/message/json.rs       | 12 ++++++++----
 crates/ruff_linter/src/message/json_lines.rs |  4 ++--
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/crates/ruff_linter/src/message/json.rs b/crates/ruff_linter/src/message/json.rs
index e1a2e3e36df988..e2d1dba4a7fc13 100644
--- a/crates/ruff_linter/src/message/json.rs
+++ b/crates/ruff_linter/src/message/json.rs
@@ -19,9 +19,9 @@ impl Emitter for JsonEmitter {
         &mut self,
         writer: &mut dyn Write,
         messages: &[Message],
-        _context: &EmitterContext,
+        context: &EmitterContext,
     ) -> anyhow::Result<()> {
-        serde_json::to_writer_pretty(writer, &ExpandedMessages { messages })?;
+        serde_json::to_writer_pretty(writer, &ExpandedMessages { messages, context })?;
 
         Ok(())
     }
@@ -29,6 +29,7 @@ impl Emitter for JsonEmitter {
 
 struct ExpandedMessages<'a> {
     messages: &'a [Message],
+    context: &'a EmitterContext<'a>,
 }
 
 impl Serialize for ExpandedMessages<'_> {
@@ -39,7 +40,7 @@ impl Serialize for ExpandedMessages<'_> {
         let mut s = serializer.serialize_seq(Some(self.messages.len()))?;
 
         for message in self.messages {
-            let value = message_to_json_value(message);
+            let value = message_to_json_value(message, self.context);
             s.serialize_element(&value)?;
         }
 
@@ -47,8 +48,9 @@ impl Serialize for ExpandedMessages<'_> {
     }
 }
 
-pub(crate) fn message_to_json_value(message: &Message) -> Value {
+pub(crate) fn message_to_json_value(message: &Message, context: &EmitterContext) -> Value {
     let source_code = message.file.to_source_code();
+    let notebook_index = context.notebook_index(message.filename());
 
     let fix = message.fix.as_ref().map(|fix| {
         json!({
@@ -61,12 +63,14 @@ pub(crate) fn message_to_json_value(message: &Message) -> Value {
     let start_location = source_code.source_location(message.start());
     let end_location = source_code.source_location(message.end());
     let noqa_location = source_code.source_location(message.noqa_offset);
+    let cell = notebook_index.map(|index| index.cell(start_location.row.get()).unwrap_or_default());
 
     json!({
         "code": message.kind.rule().noqa_code().to_string(),
         "url": message.kind.rule().url(),
         "message": message.kind.body,
         "fix": fix,
+        "cell": cell,
         "location": start_location,
         "end_location": end_location,
         "filename": message.filename(),
diff --git a/crates/ruff_linter/src/message/json_lines.rs b/crates/ruff_linter/src/message/json_lines.rs
index 360e7ec6a71a76..048cb07fd7f86e 100644
--- a/crates/ruff_linter/src/message/json_lines.rs
+++ b/crates/ruff_linter/src/message/json_lines.rs
@@ -11,11 +11,11 @@ impl Emitter for JsonLinesEmitter {
         &mut self,
         writer: &mut dyn Write,
         messages: &[Message],
-        _context: &EmitterContext,
+        context: &EmitterContext,
     ) -> anyhow::Result<()> {
         let mut w = writer;
         for message in messages {
-            serde_json::to_writer(&mut w, &message_to_json_value(message))?;
+            serde_json::to_writer(&mut w, &message_to_json_value(message, context))?;
             w.write_all(b"\n")?;
         }
         Ok(())