Skip to content

Commit eaf2ed5

Browse files
SamChou19815facebook-github-bot
SamChou19815
authored and
facebook-github-bot
committed
Support goto definition of attribute of imported module
Summary: In this diff, for attributes of imported module, we will now report its definition as `AttrDefinition::PartiallyResolvedImportedModuleAttribute(imported_module_name)` instead of `None`. Later in `lsp.rs` where we have access of transaction, we then resolve it to `TextRangeWithModuleInfo`. In this way, go-to-definition can work on this kind of attributes. For reference indexing, we can just store it into `externally_defined_variable_references`, since the following two kinds of references are equivalent: ``` import foo foo.bar # equivalent to from foo import bar ``` Reviewed By: kinto0 Differential Revision: D74487779 fbshipit-source-id: f4efa06b76e68627dbdaad056608233ad353eb55
1 parent 1f57bae commit eaf2ed5

File tree

5 files changed

+234
-56
lines changed

5 files changed

+234
-56
lines changed

‎pyrefly/lib/alt/answers.rs‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use starlark_map::Hashed;
2020
use starlark_map::small_map::SmallMap;
2121
use starlark_map::small_set::SmallSet;
2222

23+
use crate::alt::attr::AttrDefinition;
2324
use crate::alt::attr::AttrInfo;
2425
use crate::alt::traits::Solve;
2526
use crate::alt::traits::SolveRecursive;
@@ -852,18 +853,31 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
852853
definition,
853854
} in self.completions(base.clone(), Some(attribute_name), false)
854855
{
855-
if let Some(TextRangeWithModuleInfo {
856-
module_info: module,
857-
range,
858-
}) = definition
859-
&& module.path() != self.bindings().module_info().path()
860-
{
861-
index
862-
.lock()
863-
.externally_defined_attribute_references
864-
.entry(module.path().dupe())
865-
.or_default()
866-
.push((range, attribute_reference_range))
856+
match definition {
857+
Some(AttrDefinition::FullyResolved(TextRangeWithModuleInfo {
858+
module_info: module,
859+
range,
860+
})) => {
861+
if module.path() != self.bindings().module_info().path() {
862+
index
863+
.lock()
864+
.externally_defined_attribute_references
865+
.entry(module.path().dupe())
866+
.or_default()
867+
.push((range, attribute_reference_range))
868+
}
869+
}
870+
Some(AttrDefinition::PartiallyResolvedImportedModuleAttribute {
871+
module_name,
872+
}) => {
873+
index
874+
.lock()
875+
.externally_defined_variable_references
876+
.entry((module_name, attribute_name.clone()))
877+
.or_default()
878+
.push(attribute_reference_range);
879+
}
880+
None => {}
867881
}
868882
}
869883
}

‎pyrefly/lib/alt/attr.rs‎

Lines changed: 45 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1348,11 +1348,17 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
13481348
}
13491349
}
13501350

1351+
#[derive(Debug)]
1352+
pub enum AttrDefinition {
1353+
FullyResolved(TextRangeWithModuleInfo),
1354+
PartiallyResolvedImportedModuleAttribute { module_name: ModuleName },
1355+
}
1356+
13511357
#[derive(Debug)]
13521358
pub struct AttrInfo {
13531359
pub name: Name,
13541360
pub ty: Option<Type>,
1355-
pub definition: Option<TextRangeWithModuleInfo>,
1361+
pub definition: Option<AttrDefinition>,
13561362
}
13571363

13581364
impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
@@ -1381,9 +1387,8 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
13811387
res.push(AttrInfo {
13821388
name: fld.clone(),
13831389
ty: None,
1384-
definition: Some(TextRangeWithModuleInfo::new(
1385-
c.module_info().dupe(),
1386-
range,
1390+
definition: Some(AttrDefinition::FullyResolved(
1391+
TextRangeWithModuleInfo::new(c.module_info().dupe(), range),
13871392
)),
13881393
});
13891394
}
@@ -1394,9 +1399,8 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
13941399
res.push(AttrInfo {
13951400
name: expected_attribute_name.clone(),
13961401
ty: None,
1397-
definition: Some(TextRangeWithModuleInfo::new(
1398-
c.module_info().dupe(),
1399-
range,
1402+
definition: Some(AttrDefinition::FullyResolved(
1403+
TextRangeWithModuleInfo::new(c.module_info().dupe(), range),
14001404
)),
14011405
});
14021406
}
@@ -1427,7 +1431,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
14271431
res.extend(exports.wildcard(self.exports).iter().map(|x| AttrInfo {
14281432
name: x.clone(),
14291433
ty: None,
1430-
definition: None,
1434+
definition: Some(
1435+
AttrDefinition::PartiallyResolvedImportedModuleAttribute {
1436+
module_name,
1437+
},
1438+
),
14311439
}));
14321440
}
14331441
Some(expected_attribute_name) => {
@@ -1438,7 +1446,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
14381446
res.push(AttrInfo {
14391447
name: expected_attribute_name.clone(),
14401448
ty: None,
1441-
definition: None,
1449+
definition: Some(
1450+
AttrDefinition::PartiallyResolvedImportedModuleAttribute {
1451+
module_name,
1452+
},
1453+
),
14421454
});
14431455
}
14441456
}
@@ -1494,18 +1506,30 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
14941506
}
14951507
if include_types {
14961508
for info in &mut res {
1497-
if let Some(TextRangeWithModuleInfo {
1498-
module_info: _,
1499-
range,
1500-
}) = info.definition
1501-
{
1502-
info.ty =
1503-
match self.lookup_attr_from_attribute_base(base.clone(), &info.name) {
1504-
LookupResult::Found(attr) => self
1505-
.resolve_get_access(attr, range, &self.error_swallower(), None)
1506-
.ok(),
1507-
_ => None,
1508-
};
1509+
if let Some(definition) = &info.definition {
1510+
match definition {
1511+
AttrDefinition::FullyResolved(TextRangeWithModuleInfo {
1512+
module_info: _,
1513+
range,
1514+
}) => {
1515+
info.ty = match self
1516+
.lookup_attr_from_attribute_base(base.clone(), &info.name)
1517+
{
1518+
LookupResult::Found(attr) => self
1519+
.resolve_get_access(
1520+
attr,
1521+
*range,
1522+
&self.error_swallower(),
1523+
None,
1524+
)
1525+
.ok(),
1526+
_ => None,
1527+
};
1528+
}
1529+
AttrDefinition::PartiallyResolvedImportedModuleAttribute {
1530+
module_name: _,
1531+
} => {}
1532+
}
15091533
}
15101534
}
15111535
}

‎pyrefly/lib/state/lsp.rs‎

Lines changed: 38 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use ruff_text_size::TextRange;
2020
use ruff_text_size::TextSize;
2121
use starlark_map::ordered_set::OrderedSet;
2222

23+
use crate::alt::attr::AttrDefinition;
2324
use crate::alt::attr::AttrInfo;
2425
use crate::binding::binding::Binding;
2526
use crate::binding::binding::Key;
@@ -186,6 +187,30 @@ impl<'a> Transaction<'a> {
186187
}
187188
}
188189

