Skip to content

Commit fd39880

Browse files
rubmarymeta-codesync[bot]
rubmary
authored and
meta-codesync[bot]
committed
Record xrefs for symbols __all__
Summary: We currently do not record xrefs for symbols in `__all__ `which is a gap compared to the current python indexer. Reviewed By: aahanaggarwal Differential Revision: D85440096 fbshipit-source-id: 6fb0b3cecc94f361400eb5e00e2d9d8ca1b7b052
1 parent 6698c93 commit fd39880

File tree

2 files changed

+167
-59
lines changed

2 files changed

+167
-59
lines changed

‎pyrefly/lib/report/glean/convert.rs‎

Lines changed: 95 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@ use num_traits::ToPrimitive;
1717
use pyrefly_build::handle::Handle;
1818
use pyrefly_python::ast::Ast;
1919
use pyrefly_python::docstring::Docstring;
20+
use pyrefly_python::dunder;
2021
use pyrefly_python::module_name::ModuleName;
2122
use pyrefly_util::visit::Visit;
2223
use ruff_python_ast::Decorator;
2324
use ruff_python_ast::ExceptHandler;
2425
use ruff_python_ast::Expr;
2526
use ruff_python_ast::ExprAttribute;
2627
use ruff_python_ast::ExprCall;
28+
use ruff_python_ast::ExprList;
2729
use ruff_python_ast::ExprName;
2830
use ruff_python_ast::ExprStringLiteral;
2931
use ruff_python_ast::Identifier;
@@ -168,6 +170,13 @@ struct GleanState<'a> {
168170
facts: Facts,
169171
names: HashSet<Arc<String>>,
170172
locations_fqnames: HashMap<TextSize, Arc<String>>,
173+
import_names: HashMap<Arc<String>, Arc<String>>,
174+
}
175+
176+
struct AssignInfo<'a> {
177+
range: TextRange,
178+
annotation: Option<&'a Expr>,
179+
value: Option<&'a Expr>,
171180
}
172181

173182
impl Facts {
@@ -204,6 +213,7 @@ impl GleanState<'_> {
204213
),
205214
names: HashSet::new(),
206215
locations_fqnames: HashMap::new(),
216+
import_names: HashMap::new(),
207217
}
208218
}
209219

@@ -231,7 +241,7 @@ impl GleanState<'_> {
231241
let name = module.map_or(ModuleName::from_name(&component), |x: ModuleName| {
232242
x.append(&component)
233243
});
234-
self.record_name(name.to_string(), None);
244+
self.record_name(name.to_string());
235245
module = Some(name);
236246
module_names.push(name);
237247
}
@@ -248,7 +258,7 @@ impl GleanState<'_> {
248258
let name = module.map_or(ModuleName::from_name(&component), |x: ModuleName| {
249259
x.append(&component)
250260
});
251-
self.record_name(name.to_string(), None);
261+
self.record_name(name.to_string());
252262
self.facts
253263
.modules
254264
.push(python::Module::new(python::Name::new(name.to_string())));
@@ -321,19 +331,29 @@ impl GleanState<'_> {
321331
}
322332
}
323333

324-
fn record_name(&mut self, name: String, position: Option<TextSize>) -> python::Name {
334+
fn record_name(&mut self, name: String) -> Arc<String> {
325335
let arc_name = Arc::new(name.clone());
326336
if self.names.insert(arc_name.dupe()) {
327337
self.facts.name_to_sname.push(python::NameToSName::new(
328338
python::Name::new(name.clone()),
329339
create_sname(&name),
330340
));
331341
}
332-
position.map(|x| self.locations_fqnames.insert(x, arc_name.dupe()));
342+
arc_name
343+
}
333344

345+
fn record_name_with_position(&mut self, name: String, position: TextSize) -> python::Name {
346+
let arc_name = self.record_name(name.clone());
347+
self.locations_fqnames.insert(position, arc_name.dupe());
334348
python::Name::new(name)
335349
}
336350

351+
fn record_name_for_import(&mut self, name: String, import_name: &str) {
352+
let arc_name = self.record_name(name);
353+
let arc_import_name = Arc::new(import_name.to_owned());
354+
self.import_names.insert(arc_name.dupe(), arc_import_name);
355+
}
356+
337357
fn make_fq_name_for_declaration(
338358
&mut self,
339359
name: &Identifier,
@@ -362,9 +382,9 @@ impl GleanState<'_> {
362382
}
363383
};
364384
if !self.names.contains(&scope) {
365-
self.record_name(scope.clone(), None);
385+
self.record_name(scope.clone());
366386
}
367-
self.record_name(join_names(&scope, name), Some(name.range.start()))
387+
self.record_name_with_position(join_names(&scope, name), name.range.start())
368388
}
369389

