diff --git a/pyrefly/lib/export/symbols.rs b/pyrefly/lib/export/symbols.rs index f174272703..6272ed38e1 100644 --- a/pyrefly/lib/export/symbols.rs +++ b/pyrefly/lib/export/symbols.rs @@ -7,7 +7,8 @@ //! A flat, source-order table of definitions used to find nested workspace //! symbols: functions, classes, methods, type aliases, and simple module/class -//! assignments. It is cached on `Exports` for first-party modules because the +//! assignments, including attributes assigned through a method's receiver. +//! It is cached on `Exports` for first-party modules because the //! export table itself contains only top-level names and therefore cannot //! provide nested definitions. //! @@ -73,6 +74,7 @@ impl FlatSymbols { Scope { parent: None, kind: ScopeKind::Module, + receiver: None, }, &mut out, ); @@ -99,10 +101,13 @@ enum ScopeKind { /// Where a run of statements sits and which definition encloses it. #[derive(Clone, Copy)] -struct Scope { +struct Scope<'a> { /// Index of the enclosing definition; `None` at module level. parent: Option, kind: ScopeKind, + /// A bound method's first positional parameter and its enclosing class. + /// Nested functions and classes establish their own receiver scope. + receiver: Option<(&'a Name, FlatSymbolIndex)>, } fn push_symbol( @@ -116,7 +121,7 @@ fn push_symbol( idx } -fn assignment_kind(name: &Name, scope: Scope) -> SymbolKind { +fn assignment_kind(name: &Name, scope: Scope<'_>) -> SymbolKind { if is_constant_name(name) { SymbolKind::Constant } else if scope.kind == ScopeKind::Class { @@ -126,15 +131,45 @@ fn assignment_kind(name: &Name, scope: Scope) -> SymbolKind { } } -fn push_assignment_targets(out: &mut Vec, target: &Expr, scope: Scope) { - Ast::expr_lvalue(target, &mut |name| { - push_symbol( - out, - ShortIdentifier::expr_name(name), - assignment_kind(&name.id, scope), - scope.parent, - ); - }); +fn push_assignment_targets(out: &mut Vec, target: &Expr, scope: Scope<'_>) { + match target { + Expr::Name(name) + if scope.kind != ScopeKind::Function && !Ast::is_synthesized_empty_name(name) => + { + push_symbol( + out, + ShortIdentifier::expr_name(name), + assignment_kind(&name.id, scope), + scope.parent, + ); + } + Expr::Attribute(attr) => { + if let Some((receiver, class)) = scope.receiver + && let Expr::Name(value) = &*attr.value + && &value.id == receiver + && !Ast::is_synthesized_empty_identifier(&attr.attr) + { + push_symbol( + out, + ShortIdentifier::new(&attr.attr), + SymbolKind::Attribute, + Some(class), + ); + } + } + Expr::Tuple(tuple) => { + for target in &tuple.elts { + push_assignment_targets(out, target, scope); + } + } + Expr::List(list) => { + for target in &list.elts { + push_assignment_targets(out, target, scope); + } + } + Expr::Starred(starred) => push_assignment_targets(out, &starred.value, scope), + _ => {} + } } /// Walk `stmts` appending symbols to `out`. @@ -142,7 +177,7 @@ fn push_assignment_targets(out: &mut Vec, target: &Expr, scope: Scop /// Functions and classes are recorded at any depth, including inside a function /// body. Control-flow statements are descended into with the scope unchanged, so /// a class attribute guarded by an `if` still attaches to its class. -fn build(stmts: &[Stmt], scope: Scope, out: &mut Vec) { +fn build<'a>(stmts: &'a [Stmt], scope: Scope<'a>, out: &mut Vec) { for stmt in stmts { match stmt { Stmt::FunctionDef(f) => { @@ -164,11 +199,32 @@ fn build(stmts: &[Stmt], scope: Scope, out: &mut Vec) { scope.parent, )) }; + // The export pass is syntactic: recognize spelled-out staticmethod + // decorators without resolving imports or evaluating decorators. + let is_staticmethod = + f.decorator_list + .iter() + .any(|decorator| match &decorator.expression { + Expr::Name(name) => name.id == "staticmethod", + Expr::Attribute(attr) => attr.attr.id == "staticmethod", + _ => false, + }); + let receiver = if scope.kind == ScopeKind::Class && !is_staticmethod { + f.parameters + .posonlyargs + .first() + .or_else(|| f.parameters.args.first()) + .zip(scope.parent) + .map(|(parameter, class)| (¶meter.parameter.name.id, class)) + } else { + None + }; build( &f.body, Scope { parent: body_parent, kind: ScopeKind::Function, + receiver, }, out, ); @@ -189,16 +245,17 @@ fn build(stmts: &[Stmt], scope: Scope, out: &mut Vec) { Scope { parent: body_parent, kind: ScopeKind::Class, + receiver: None, }, out, ); } - Stmt::Assign(a) if scope.kind != ScopeKind::Function => { + Stmt::Assign(a) => { for target in &a.targets { push_assignment_targets(out, target, scope); } } - Stmt::AnnAssign(a) if scope.kind != ScopeKind::Function => { + Stmt::AnnAssign(a) => { push_assignment_targets(out, &a.target, scope); } Stmt::TypeAlias(t) if scope.kind != ScopeKind::Function => { @@ -282,6 +339,99 @@ mod tests { ); } + #[test] + fn test_receiver_attribute_assignment_targets() { + assert_eq!( + walk( + r#" +class Example: + def __init__(this, /): + this._private = this.public = 1 + this.annotated: int = 2 + this.first, [this.second, *this.rest] = values + if condition: + this.conditional = 3 + local = 4 + other.unrelated = 5 + this.child.unrelated = 6 + this.items[0] = 7 +"# + ), + vec![ + "Class Example", + ".Method __init__", + ".Attribute _private", + ".Attribute public", + ".Attribute annotated", + ".Attribute first", + ".Attribute second", + ".Attribute rest", + ".Attribute conditional", + ] + ); + } + + #[test] + fn test_receiver_scope_does_not_leak_into_nested_definitions() { + assert_eq!( + walk( + r#" +def top_level(self): + self.unrelated = 1 +class Example: + def method(self): + def nested(self): + self.unrelated = 2 + class Inner: + self.unrelated = 3 + def method(this): + this.inner = 4 + self.outer = 5 +"# + ), + vec![ + "Function top_level", + "Class Example", + ".Method method", + "..Function nested", + "..Class Inner", + "...Method method", + "...Attribute inner", + ".Attribute outer", + ] + ); + } + + #[test] + fn test_static_methods_have_no_receiver() { + assert_eq!( + walk( + r#" +class Example: + @staticmethod + def static(self): + self.unrelated = 1 + @builtins.staticmethod + def qualified_static(self): + self.unrelated = 2 + def no_positional_parameter(*, self): + self.unrelated = 3 + @classmethod + def class_method(cls): + cls.class_attribute = 4 +"# + ), + vec![ + "Class Example", + ".Method static", + ".Method qualified_static", + ".Method no_positional_parameter", + ".Method class_method", + ".Attribute class_attribute", + ] + ); + } + #[test] fn test_control_flow_does_not_change_the_enclosing_scope() { assert_eq!( diff --git a/pyrefly/lib/test/lsp.rs b/pyrefly/lib/test/lsp.rs index 38e0f9b16a..45d72bb130 100644 --- a/pyrefly/lib/test/lsp.rs +++ b/pyrefly/lib/test/lsp.rs @@ -29,3 +29,4 @@ mod rename; mod semantic_tokens; mod signature_help; mod type_definition; +mod workspace_symbols; diff --git a/pyrefly/lib/test/lsp/workspace_symbols.rs b/pyrefly/lib/test/lsp/workspace_symbols.rs new file mode 100644 index 0000000000..c9c9b10427 --- /dev/null +++ b/pyrefly/lib/test/lsp/workspace_symbols.rs @@ -0,0 +1,105 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use lsp_types::SymbolKind; + +use crate::state::require::Require; +use crate::test::util::TestEnv; + +#[test] +fn test_workspace_symbols_underscore_prefixed_methods() { + let code = r#" +class Example: + def _private_method(self) -> None: + pass + + def public_method(self) -> None: + pass +"#; + let (state, _) = TestEnv::one("main", code).to_state(); + let transaction = state.transaction(); + for name in ["_private_method", "public_method"] { + let symbols = transaction.workspace_symbols(name, None).unwrap(); + assert_eq!(symbols.len(), 1, "expected {name} in workspace symbols"); + let symbol = &symbols[0]; + assert_eq!(symbol.name, name); + assert_eq!(symbol.kind, SymbolKind::METHOD); + assert_eq!(symbol.container_name.as_deref(), Some("Example")); + assert_eq!(symbol.location.module.code_at(symbol.location.range), name); + } +} + +// https://github.com/facebook/pyrefly/issues/4688 +#[test] +fn test_workspace_symbols_instance_attributes() { + let code = r#" +class Example: + def __init__(self) -> None: + self._private_member = 1 + self.public_member = 2 +"#; + for require in [Require::Everything, Require::Indexing] { + let (state, _) = TestEnv::one("main", code) + .with_run_require(require) + .to_state(); + let transaction = state.transaction(); + for name in ["_private_member", "public_member"] { + let symbols = transaction.workspace_symbols(name, None).unwrap(); + assert_eq!(symbols.len(), 1, "expected {name} with {require:?}"); + let symbol = &symbols[0]; + assert_eq!(symbol.name, name); + assert_eq!(symbol.kind, SymbolKind::FIELD); + assert_eq!(symbol.container_name.as_deref(), Some("Example")); + assert_eq!(symbol.location.module.code_at(symbol.location.range), name); + assert!( + transaction + .search_exports_fuzzy(name, None) + .unwrap() + .is_empty() + ); + } + } +} + +#[test] +fn test_workspace_symbols_instance_attribute_deduplication() { + let code = r#" +class Example: + def __init__(self) -> None: + self._private_member = 1 + + def reset(self) -> None: + self._private_member = 2 + +class Other: + def __init__(self) -> None: + self._private_member = 3 +"#; + let (state, _) = TestEnv::one("main", code).to_state(); + let transaction = state.transaction(); + let symbols = transaction + .workspace_symbols("_private_member", None) + .unwrap(); + let mut containers = symbols + .iter() + .map(|symbol| { + assert_eq!(symbol.name, "_private_member"); + assert_eq!(symbol.kind, SymbolKind::FIELD); + symbol.container_name.as_deref().unwrap() + }) + .collect::>(); + containers.sort_unstable(); + assert_eq!(containers, ["Example", "Other"]); + let symbol = symbols + .iter() + .find(|symbol| symbol.container_name.as_deref() == Some("Example")) + .unwrap(); + assert_eq!( + symbol.location.range.start().to_usize(), + code.find("_private_member").unwrap(), + ); +}