190+
fn resolve_attribute_definition(
191+
&self,
192+
handle: &Handle,
193+
attr_name: &Name,
194+
definition: AttrDefinition,
195+
) -> Option<(TextRangeWithModuleInfo, Option<DocString>)> {
196+
match definition {
197+
AttrDefinition::FullyResolved(text_range_with_module_info) => {
198+
// TODO(kylei): attribute docstrings
199+
Some((text_range_with_module_info, None))
200+
}
201+
AttrDefinition::PartiallyResolvedImportedModuleAttribute { module_name } => {
202+
let mut gas = INITIAL_GAS;
203+
let (handle, export) =
204+
self.resolve_named_import(handle, module_name, attr_name.clone(), &mut gas)?;
205+
let module_info = self.get_module_info(&handle)?;
206+
Some((
207+
TextRangeWithModuleInfo::new(module_info, export.location),
208+
export.docstring,
209+
))
210+
}
211+
}
212+
}
213+
189214
fn key_to_export(&self, handle: &Handle, key: &Key, mut gas: Gas) -> Option<(Handle, Export)> {
190215
let bindings = self.get_bindings(handle)?;
191216
let intermediate_definition = key_to_intermediate_definition(&bindings, key, &mut gas)?;
@@ -261,8 +286,9 @@ impl<'a> Transaction<'a> {
261286
let items = solver.completions(base_type.arc_clone(), Some(&attribute.attr.id), false);
262287
items.into_iter().find_map(|x| {
263288
if x.name == attribute.attr.id {
264-
// TODO(kylei): attribute docstrings
265-
Some((DefinitionMetadata::Attribute(x.name), x.definition?, None))
289+
let (definition, docstring) =
290+
self.resolve_attribute_definition(handle, &x.name, x.definition?)?;
291+
Some((DefinitionMetadata::Attribute(x.name), definition, docstring))
266292
} else {
267293
None
268294
}
@@ -383,16 +409,20 @@ impl<'a> Transaction<'a> {
383409
&& let Some(base_type) = answers.get_type_trace(attribute.value.range())
384410
{
385411
for AttrInfo {
386-
name: _,
412+
name,
387413
ty: _,
388414
definition: attribute_definition,
389415
} in solver.completions(base_type.arc_clone(), Some(expected_name), false)
390416
{
391-
if let Some(TextRangeWithModuleInfo {
392-
module_info: module,
393-
range,
394-
}) = attribute_definition
395-
&& module.path() == definition.module_info.path()
417+
if let Some((
418+
TextRangeWithModuleInfo {
419+
module_info: module,
420+
range,
421+
},
422+
_,
423+
)) = attribute_definition.and_then(|definition| {
424+
self.resolve_attribute_definition(handle, &name, definition)
425+
}) && module.path() == definition.module_info.path()
396426
&& range == definition.range
397427
{
398428
references.push(attribute.attr.range());

‎pyrefly/lib/test/lsp/definition.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -663,3 +663,35 @@ Definition Result:
663663
report.trim(),
664664
);
665665
}
666+
667+
#[test]
668+
fn module_attribute_test() {
669+
let code_class_provider = r#"
670+
class MyClass:
671+
x = 5
672+
"#;
673+
let code = r#"
674+
import my_class
675+
my_class.MyClass
676+
# ^
677+
"#;
678+
let report = get_batched_lsp_operations_report(
679+
&[("main", code), ("my_class", code_class_provider)],
680+
get_test_report,
681+
);
682+
assert_eq!(
683+
r#"
684+
# main.py
685+
3 | my_class.MyClass
686+
^
687+
Definition Result:
688+
2 | class MyClass:
689+
^^^^^^^
690+
691+
692+
# my_class.py
693+
"#
694+
.trim(),
695+
report.trim(),
696+
);
697+
}

0 commit comments

Comments
 (0)