370390
fn make_fq_names_for_expr(&self, expr: &Expr) -> Vec<python::Name> {
@@ -806,11 +826,33 @@ impl GleanState<'_> {
806826
decl_infos
807827
}
808828

829+
fn add_xrefs_for_explicit_exports(&mut self, exprs_list: &[Expr]) {
830+
let exports = exprs_list
831+
.iter()
832+
.filter_map(|e| e.as_string_literal_expr())
833+
.map(|str_lit| (str_lit.value.to_str(), str_lit.range()));
834+
835+
for (local_name, source) in exports {
836+
let name = join_names(self.module_name.as_str(), local_name);
837+
if self.names.contains(&name) {
838+
let fqname = self
839+
.import_names
840+
.get(&name)
841+
.map(|n| (**n).clone())
842+
.unwrap_or(name);
843+
844+
self.add_xref(python::XRefViaName {
845+
target: python::Name::new(fqname),
846+
source: to_span(source),
847+
});
848+
}
849+
}
850+
}
851+
809852
fn variable_facts(
810853
&mut self,
811854
expr: &Expr,
812-
range: TextRange,
813-
annotation: Option<&Expr>,
855+
info: &AssignInfo,
814856
ctx: &NodeContext,
815857
next: Option<&Stmt>,
816858
def_infos: &mut Vec<DeclarationInfo>,
@@ -830,15 +872,20 @@ impl GleanState<'_> {
830872
next.and_then(|stmt| Docstring::range_from_stmts(slice::from_ref(stmt)));
831873
def_infos.push(self.variable_info(
832874
fqname,
833-
range,
834-
self.type_info(annotation),
875+
info.range,
876+
self.type_info(info.annotation),
835877
docstring_range,
836878
ctx,
837879
));
880+
881+
if name.id() == &dunder::ALL
882+
&& let Some(value) = info.value
883+
&& let Expr::List(ExprList { elts, .. }) = value
884+
{
885+
self.add_xrefs_for_explicit_exports(elts);
886+
}
838887
}
839-
expr.recurse(&mut |expr| {
840-
self.variable_facts(expr, range, annotation, ctx, next, def_infos)
841-
});
888+
expr.recurse(&mut |expr| self.variable_facts(expr, info, ctx, next, def_infos));
842889
}
843890

