55 * LICENSE file in the root directory of this source tree.
66 */
77
8+ use std:: iter;
9+
810use pyrefly_python:: ast:: Ast ;
911use pyrefly_python:: docstring:: Docstring ;
1012use pyrefly_python:: module:: Module ;
1113use ruff_python_ast:: Expr ;
1214use ruff_python_ast:: ModModule ;
1315use ruff_python_ast:: Stmt ;
16+ use ruff_python_ast:: StmtIf ;
1417use ruff_text_size:: Ranged ;
1518use ruff_text_size:: TextRange ;
1619
@@ -55,6 +58,14 @@ impl AttributeContext {
5558 Stmt :: FunctionDef ( function) if function. name . range == target_range => {
5659 AttributeKind :: Function
5760 }
61+ Stmt :: If ( if_stmt) => {
62+ if let Some ( context) = if_branches ( if_stmt) . find_map ( |branch| {
63+ Self :: from_body ( branch, & mut parents. clone ( ) , target_range)
64+ } ) {
65+ return Some ( context) ;
66+ }
67+ continue ;
68+ }
5869 Stmt :: Assign ( assign)
5970 if assign. targets . iter ( ) . any (
6071 |target| matches ! ( target, Expr :: Name ( name) if name. range( ) == target_range) ,
@@ -92,26 +103,49 @@ pub(crate) fn definition_from_ast(
92103 context : & AttributeContext ,
93104 attr_name : & ruff_python_ast:: name:: Name ,
94105) -> Option < ( TextRange , Option < TextRange > ) > {
95- let mut body = ast. body . as_slice ( ) ;
106+ let mut body = flatten_conditionals ( & ast. body ) ;
96107 for class_name in & context. parent_classes {
97108 let class_def = body. iter ( ) . find_map ( |stmt| match stmt {
98109 Stmt :: ClassDef ( class_def) if & class_def. name . id == class_name => Some ( class_def) ,
99110 _ => None ,
100111 } ) ?;
101- body = class_def. body . as_slice ( ) ;
112+ body = flatten_conditionals ( & class_def. body ) ;
102113 }
103114 match context. kind {
104- AttributeKind :: Function => definition_from_function ( body, attr_name) ,
105- AttributeKind :: Class => definition_from_class ( body, attr_name) ,
106- AttributeKind :: Assignment => definition_from_assignment ( body, attr_name) ,
115+ AttributeKind :: Function => definition_from_function ( & body, attr_name) ,
116+ AttributeKind :: Class => definition_from_class ( & body, attr_name) ,
117+ AttributeKind :: Assignment => definition_from_assignment ( & body, attr_name) ,
107118 }
108119}
109120
121+ /// The bodies of the `if`, `elif`, and `else` branches of `if_stmt`.
122+ fn if_branches ( if_stmt : & StmtIf ) -> impl Iterator < Item = & [ Stmt ] > {
123+ iter:: once ( if_stmt. body . as_slice ( ) ) . chain (
124+ if_stmt
125+ . elif_else_clauses
126+ . iter ( )
127+ . map ( |clause| clause. body . as_slice ( ) ) ,
128+ )
129+ }
130+
131+ /// Lists the statements of `body`, including those nested in `if` blocks such as
132+ /// `if sys.version_info >= ...:`, which do not introduce a scope.
133+ fn flatten_conditionals ( body : & [ Stmt ] ) -> Vec < & Stmt > {
134+ body. iter ( )
135+ . flat_map ( |stmt| match stmt {
136+ Stmt :: If ( if_stmt) => if_branches ( if_stmt)
137+ . flat_map ( flatten_conditionals)
138+ . collect ( ) ,
139+ _ => vec ! [ stmt] ,
140+ } )
141+ . collect ( )
142+ }
143+
110144fn definition_from_function (
111- body : & [ Stmt ] ,
145+ body : & [ & Stmt ] ,
112146 attr_name : & ruff_python_ast:: name:: Name ,
113147) -> Option < ( TextRange , Option < TextRange > ) > {
114- for stmt in body {
148+ for & stmt in body {
115149 if let Stmt :: FunctionDef ( func_def) = stmt
116150 && & func_def. name . id == attr_name
117151 {
@@ -123,10 +157,10 @@ fn definition_from_function(
123157}
124158
125159fn definition_from_class (
126- body : & [ Stmt ] ,
160+ body : & [ & Stmt ] ,
127161 attr_name : & ruff_python_ast:: name:: Name ,
128162) -> Option < ( TextRange , Option < TextRange > ) > {
129- for stmt in body {
163+ for & stmt in body {
130164 if let Stmt :: ClassDef ( class_def) = stmt
131165 && & class_def. name . id == attr_name
132166 {
@@ -138,10 +172,10 @@ fn definition_from_class(
138172}
139173
140174fn definition_from_assignment (
141- body : & [ Stmt ] ,
175+ body : & [ & Stmt ] ,
142176 attr_name : & ruff_python_ast:: name:: Name ,
143177) -> Option < ( TextRange , Option < TextRange > ) > {
144- for ( idx, stmt) in body. iter ( ) . enumerate ( ) {
178+ for ( idx, & stmt) in body. iter ( ) . enumerate ( ) {
145179 let ( matches, def_range) = match stmt {
146180 Stmt :: Assign ( assign) => {
147181 let target_opt = assign
@@ -157,7 +191,7 @@ fn definition_from_assignment(
157191 _ => ( false , None ) ,
158192 } ;
159193 if matches {
160- let docstring_range = if let Some ( Stmt :: Expr ( expr_stmt) ) = body. get ( idx + 1 )
194+ let docstring_range = if let Some ( Stmt :: Expr ( expr_stmt) ) = body. get ( idx + 1 ) . copied ( )
161195 && matches ! ( expr_stmt. value. as_ref( ) , Expr :: StringLiteral ( _) )
162196 {
163197 Some ( expr_stmt. range ( ) )
@@ -169,3 +203,33 @@ fn definition_from_assignment(
169203 }
170204 None
171205}
206+
207+ #[ cfg( test) ]
208+ mod tests {
209+ use std:: sync:: Arc ;
210+
211+ use pyrefly_python:: module_name:: ModuleName ;
212+ use pyrefly_python:: module_path:: ModulePath ;
213+ use ruff_python_ast:: name:: Name ;
214+
215+ use super :: * ;
216+
217+ /// A method declared under `if sys.version_info` in a stub maps to the unconditional
218+ /// method in the source, and the source docstring is found.
219+ #[ test]
220+ fn test_method_in_conditional_block ( ) {
221+ let stub = "class C:\n if sys.version_info >= (3, 12):\n def f(self): ...\n else:\n def f(self, x): ...\n " ;
222+ let source = "class C:\n def f(self):\n \" \" \" Doc.\" \" \" \n " ;
223+ let stub_module = Module :: new (
224+ ModuleName :: from_str ( "m" ) ,
225+ ModulePath :: filesystem ( "m.pyi" . into ( ) ) ,
226+ Arc :: new ( stub. to_owned ( ) ) ,
227+ ) ;
228+ let offset = stub. find ( "f(self)" ) . unwrap ( ) ;
229+ let range = TextRange :: at ( u32:: try_from ( offset) . unwrap ( ) . into ( ) , 1 . into ( ) ) ;
230+ let context = AttributeContext :: from_module ( & stub_module, range) . unwrap ( ) ;
231+ let ast = Ast :: parse ( source, stub_module. source_type ( ) ) . 0 ;
232+ let ( _, docstring) = definition_from_ast ( & ast, & context, & Name :: new ( "f" ) ) . unwrap ( ) ;
233+ assert ! ( docstring. is_some( ) ) ;
234+ }
235+ }
0 commit comments