Add comments option to playground (#6911)

Co-authored-by: Micha Reiser <micha@reiser.io>

Cargo.lock

@@ -2532,6 +2532,7 @@ dependencies = [
  "log",
  "ruff",
  "ruff_diagnostics",
+ "ruff_formatter",
  "ruff_python_ast",
  "ruff_python_codegen",
  "ruff_python_formatter",

crates/ruff_python_formatter/src/lib.rs

@@ -14,7 +14,7 @@ use ruff_text_size::TextLen;
 use crate::comments::{
     dangling_comments, leading_comments, trailing_comments, Comments, SourceComment,
 };
-use crate::context::PyFormatContext;
+pub use crate::context::PyFormatContext;
 pub use crate::options::{MagicTrailingComma, PyFormatOptions, QuoteStyle};
 use crate::verbatim::suppressed_node;
 
@@ -167,6 +167,17 @@ pub fn format_node<'a>(
     Ok(formatted)
 }
 
+/// Public function for generating a printable string of the debug comments.
+pub fn pretty_comments(formatted: &Formatted<PyFormatContext>, source: &str) -> String {
+    let comments = formatted.context().comments();
+
+    // When comments are empty we'd display an empty map '{}'
+    std::format!(
+        "{comments:#?}",
+        comments = comments.debug(SourceCode::new(source))
+    )
+}
+
 pub(crate) struct NotYetImplementedCustomText<'a> {
     text: &'static str,
     node: AnyNodeRef<'a>,

crates/ruff_wasm/Cargo.toml

@@ -22,6 +22,7 @@ ruff = { path = "../ruff" }
 ruff_diagnostics = { path = "../ruff_diagnostics" }
 ruff_python_ast = { path = "../ruff_python_ast" }
 ruff_python_codegen = { path = "../ruff_python_codegen" }
+ruff_formatter = { path = "../ruff_formatter" }
 ruff_python_formatter = { path = "../ruff_python_formatter" }
 ruff_python_index = { path = "../ruff_python_index" }
 ruff_python_parser = { path = "../ruff_python_parser" }

crates/ruff_wasm/src/lib.rs

@@ -10,10 +10,11 @@ use ruff::linter::{check_path, LinterResult};
 use ruff::registry::AsRule;
 use ruff::settings::types::PythonVersion;
 use ruff::settings::{defaults, flags, Settings};
-use ruff_python_ast::PySourceType;
+use ruff_formatter::{FormatResult, Formatted};
+use ruff_python_ast::{Mod, PySourceType};
 use ruff_python_codegen::Stylist;
-use ruff_python_formatter::{format_module, format_node, PyFormatOptions};
+use ruff_python_formatter::{format_node, pretty_comments, PyFormatContext, PyFormatOptions};
-use ruff_python_index::{CommentRangesBuilder, Indexer};
+use ruff_python_index::{CommentRanges, CommentRangesBuilder, Indexer};
 use ruff_python_parser::lexer::LexResult;
 use ruff_python_parser::AsMode;
 use ruff_python_parser::{parse_tokens, Mode};
@@ -230,32 +231,28 @@ impl Workspace {
     }
 
     pub fn format(&self, contents: &str) -> Result<String, Error> {
-        // TODO(konstin): Add an options for py/pyi to the UI (1/2)
+        let parsed = ParsedModule::from_source(contents)?;
-        let options = PyFormatOptions::from_source_type(PySourceType::default());
+        let formatted = parsed.format().map_err(into_error)?;
-        let printed = format_module(contents, options).map_err(into_error)?;
+        let printed = formatted.print().map_err(into_error)?;
 
         Ok(printed.into_code())
     }
 
     pub fn format_ir(&self, contents: &str) -> Result<String, Error> {
-        let tokens: Vec<_> = ruff_python_parser::lexer::lex(contents, Mode::Module).collect();
+        let parsed = ParsedModule::from_source(contents)?;
-        let mut comment_ranges = CommentRangesBuilder::default();
+        let formatted = parsed.format().map_err(into_error)?;
-
-        for (token, range) in tokens.iter().flatten() {
-            comment_ranges.visit_token(token, *range);
-        }
-
-        let comment_ranges = comment_ranges.finish();
-        let module = parse_tokens(tokens, Mode::Module, ".").map_err(into_error)?;
-
-        // TODO(konstin): Add an options for py/pyi to the UI (2/2)
-        let options = PyFormatOptions::from_source_type(PySourceType::default());
-        let formatted =
-            format_node(&module, &comment_ranges, contents, options).map_err(into_error)?;
 
         Ok(format!("{formatted}"))
     }
 
+    pub fn comments(&self, contents: &str) -> Result<String, Error> {
+        let parsed = ParsedModule::from_source(contents)?;
+        let formatted = parsed.format().map_err(into_error)?;
+        let comments = pretty_comments(&formatted, contents);
+
+        Ok(comments)
+    }
+
     /// Parses the content and returns its AST
     pub fn parse(&self, contents: &str) -> Result<String, Error> {
         let parsed = ruff_python_parser::parse(contents, Mode::Module, ".").map_err(into_error)?;
@@ -273,3 +270,40 @@ impl Workspace {
 pub(crate) fn into_error<E: std::fmt::Display>(err: E) -> Error {
     Error::new(&err.to_string())
 }
+
+struct ParsedModule<'a> {
+    source_code: &'a str,
+    module: Mod,
+    comment_ranges: CommentRanges,
+}
+
+impl<'a> ParsedModule<'a> {
+    fn from_source(source: &'a str) -> Result<Self, Error> {
+        let tokens: Vec<_> = ruff_python_parser::lexer::lex(source, Mode::Module).collect();
+        let mut comment_ranges = CommentRangesBuilder::default();
+
+        for (token, range) in tokens.iter().flatten() {
+            comment_ranges.visit_token(token, *range);
+        }
+        let comment_ranges = comment_ranges.finish();
+        let module = parse_tokens(tokens, Mode::Module, ".").map_err(into_error)?;
+
+        Ok(Self {
+            source_code: source,
+            comment_ranges,
+            module,
+        })
+    }
+
+    fn format(&self) -> FormatResult<Formatted<PyFormatContext>> {
+        // TODO(konstin): Add an options for py/pyi to the UI (2/2)
+        let options = PyFormatOptions::from_source_type(PySourceType::default());
+
+        format_node(
+            &self.module,
+            &self.comment_ranges,
+            self.source_code,
+            options,
+        )
+    }
+}

playground/src/Editor/Editor.tsx

@@ -139,6 +139,13 @@ export default function Editor() {
             };
             break;
 
+          case "Comments":
+            secondary = {
+              status: "ok",
+              content: workspace.comments(pythonSource),
+            };
+            break;
+
           case "Tokens":
             secondary = {
               status: "ok",

playground/src/Editor/Icons.tsx

@@ -124,3 +124,19 @@ export function FormatterIRIcon() {
     </svg>
   );
 }
+export function CommentsIcon() {
+  return (
+    <svg
+      viewBox="0 0 24 24"
+      xmlns="http://www.w3.org/2000/svg"
+      width={24}
+      height={24}
+    >
+      <path
+        d="M7.7,18.3H19.4a2.1,2.1,0,0,0,2.1-2.1V4.6a2.1,2.1,0,0,0-2.1-2.1H4.6A2.1,2.1,0,0,0,2.5,4.6V21.5Z"
+        stroke="#ffffff"
+        fill="none"
+      ></path>
+    </svg>
+  );
+}

playground/src/Editor/SecondaryPanel.tsx

@@ -6,6 +6,7 @@ export enum SecondaryTool {
   "AST" = "AST",
   "Tokens" = "Tokens",
   "FIR" = "FIR",
+  "Comments" = "Comments",
 }
 
 export type SecondaryPanelResult =
@@ -64,6 +65,10 @@ function Content({
           case "FIR":
             language = "fir";
             break;
+
+          case "Comments":
+            language = "Comments";
+            break;
         }
 
         return (

playground/src/Editor/SecondarySideBar.tsx

@@ -4,6 +4,7 @@ import {
   FormatterIRIcon,
   StructureIcon,
   TokensIcon,
+  CommentsIcon,
 } from "./Icons";
 import { SecondaryTool } from "./SecondaryPanel";
 
@@ -53,6 +54,15 @@ export default function SecondarySideBar({
       >
         <FormatterIRIcon />
       </SideBarEntry>
+
+      <SideBarEntry
+        title="Formatter comments"
+        position={"right"}
+        selected={selected === SecondaryTool.Comments}
+        onClick={() => onSelected(SecondaryTool.Comments)}
+      >
+        <CommentsIcon />
+      </SideBarEntry>
     </SideBar>
   );
 }

playground/src/Editor/setupMonaco.tsx

@@ -30,6 +30,7 @@ export function setupMonaco(monaco: Monaco) {
   defineFirLanguage(monaco);
   defineRustPythonTokensLanguage(monaco);
   defineRustPythonAstLanguage(monaco);
+  defineCommentsLanguage(monaco);
 }
 
 function defineAyuThemes(monaco: Monaco) {
@@ -620,6 +621,62 @@ function defineRustPythonAstLanguage(monaco: Monaco) {
   });
 }
 
+// Modeled after 'RustPythonAst'
+function defineCommentsLanguage(monaco: Monaco) {
+  monaco.languages.register({
+    id: "Comments",
+  });
+
+  monaco.languages.setMonarchTokensProvider("Comments", {
+    keywords: ["None", "Err"],
+    tokenizer: {
+      root: [
+        [
+          /[a-zA-Z_$][\w$]*/,
+          {
+            cases: {
+              "@keywords": "keyword",
+              "@default": "identifier",
+            },
+          },
+        ],
+
+        // Whitespace
+        [/[ \t\r\n]+/, "white"],
+
+        // Strings
+        [/"/, { token: "string.quote", bracket: "@open", next: "@string" }],
+
+        [/\d+/, "number"],
+
+        [/[{}()[\]]/, "@brackets"],
+      ],
+      string: [
+        [/[^\\"]+/, "string"],
+        [/\\[\\"]/, "string.escape"],
+        [/"/, { token: "string.quote", bracket: "@close", next: "@pop" }],
+      ],
+    },
+    brackets: [
+      {
+        open: "(",
+        close: ")",
+        token: "delimiter.parenthesis",
+      },
+      {
+        open: "{",
+        close: "}",
+        token: "delimiter.curly",
+      },
+      {
+        open: "[",
+        close: "]",
+        token: "delimiter.bracket",
+      },
+    ],
+  });
+}
+
 function defineRustPythonTokensLanguage(monaco: Monaco) {
   monaco.languages.register({
     id: "RustPythonTokens",