844891
fn find_fqname_definition_at_position(&self, position: TextSize) -> Vec<String> {
@@ -862,12 +909,8 @@ impl GleanState<'_> {
862909
resolve_original_name: bool,
863910
) -> DeclarationInfo {
864911
let as_name_fqname = join_names(self.module_name.as_str(), as_name);
865-
self.record_name(from_name.to_owned(), Some(as_name_range.start()));
866-
self.record_name(as_name_fqname.clone(), None);
867-
868912
let from_name_fact = python::Name::new(from_name.to_owned());
869-
let as_name_fact = python::Name::new(as_name_fqname);
870-
913+
let as_name_fact = python::Name::new(as_name_fqname.clone());
871914
let original_name = if resolve_original_name
872915
&& let Some(name) = self
873916
.find_fqname_definition_at_position(from_name_range.start())
@@ -877,6 +920,9 @@ impl GleanState<'_> {
877920
} else {
878921
from_name.to_owned()
879922
};
923+
924+
self.record_name_with_position(from_name.to_owned(), as_name_range.start());
925+
self.record_name_for_import(as_name_fqname, &original_name);
880926
self.add_xref(python::XRefViaName {
881927
target: python::Name::new(original_name),
882928
source: to_span(from_name_range),
@@ -1161,39 +1207,33 @@ impl GleanState<'_> {
11611207
decl_infos.append(&mut func_decl_infos);
11621208
}
11631209
Stmt::Assign(assign) => {
1210+
let info = AssignInfo {
1211+
range: assign.range(),
1212+
annotation: None,
1213+
value: Some(assign.value.as_ref()),
1214+
};
11641215
assign.targets.visit(&mut |target| {
1165-
self.variable_facts(
1166-
target,
1167-
assign.range(),
1168-
None,
1169-
context,
1170-
next,
1171-
&mut decl_infos,
1172-
)
1216+
self.variable_facts(target, &info, context, next, &mut decl_infos)
11731217
});
11741218
self.visit_exprs(&assign.value, container);
11751219
}
11761220
Stmt::AnnAssign(assign) => {
1177-
self.variable_facts(
1178-
&assign.target,
1179-
assign.range(),
1180-
Some(&assign.annotation),
1181-
context,
1182-
next,
1183-
&mut decl_infos,
1184-
);
1221+
let info = AssignInfo {
1222+
range: assign.range(),
1223+
annotation: Some(&assign.annotation),
1224+
value: assign.value.as_ref().map(|v| v.as_ref()),
1225+
};
1226+
self.variable_facts(&assign.target, &info, context, next, &mut decl_infos);
11851227
self.visit_exprs(&assign.annotation, container);
11861228
self.visit_exprs(&assign.value, container);
11871229
}
11881230
Stmt::AugAssign(assign) => {
1189-
self.variable_facts(
1190-
&assign.target,
1191-
assign.range(),
1192-
None,
1193-
context,
1194-
next,
1195-
&mut decl_infos,
1196-
);
1231+
let info = AssignInfo {
1232+
range: assign.range(),
1233+
annotation: None,
1234+
value: Some(assign.value.as_ref()),
1235+
};
1236+
self.variable_facts(&assign.target, &info, context, next, &mut decl_infos);
11971237
self.visit_exprs(&assign.value, container);
11981238
}
11991239
Stmt::Import(import) => {
@@ -1206,14 +1246,12 @@ impl GleanState<'_> {
12061246
}
12071247
Stmt::For(stmt_for) => {
12081248
stmt_for.target.visit(&mut |target| {
1209-
self.variable_facts(
1210-
target,
1211-
target.range(),
1212-
None,
1213-
context,
1214-
next,
1215-
&mut decl_infos,
1216-
)
1249+
let info = AssignInfo {
1250+
range: target.range(),
1251+
annotation: None,
1252+
value: None,
1253+
};
1254+
self.variable_facts(target, &info, context, next, &mut decl_infos)
12171255
});
12181256
self.visit_exprs(&stmt_for.iter, container);
12191257
}
@@ -1228,14 +1266,12 @@ impl GleanState<'_> {
12281266
for item in &stmt_with.items {
12291267
self.visit_exprs(&item.context_expr, container);
12301268
item.optional_vars.visit(&mut |target| {
1231-
self.variable_facts(
1232-
target,
1233-
target.range(),
1234-
None,
1235-
context,
1236-
next,
1237-
&mut decl_infos,
1238-
)
1269+
let info = AssignInfo {
1270+
range: target.range(),
1271+
annotation: None,
1272+
value: None,
1273+
};
1274+
self.variable_facts(target, &info, context, next, &mut decl_infos)
12391275
});
12401276
}
12411277
}
@@ -1281,7 +1317,7 @@ impl Glean {
12811317
let ast = &*transaction.get_ast(handle).unwrap();
12821318
let mut glean_state = GleanState::new(transaction, handle);
12831319

1284-
glean_state.record_name("".to_owned(), None);
1320+
glean_state.record_name("".to_owned());
12851321
let file_language_fact =
12861322
src::FileLanguage::new(glean_state.file_fact(), src::Language::Python);
12871323
let digest_fact = glean_state.digest_fact();

‎pyrefly/lib/report/glean/snapshots/exports_all.json‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -610,6 +610,29 @@
610610
"key": "exports_all.py"
611611
},
612612
"spans": [
613+
{
614+
"length": 3,
615+
"start": 54
616+
}
617+
],
618+
"target": {
619+
"id": 0,
620+
"key": "exports_all.x"
621+
}
622+
}
623+
},
624+
{
625+
"id": 0,
626+
"key": {
627+
"file": {
628+
"id": 0,
629+
"key": "exports_all.py"
630+
},
631+
"spans": [
632+
{
633+
"length": 3,
634+
"start": 59
635+
},
613636
{
614637
"length": 5,
615638
"start": 19
@@ -621,6 +644,25 @@
621644
}
622645
}
623646
},
647+
{
648+
"id": 0,
649+
"key": {
650+
"file": {
651+
"id": 0,
652+
"key": "exports_all.py"
653+
},
654+
"spans": [
655+
{
656+
"length": 3,
657+
"start": 77
658+
}
659+
],
660+
"target": {
661+
"id": 0,
662+
"key": "exports_all.y"
663+
}
664+
}
665+
},
624666
{
625667
"id": 0,
626668
"key": {
@@ -653,6 +695,36 @@
653695
"key": "exports_all.py"
654696
},
655697
"xrefs": [
698+
{
699+
"source": {
700+
"length": 3,
701+
"start": 54
702+
},
703+
"target": {
704+
"id": 0,
705+
"key": "exports_all.x"
706+
}
707+
},
708+
{
709+
"source": {
710+
"length": 3,
711+
"start": 59
712+
},
713+
"target": {
714+
"id": 0,
715+
"key": "typing.Union"
716+
}
717+
},
718+
{
719+
"source": {
720+
"length": 3,
721+
"start": 77
722+
},
723+
"target": {
724+
"id": 0,
725+
"key": "exports_all.y"
726+
}
727+
},
656728
{
657729
"source": {
658730
"length": 5,

0 commit comments

Comments
 (0)