Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
235 changes: 203 additions & 32 deletions pyrefly/lib/commands/coverage/collect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use pyrefly_python::dunder;
use pyrefly_python::ignore::Ignore;
use pyrefly_python::module::Module;
use pyrefly_python::module_name::ModuleName;
use pyrefly_python::module_path::ModulePathDetails;
use pyrefly_python::module_path::ModuleStyle;
use pyrefly_python::nesting_context::NestingContext;
use pyrefly_python::short_identifier::ShortIdentifier;
Expand Down Expand Up @@ -325,7 +326,7 @@ fn trace_export_origin(
handle: &Handle,
mut cur_name: Name,
transaction: &Transaction,
) -> Option<String> {
) -> Option<(Handle, String)> {
let mut seen = SmallSet::new();
let mut cur_handle = handle.clone();

Expand All @@ -337,7 +338,7 @@ fn trace_export_origin(

match transaction.get_exports(&cur_handle).get(&cur_name) {
Some(ExportLocation::ThisModule(_)) | None => {
return Some(format!("{module_name}.{cur_name}"));
return Some((cur_handle, format!("{module_name}.{cur_name}")));
}
Some(ExportLocation::OtherModule(other_module, alias)) => {
if let Some(alias) = alias {
Expand All @@ -351,39 +352,118 @@ fn trace_export_origin(
}
}

fn public_exported_names(handle: &Handle, transaction: &Transaction) -> Vec<Name> {
let exports_data = transaction.get_exports_data(handle);
let exports = transaction.get_exports(handle);

// prioritize `__all__` if present, otherwise local defs + `import x as x`
if let Some(all_iter) = exports_data.get_explicit_dunder_all_names_iter() {
all_iter.cloned().collect()
} else {
exports
.iter()
.filter_map(|(name, loc)| {
let is_local = matches!(loc, ExportLocation::ThisModule(_));
let is_reexport = exports_data.is_explicit_reexport(name);
(is_public_name(name.as_str()) && (is_local || is_reexport)).then_some(name.clone())
})
.collect()
}
}

fn same_handle(left: &Handle, right: &Handle) -> bool {
left.module() == right.module() && left.path() == right.path()
}

fn contains_handle(handles: &[Handle], candidate: &Handle) -> bool {
handles.iter().any(|handle| same_handle(handle, candidate))
}
Comment on lines +374 to +380

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handle already derives PartialEq, Eq, and Hash, so these are not needed, and even hurt performance (+ O(n^2)). Assuming that you ignored sys_info by mistake, then both functions can be removed and a SmallSet of Handle could be used instead.


fn is_project_reexport_origin(
public_handle: &Handle,
origin_handle: &Handle,
config_finder: &ConfigFinder,
) -> bool {
if origin_handle.path().is_bundled()
|| !matches!(
origin_handle.path().details(),
ModulePathDetails::FileSystem(_) | ModulePathDetails::Memory(_)
)
{
return false;
}

let origin_path = origin_handle.path().as_path();
let config = config_finder.python_file(public_handle.module_kind(), public_handle.path());
if config
.site_package_path()
.any(|site_package| origin_path.starts_with(site_package))
{
return false;
}

if public_handle.module() == ModuleName::unknown() {
return public_handle
.path()
.as_path()
.parent()
.is_some_and(|parent| origin_path.starts_with(parent));
}

if let (Some(public_root), Some(origin_root)) = (
public_handle.path().root_of(public_handle.module()),
origin_handle.path().root_of(origin_handle.module()),
) {
return public_root == origin_root;
}

config
.search_path()
.any(|search_root| origin_path.starts_with(search_root))
}
Comment on lines +382 to +423

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are a whole bunch of unmotivated heuristics here: a bundle check, path-details matching, site-package rejection, __unknown__ special-casing, parent-dir special-casing, double root_of comparisons, and then a search-path fallback.

Did you verify that each of these heuristics is actually needed here? If so, then please explain why exactly that is (as a comment). If not, then it should be removed.


fn public_reexport_origin_handles(
handles: &[Handle],
transaction: &Transaction,
config_finder: &ConfigFinder,
) -> Vec<Handle> {
let mut origin_handles = Vec::new();
for handle in handles.iter().filter(|h| is_public_module(h.module())) {
for name in public_exported_names(handle, transaction) {
if EXCLUDED_MODULE_DUNDERS.contains(&name.as_str()) {
continue;
}
let Some((origin_handle, _)) = trace_export_origin(handle, name, transaction) else {
continue;
};
if same_handle(handle, &origin_handle)
|| contains_handle(handles, &origin_handle)
|| contains_handle(&origin_handles, &origin_handle)
|| !is_project_reexport_origin(handle, &origin_handle, config_finder)
{
continue;
}
origin_handles.push(origin_handle);
}
}
origin_handles
}
Comment on lines +425 to +450

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This iterates the same public modules as compute_public_fqns below, applies the same EXCLUDED_MODULE_DUNDERS filter, and traces the same export. You could produce both the origin-handle list and the FQN set in a single pass, and avoid this introduced code duplication.


/// Collect origin FQNs of all publicly exported names across public modules.
fn compute_public_fqns(handles: &[Handle], transaction: &Transaction) -> HashSet<String> {
handles
.iter()
.filter(|h| is_public_module(h.module()))
.flat_map(|handle| {
let exports_data = transaction.get_exports_data(handle);
let exports = transaction.get_exports(handle);

// prioritize `__all__` if present, otherwise local defs + `import x as x`
let names: Vec<Name> =
if let Some(all_iter) = exports_data.get_explicit_dunder_all_names_iter() {
all_iter.cloned().collect()
} else {
exports
.iter()
.filter_map(|(name, loc)| {
let is_local = matches!(loc, ExportLocation::ThisModule(_));
let is_reexport = exports_data.is_explicit_reexport(name);
(is_public_name(name.as_str()) && (is_local || is_reexport))
.then_some(name.clone())
})
.collect()
};

// emit both the local FQN and the traced origin FQN so a file-scoped run matches
// whichever module was requested
names
public_exported_names(handle, transaction)
.into_iter()
.filter(|n| !EXCLUDED_MODULE_DUNDERS.contains(&n.as_str()))
.flat_map(move |name| {
let local = format!("{}.{}", handle.module(), name);
let origin = trace_export_origin(handle, name, transaction);
let origin = trace_export_origin(handle, name, transaction)
.map(|(_, origin_fqn)| origin_fqn);
std::iter::once(local).chain(origin)
})
})
Expand Down Expand Up @@ -1658,23 +1738,41 @@ pub fn collect_module_reports(
let mut module_reports: Vec<ModuleReport> = Vec::new();
let mut errors: Vec<Error> = Vec::new();
transaction.run(handles.as_slice(), Require::Everything, None);
let public_fqns = public_only.then(|| compute_public_fqns(&handles, transaction));
let mut report_handles = handles.clone();
let public_origin_handles = if public_only {
public_reexport_origin_handles(&handles, transaction, holder.as_ref().config_finder())
} else {
Vec::new()
};
let public_fqns = if public_only {
for origin_handle in public_origin_handles.iter().cloned() {
if !contains_handle(&report_handles, &origin_handle) {
report_handles.push(origin_handle);
}
}
if report_handles.len() != handles.len() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if report_handles.len() != handles.len() {
if !public_origin_handles.is_empty() {

transaction.run(report_handles.as_slice(), Require::Everything, None);
}
Some(compute_public_fqns(&handles, transaction))
} else {
None
};

let shadowed = if prefer_stubs {
py_paths_shadowed_by_pyi(&handles)
py_paths_shadowed_by_pyi(&report_handles)
} else {
SmallSet::new()
};

// Map each .pyi to its corresponding .py: first co-located,
// then by module-name lookup in site-package-path.
let pyi_to_py: HashMap<PathBuf, Handle> = if prefer_stubs {
let py_by_path: HashMap<PathBuf, &Handle> = handles
let py_by_path: HashMap<PathBuf, &Handle> = report_handles
.iter()
.filter(|h| !h.path().is_interface())
.map(|h| (h.path().as_path().to_path_buf(), h))
.collect();
let mut map: HashMap<PathBuf, Handle> = handles
let mut map: HashMap<PathBuf, Handle> = report_handles
.iter()
.filter(|h| h.path().is_interface())
.filter_map(|h| {
Expand All @@ -1686,7 +1784,7 @@ pub fn collect_module_reports(
.collect();
// Fall back to site-package-path for stubs-only packages.
let mut external_handles = Vec::new();
for h in handles.iter().filter(|h| h.path().is_interface()) {
for h in report_handles.iter().filter(|h| h.path().is_interface()) {
let pyi_path = h.path().as_path().to_path_buf();
if map.contains_key(&pyi_path) {
continue;
Expand Down Expand Up @@ -1728,7 +1826,7 @@ pub fn collect_module_reports(

// gh-3632: skip files whose module name isn't importable (shadowed parent).
let module = handle.module();
if module != ModuleName::unknown() {
if module != ModuleName::unknown() && !contains_handle(&public_origin_handles, handle) {
let config = config_finder.python_file(handle.module_kind(), handle.path());
find_import_filtered(&config, module, None, None, &dir_cache, None).finding()?;
}
Expand Down Expand Up @@ -1787,7 +1885,7 @@ pub fn collect_module_reports(
};
let collected: Vec<(ModuleReport, Vec<Error>)> = holder
.as_ref()
.install(|| handles.par_iter().filter_map(collect_one).collect());
.install(|| report_handles.par_iter().filter_map(collect_one).collect());
for (report, module_errors) in collected {
module_reports.push(report);
errors.extend(module_errors);
Expand All @@ -1814,11 +1912,18 @@ mod tests {

use dupe::Dupe;
use pyrefly_build::handle::Handle;
use pyrefly_config::args::ConfigOverrideArgs;
use pyrefly_config::config::ConfigFile;
use pyrefly_python::module_name::ModuleName;
use pyrefly_python::module_path::ModulePath;
use pyrefly_python::sys_info::SysInfo;
use pyrefly_util::globs::FilteredGlobs;
use pyrefly_util::globs::Globs;
use pyrefly_util::globs::HiddenDirFilter;
use pyrefly_util::thread_pool::TEST_THREAD_COUNT;

use super::*;
use crate::commands::config_finder::default_config_finder_with_overrides;
use crate::state::require::Require;
use crate::test::util::TestEnv;

Expand Down Expand Up @@ -2098,8 +2203,6 @@ mod tests {
/// gh-3632: skip a file shadowed by a same-named module (it can't be imported).
#[test]
fn test_report_skips_unimportable_shadowed_module() {
use pyrefly_config::config::ConfigFile;

let site = tempfile::TempDir::new().unwrap();
let dir = site.path();
std::fs::write(dir.join("lapack_lite.pyi"), "def f() -> int: ...\n").unwrap();
Expand All @@ -2121,6 +2224,74 @@ mod tests {
assert!(!find("lapack_lite.fortran"), "shadowed file skipped");
}

fn collect_public_reexport_sample(project_excludes: Vec<String>) -> Vec<ModuleReport> {
let tempdir = tempfile::TempDir::new().unwrap();
let root = tempdir.path();
let package = root.join("foo");
std::fs::create_dir(&package).unwrap();
std::fs::write(
package.join("__init__.py"),
"from foo._something_private import bar\n__all__ = ['bar']\n",
)
.unwrap();
std::fs::write(
package.join("_something_private.py"),
"def bar(a) -> None:\n return None\n",
)
.unwrap();

let files_to_check = Globs::new(vec![package.display().to_string()]).unwrap();
let project_excludes = Globs::new_with_root(root, project_excludes).unwrap();
let roots = files_to_check.roots();
let files_to_check = Box::new(FilteredGlobs::new(
files_to_check,
project_excludes.clone(),
None,
HiddenDirFilter::RelativeTo(roots.clone()),
));
let config_finder =
default_config_finder_with_overrides(ConfigOverrideArgs::default(), false, None);
if let Some(root) = roots.first() {
config_finder.directory(root);
}
let (reports, errors) = collect_module_reports(
files_to_check,
config_finder,
true,
None,
true,
None,
TEST_THREAD_COUNT,
)
.unwrap();
assert!(errors.is_empty());
reports
}

#[test]
fn test_public_only_keeps_project_excluded_public_reexport_origin() {
let reports = collect_public_reexport_sample(vec!["foo/_something_private.py".to_owned()]);
let private_report = reports
.iter()
.find(|report| report.name == "foo._something_private")
.expect("publicly re-exported origin should still be reported");

assert!(
private_report
.symbol_reports
.iter()
.any(|symbol| symbol.name() == "foo._something_private.bar")
);
assert!(private_report.coverage < 100.0);
}

#[test]
fn test_public_only_excludes_origin_when_public_reexport_is_excluded() {
let reports = collect_public_reexport_sample(vec!["foo/__init__.py".to_owned()]);

assert!(reports.is_empty());
}

/// When both test.py and test.pyi exist, the .py file is shadowed.
#[test]
fn test_pyi_shadows_py_in_report() {
Expand Down
Loading