diff --git a/src/rust/engine/dep_inference/src/python/mod.rs b/src/rust/engine/dep_inference/src/python/mod.rs index 26a083e550a8..f80a564a6825 100644 --- a/src/rust/engine/dep_inference/src/python/mod.rs +++ b/src/rust/engine/dep_inference/src/python/mod.rs @@ -50,7 +50,11 @@ pub fn get_dependencies( } let mut new_key_parts = path_parts[0..((path_parts.len() - level) + 1)].to_vec(); - new_key_parts.push(nonrelative); + if !nonrelative.is_empty() { + // an import like `from .. import *` can end up with key == '..', and hence nonrelative == ""; + // the result should just be the raw parent traversal, without a suffix part + new_key_parts.push(nonrelative); + } let old_value = import_map.remove(&key).unwrap(); import_map.insert(new_key_parts.join("."), old_value); @@ -152,6 +156,7 @@ impl ImportCollector<'_> { fn normalize_import_node(node: tree_sitter::Node) -> Option { match node.kind_id() { KindID::ALIASED_IMPORT => node.named_child(0), + KindID::WILDCARD_IMPORT => None, KindID::ERROR => None, _ => Some(node), } @@ -224,8 +229,27 @@ impl Visitor for ImportCollector<'_> { fn visit_import_from_statement(&mut self, node: tree_sitter::Node) -> ChildBehavior { if !self.is_pragma_ignored(node) { + // the grammar is something like `from $module_name import $($name),* | '*'`, where $... is a field + // name. + let module_name = node + .child_by_field_name("module_name") + .expect("`from ... import ...` must have module_name"); + + let mut any_names = false; for child in node.children_by_field_name("name", &mut node.walk()) { - self.insert_import(node.named_child(0).unwrap(), Some(child), false); + self.insert_import(module_name, Some(child), false); + any_names = true; + } + + if !any_names { + // There's no names (i.e. it's probably not `from ... import some, names`), let's look for + // the * in a wildcard import. (It doesn't have a field name, so we have to search for it + // manually.) + for child in node.children(&mut node.walk()) { + if child.kind_id() == KindID::WILDCARD_IMPORT { + self.insert_import(module_name, Some(child), false); + } + } } } ChildBehavior::Ignore diff --git a/src/rust/engine/dep_inference/src/python/tests.rs b/src/rust/engine/dep_inference/src/python/tests.rs index 7599f65759a3..27e87c77328e 100644 --- a/src/rust/engine/dep_inference/src/python/tests.rs +++ b/src/rust/engine/dep_inference/src/python/tests.rs @@ -109,6 +109,7 @@ fn pragma_ignore() { assert_imports("import a.b # pants: no-infer-dep", &[]); assert_imports("import a.b as d # pants: no-infer-dep", &[]); assert_imports("from a import b # pants: no-infer-dep", &[]); + assert_imports("from a import * # pants: no-infer-dep", &[]); assert_imports("from a import b, c # pants: no-infer-dep", &[]); assert_imports("from a import b, c as d # pants: no-infer-dep", &[]); assert_imports( @@ -188,6 +189,12 @@ fn pragma_ignore() { )", &[], ); + assert_imports( + r" + from a.b import \ + * # pants: no-infer-dep", + &[], + ); } #[test] @@ -532,11 +539,16 @@ fn syntax_errors_and_other_fun() { assert_imports("import .b", &[]); assert_imports("import a....b", &["a....b"]); assert_imports("import a.", &[]); + assert_imports("import *", &[]); assert_imports("from a import", &[]); assert_imports("from a imp x", &[]); assert_imports("from from import a as .as", &[]); assert_imports("from a import ......g", &["a.g"]); assert_imports("from a. import b", &[]); + assert_imports("from a import *, b", &["a"]); + assert_imports("from a import b, *", &["a.b"]); + assert_imports("from a import (*)", &[]); + assert_imports("from * import b", &["b"]); assert_imports("try:...\nexcept:import a", &["a"]); assert_imports("try:...\nexcept 1:import a", &["a"]); assert_imports("try:...\nexcept x=1:import a", &["a"]);