range without decorators

Summary: in screenshot 1 you see that the range includes decorators which isn't great. this diff takes those off of classes/functions

Reviewed By: stroxler

Differential Revision: D85865653

fbshipit-source-id: 0efa95e68339d2fcdb1c403624962f1e63752dfe

Author: kinto0

pyrefly/lib/state/state.rs

@@ -1633,6 +1633,19 @@ fn collect_folding_ranges(body: &[Stmt]) -> Vec<(TextRange, Option<FoldingRangeK
     use ruff_python_ast::ExceptHandler;
     use ruff_text_size::Ranged;
 
+    fn range_without_decorators(
+        range: TextRange,
+        decorators: &[ruff_python_ast::Decorator],
+    ) -> TextRange {
+        let decorators_range = decorators
+            .first()
+            .map(|first| first.range().cover(decorators.last().unwrap().range()));
+
+        decorators_range.map_or(range, |x| {
+            range.add_start(x.len() + ruff_text_size::TextSize::from(1))
+        })
+    }
+
     struct FoldingRangeCollector {
         ranges: Vec<(TextRange, Option<FoldingRangeKind>)>,
     }
@@ -1649,12 +1662,14 @@ fn collect_folding_ranges(body: &[Stmt]) -> Vec<(TextRange, Option<FoldingRangeK
             match stmt {
                 Stmt::FunctionDef(func) => {
                     if !func.body.is_empty() {
-                        self.ranges.push((func.range, None));
+                        let range = range_without_decorators(func.range, &func.decorator_list);
+                        self.ranges.push((range, None));
                     }
                 }
                 Stmt::ClassDef(class) => {
                     if !class.body.is_empty() {
-                        self.ranges.push((class.range, None));
+                        let range = range_without_decorators(class.range, &class.decorator_list);
+                        self.ranges.push((range, None));
                     }
                 }
                 Stmt::If(if_stmt) => {

pyrefly/lib/test/lsp/folding_ranges.rs

@@ -478,3 +478,77 @@ if True:
         report.trim(),
     );
 }
+
+#[test]
+fn folding_ranges_with_decorators() {
+    let code = r#"
+from typing import overload, Any
+
+@overload
+def foo(x: int) -> str: ...
+@overload
+def foo(x: str) -> int: ...
+def foo(*args, **kwargs) -> Any:
+    pass
+
+def decorator1(cls):
+    return cls
+
+@decorator1
+class MyClass:
+    """Class docstring"""
+    
+    @staticmethod
+    def static_method():
+        pass
+    
+    @classmethod
+    def class_method(cls):
+        pass
+"#;
+
+    let report =
+        get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_folding_ranges_report);
+
+    assert_eq!(
+        r#"# main.py
+
+[
+  {
+    "start_line": 4,
+    "end_line": 4
+  },
+  {
+    "start_line": 6,
+    "end_line": 6
+  },
+  {
+    "start_line": 7,
+    "end_line": 8
+  },
+  {
+    "start_line": 10,
+    "end_line": 11
+  },
+  {
+    "start_line": 14,
+    "end_line": 23
+  },
+  {
+    "start_line": 15,
+    "end_line": 15,
+    "kind": "comment"
+  },
+  {
+    "start_line": 18,
+    "end_line": 19
+  },
+  {
+    "start_line": 22,
+    "end_line": 23
+  }
+]"#
+        .trim(),
+        report.trim(),
+    );
+}