folding ranges for expressions

Summary: literals + calls

Reviewed By: stroxler

Differential Revision: D85867049

fbshipit-source-id: e32f6e6fac5606dd9405c2233b61d84e59c8a217

Author: kinto0

pyrefly/lib/state/state.rs

@@ -58,10 +58,12 @@ use pyrefly_util::uniques::UniqueFactory;
 use pyrefly_util::upgrade_lock::UpgradeLock;
 use pyrefly_util::upgrade_lock::UpgradeLockExclusiveGuard;
 use pyrefly_util::upgrade_lock::UpgradeLockWriteGuard;
+use ruff_python_ast::Expr;
 use ruff_python_ast::Stmt;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::visitor::walk_body;
+use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 use starlark_map::Hashed;
 use starlark_map::small_map::SmallMap;
@@ -507,7 +509,8 @@ impl<'a> Transaction<'a> {
         handle: &Handle,
     ) -> Option<Vec<(TextRange, Option<FoldingRangeKind>)>> {
         let ast = self.get_ast(handle)?;
-        let mut ranges = collect_folding_ranges(&ast.body);
+        let module_info = self.get_module_info(handle)?;
+        let mut ranges = collect_folding_ranges(&ast.body, &module_info);
         ranges.sort_by_key(|(range, _)| range.start());
         ranges.dedup();
         Some(ranges)
@@ -1629,7 +1632,10 @@ impl<'a> Transaction<'a> {
     }
 }
 
-fn collect_folding_ranges(body: &[Stmt]) -> Vec<(TextRange, Option<FoldingRangeKind>)> {
+fn collect_folding_ranges(
+    body: &[Stmt],
+    module: &Module,
+) -> Vec<(TextRange, Option<FoldingRangeKind>)> {
     use ruff_python_ast::ExceptHandler;
     use ruff_text_size::Ranged;
 
@@ -1646,11 +1652,12 @@ fn collect_folding_ranges(body: &[Stmt]) -> Vec<(TextRange, Option<FoldingRangeK
         })
     }
 
-    struct FoldingRangeCollector {
+    struct FoldingRangeCollector<'a> {
         ranges: Vec<(TextRange, Option<FoldingRangeKind>)>,
+        module: &'a Module,
     }
 
-    impl Visitor<'_> for FoldingRangeCollector {
+    impl Visitor<'_> for FoldingRangeCollector<'_> {
         fn visit_body(&mut self, body: &[Stmt]) {
             if let Some(range) = Docstring::range_from_stmts(body) {
                 self.ranges.push((range, Some(FoldingRangeKind::Comment)));
@@ -1720,9 +1727,31 @@ fn collect_folding_ranges(body: &[Stmt]) -> Vec<(TextRange, Option<FoldingRangeK
             }
             ruff_python_ast::visitor::walk_stmt(self, stmt);
         }
+
+        fn visit_expr(&mut self, expr: &Expr) {
+            let range = match expr {
+                Expr::Call(call) => Some(call.arguments.range),
+                Expr::Dict(dict) => Some(dict.range),
+                Expr::List(list) => Some(list.range),
+                Expr::Set(set) => Some(set.range),
+                Expr::Tuple(tuple) => Some(tuple.range),
+                _ => None,
+            };
+
+            if let Some(range) = range {
+                let lsp_range = self.module.lined_buffer().to_lsp_range(range);
+                if lsp_range.start.line != lsp_range.end.line {
+                    self.ranges.push((range, None));
+                }
+            }
+            ruff_python_ast::visitor::walk_expr(self, expr);
+        }
     }
 
-    let mut collector = FoldingRangeCollector { ranges: Vec::new() };
+    let mut collector = FoldingRangeCollector {
+        ranges: Vec::new(),
+        module,
+    };
 
     if let Some(range) = Docstring::range_from_stmts(body) {
         collector

pyrefly/lib/test/lsp/folding_ranges.rs

@@ -552,3 +552,180 @@ class MyClass:
         report.trim(),
     );
 }
+
+#[test]
+fn folding_ranges_for_calls() {
+    let code = r#"
+def foo(*args):
+    pass
+
+def short_call(*args):
+    pass
+
+def multi_line_nested(*args):
+    pass
+
+def bar(*args):
+    pass
+
+def very_long_function_name(**kwargs):
+    pass
+
+nested_arg = 1
+another_arg = 2
+second_arg = 3
+value1 = 1
+value2 = 2
+value3 = 3
+
+result = foo(
+    1,
+    2,
+    3,
+    4
+)
+
+short_call(1, 2, 3)
+
+multi_line_nested(
+    bar(
+        nested_arg,
+        another_arg
+    ),
+    second_arg
+)
+
+x = very_long_function_name(
+    param1=value1,
+    param2=value2,
+    param3=value3
+)
+"#;
+
+    let report =
+        get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_folding_ranges_report);
+
+    assert_eq!(
+        r#"# main.py
+
+[
+  {
+    "start_line": 1,
+    "end_line": 2
+  },
+  {
+    "start_line": 4,
+    "end_line": 5
+  },
+  {
+    "start_line": 7,
+    "end_line": 8
+  },
+  {
+    "start_line": 10,
+    "end_line": 11
+  },
+  {
+    "start_line": 13,
+    "end_line": 14
+  },
+  {
+    "start_line": 23,
+    "end_line": 28
+  },
+  {
+    "start_line": 32,
+    "end_line": 38
+  },
+  {
+    "start_line": 33,
+    "end_line": 36
+  },
+  {
+    "start_line": 40,
+    "end_line": 44
+  }
+]"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn folding_ranges_for_literals() {
+    let code = r#"
+data = {
+    "key1": "value1",
+    "key2": "value2",
+    "key3": "value3"
+}
+
+items = [
+    1,
+    2,
+    3,
+    4
+]
+
+my_set = {
+    "a",
+    "b",
+    "c"
+}
+
+coordinates = (
+    10,
+    20,
+    30
+)
+
+nested = {
+    "inner_dict": {
+        "nested_key": "value"
+    },
+    "inner_list": [
+        1, 2, 3
+    ]
+}
+"#;
+
+    let report =
+        get_batched_lsp_operations_report_no_cursor(&[("main", code)], get_folding_ranges_report);
+
+    assert_eq!(
+        r#"# main.py
+
+[
+  {
+    "start_line": 1,
+    "end_line": 5
+  },
+  {
+    "start_line": 7,
+    "end_line": 12
+  },
+  {
+    "start_line": 14,
+    "end_line": 18
+  },
+  {
+    "start_line": 20,
+    "end_line": 24
+  },
+  {
+    "start_line": 26,
+    "end_line": 33
+  },
+  {
+    "start_line": 27,
+    "end_line": 29
+  },
+  {
+    "start_line": 30,
+    "end_line": 32
+  }
+]"#
+        .trim(),
+        report.trim(),
+    );
+}