Skip to content

Commit 72742e7

Browse files
iksuddleiksuddle
authored
[ty] Improve folding for decorators (#23543)
<!-- Thank you for contributing to Ruff/ty! To help us out with reviewing, please consider the following: - Does this pull request include a summary of the change? (See below.) - Does this pull request include a descriptive title? (Please prefix with `[ty]` for ty pull requests.) - Does this pull request include references to any relevant issues? --> ## Summary This change removes decorators from the folding ranges of classes/functions. Resolves astral-sh/ty#2861 ## Test Plan Added snapshot tests to check the returned folding ranges.
1 parent f8d1d29 commit 72742e7

1 file changed

Lines changed: 260 additions & 3 deletions

File tree

‎crates/ty_ide/src/folding_range.rs‎

Lines changed: 260 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
use ruff_db::files::File;
22
use ruff_db::parsed::parsed_module;
33
use ruff_db::source::source_text;
4+
use ruff_python_ast::token::{TokenKind, Tokens};
45
use ruff_python_ast::visitor::source_order::{SourceOrderVisitor, TraversalSignal, walk_body};
5-
use ruff_python_ast::{AnyNodeRef, Stmt};
6+
use ruff_python_ast::{AnyNodeRef, Stmt, StmtClassDef, StmtFunctionDef};
67
use ruff_source_file::{Line, UniversalNewlines};
78
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
89

@@ -51,6 +52,7 @@ pub fn folding_ranges(db: &dyn Db, file: File) -> Vec<FoldingRange> {
5152
let mut visitor = FoldingRangeVisitor {
5253
source: source.as_str(),
5354
ranges: vec![],
55+
tokens: parsed.tokens(),
5456
};
5557
visitor.visit_body(parsed.suite());
5658

@@ -67,6 +69,7 @@ pub fn folding_ranges(db: &dyn Db, file: File) -> Vec<FoldingRange> {
6769
struct FoldingRangeVisitor<'a> {
6870
source: &'a str,
6971
ranges: Vec<FoldingRange>,
72+
tokens: &'a Tokens,
7073
}
7174

7275
impl<'a> FoldingRangeVisitor<'a> {
@@ -223,14 +226,54 @@ impl<'a> FoldingRangeVisitor<'a> {
223226
}
224227
self.add_range(FoldingRange::from(first_stmt.range()).with_kind(FoldingRangeKind::Comment));
225228
}
229+
230+
/// Add a folding range for the function or class definition.
231+
///
232+
/// `target` is checked for in `search_range`, and is used as the start if found.
233+
fn add_def_range(&mut self, target: TokenKind, search_range: TextRange, end: TextSize) {
234+
let target_token = self
235+
.tokens
236+
.in_range(search_range)
237+
.iter()
238+
.find(|tok| tok.kind() == target);
239+
if let Some(tok) = target_token {
240+
let range = TextRange::new(tok.start(), end);
241+
self.add_range(range);
242+
}
243+
}
244+
245+
/// Add a folding range for function definitions, excluding decorators.
246+
fn add_function_def_range(&mut self, func: &StmtFunctionDef) {
247+
if let Some(decorator) = func.decorator_list.last() {
248+
let target = if func.is_async {
249+
TokenKind::Async
250+
} else {
251+
TokenKind::Def
252+
};
253+
let search_range = TextRange::new(decorator.end(), func.name.start());
254+
self.add_def_range(target, search_range, func.end());
255+
} else {
256+
self.add_range(func.range());
257+
}
258+
}
259+
260+
/// Add a folding range for class definitions, excluding decorators.
261+
fn add_class_def_range(&mut self, class: &StmtClassDef) {
262+
if let Some(decorator) = class.decorator_list.last() {
263+
let search_range = TextRange::new(decorator.end(), class.name.start());
264+
self.add_def_range(TokenKind::Class, search_range, class.end());
265+
} else {
266+
self.add_range(class.range());
267+
}
268+
}
226269
}
227270

228271
impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
229272
fn enter_node(&mut self, node: AnyNodeRef<'_>) -> TraversalSignal {
230273
match node {
231274
// Compound statements that create folding regions
232275
AnyNodeRef::StmtFunctionDef(func) => {
233-
self.add_range(func.range());
276+
self.add_function_def_range(func);
234277
// Note that this may be duplicative with folding
235278
// ranges added for string literals. But I don't think
236279
// the LSP protocol specifies that this is a problem.
@@ -241,7 +284,7 @@ impl SourceOrderVisitor<'_> for FoldingRangeVisitor<'_> {
241284
self.add_docstring_range(&func.body);
242285
}
243286
AnyNodeRef::StmtClassDef(class) => {
244-
self.add_range(class.range());
287+
self.add_class_def_range(class);
245288
// See comment above for class docstrings about this
246289
// being duplicative with adding folding ranges for
247290
// string literals.
@@ -1683,6 +1726,220 @@ with open("file.txt") as f:
16831726
}
16841727
}
16851728

1729+
#[test]
1730+
fn test_folding_range_decorated_function_single() {
1731+
let test = CursorTest::builder()
1732+
.source(
1733+
"main.py",
1734+
r#"
1735+
@decorator
1736+
def my_function():
1737+
pass
1738+
<CURSOR>
1739+
"#,
1740+
)
1741+
.build();
1742+
1743+
assert_snapshot!(test.folding_ranges(), @"
1744+
info[folding-range]: Folding Range
1745+
--> main.py:3:1
1746+
|
1747+
2 | @decorator
1748+
3 | / def my_function():
1749+
4 | | pass
1750+
| |________^
1751+
|
1752+
");
1753+
}
1754+
1755+
#[test]
1756+
fn test_folding_range_decorated_function_multiple() {
1757+
let test = CursorTest::builder()
1758+
.source(
1759+
"main.py",
1760+
r#"
1761+
@first
1762+
@second
1763+
@third
1764+
def my_function():
1765+
pass
1766+
<CURSOR>
1767+
"#,
1768+
)
1769+
.build();
1770+
1771+
assert_snapshot!(test.folding_ranges(), @"
1772+
info[folding-range]: Folding Range
1773+
--> main.py:5:1
1774+
|
1775+
3 | @second
1776+
4 | @third
1777+
5 | / def my_function():
1778+
6 | | pass
1779+
| |________^
1780+
|
1781+
");
1782+
}
1783+
1784+
#[test]
1785+
fn test_folding_range_decorated_class_single() {
1786+
let test = CursorTest::builder()
1787+
.source(
1788+
"main.py",
1789+
r#"
1790+
@dataclass
1791+
class MyClass:
1792+
value: int
1793+
name: str
1794+
<CURSOR>
1795+
"#,
1796+
)
1797+
.build();
1798+
1799+
// Single decorator is one line, so no decorator folding range is emitted.
1800+
assert_snapshot!(test.folding_ranges(), @"
1801+
info[folding-range]: Folding Range
1802+
--> main.py:3:1
1803+
|
1804+
2 | @dataclass
1805+
3 | / class MyClass:
1806+
4 | | value: int
1807+
5 | | name: str
1808+
| |_____________^
1809+
|
1810+
");
1811+
}
1812+
1813+
#[test]
1814+
fn test_folding_range_decorated_class_multiple() {
1815+
let test = CursorTest::builder()
1816+
.source(
1817+
"main.py",
1818+
r#"
1819+
@decorator_a
1820+
@decorator_b
1821+
class MyClass:
1822+
value: int
1823+
<CURSOR>
1824+
"#,
1825+
)
1826+
.build();
1827+
1828+
assert_snapshot!(test.folding_ranges(), @"
1829+
info[folding-range]: Folding Range
1830+
--> main.py:4:1
1831+
|
1832+
2 | @decorator_a
1833+
3 | @decorator_b
1834+
4 | / class MyClass:
1835+
5 | | value: int
1836+
| |______________^
1837+
|
1838+
");
1839+
}
1840+
1841+
#[test]
1842+
fn test_folding_range_decorated_async_function() {
1843+
let test = CursorTest::builder()
1844+
.source(
1845+
"main.py",
1846+
r#"
1847+
@decorator
1848+
async def my_async_function():
1849+
pass
1850+
<CURSOR>
1851+
"#,
1852+
)
1853+
.build();
1854+
1855+
assert_snapshot!(test.folding_ranges(), @"
1856+
info[folding-range]: Folding Range
1857+
--> main.py:3:1
1858+
|
1859+
2 | @decorator
1860+
3 | / async def my_async_function():
1861+
4 | | pass
1862+
| |________^
1863+
|
1864+
");
1865+
}
1866+
1867+
#[test]
1868+
fn test_folding_range_decorated_nested_function() {
1869+
let test = CursorTest::builder()
1870+
.source(
1871+
"main.py",
1872+
r#"
1873+
def outer_function():
1874+
@decorator
1875+
def inner_function():
1876+
pass
1877+
<CURSOR>
1878+
"#,
1879+
)
1880+
.build();
1881+
1882+
assert_snapshot!(test.folding_ranges(), @"
1883+
info[folding-range]: Folding Range
1884+
--> main.py:2:1
1885+
|
1886+
2 | / def outer_function():
1887+
3 | | @decorator
1888+
4 | | def inner_function():
1889+
5 | | pass
1890+
| |____________^
1891+
|
1892+
1893+
info[folding-range]: Folding Range
1894+
--> main.py:4:5
1895+
|
1896+
2 | def outer_function():
1897+
3 | @decorator
1898+
4 | / def inner_function():
1899+
5 | | pass
1900+
| |____________^
1901+
|
1902+
");
1903+
}
1904+
1905+
#[test]
1906+
fn test_folding_range_decorated_async_method() {
1907+
let test = CursorTest::builder()
1908+
.source(
1909+
"main.py",
1910+
r#"
1911+
class MyClass:
1912+
@decorator
1913+
async def my_async_method(self):
1914+
pass
1915+
<CURSOR>
1916+
"#,
1917+
)
1918+
.build();
1919+
1920+
assert_snapshot!(test.folding_ranges(), @"
1921+
info[folding-range]: Folding Range
1922+
--> main.py:2:1
1923+
|
1924+
2 | / class MyClass:
1925+
3 | | @decorator
1926+
4 | | async def my_async_method(self):
1927+
5 | | pass
1928+
| |____________^
1929+
|
1930+
1931+
info[folding-range]: Folding Range
1932+
--> main.py:4:5
1933+
|
1934+
2 | class MyClass:
1935+
3 | @decorator
1936+
4 | / async def my_async_method(self):
1937+
5 | | pass
1938+
| |____________^
1939+
|
1940+
");
1941+
}
1942+
16861943
struct FoldingRangeDiagnostic {
16871944
file: File,
16881945
folding_range: FoldingRange,

0 commit comments

Comments
 (0)