Preserve attributes for imports_granularity=Item

Fixes #5030
rust-lang · Feb 25, 2022 · ac910f4 · ac910f4
1 parent 89ca3f3
commit ac910f4
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/src/imports.rs b/src/imports.rs
@@ -108,6 +108,12 @@ pub(crate) struct UseTree {
     attrs: Option<Vec<ast::Attribute>>,
 }
 
+#[derive(Copy, Clone)]
+enum Attributes {
+    Preserve,
+    Discard,
+}
+
 impl PartialEq for UseTree {
     fn eq(&self, other: &UseTree) -> bool {
         self.path == other.path
@@ -190,7 +196,7 @@ pub(crate) fn merge_use_trees(use_trees: Vec<UseTree>, merge_by: SharedPrefix) -
             continue;
         }
 
-        for flattened in use_tree.flatten() {
+        for flattened in use_tree.flatten(Attributes::Discard) {
             if let Some(tree) = result
                 .iter_mut()
                 .find(|tree| tree.share_prefix(&flattened, merge_by))
@@ -207,7 +213,7 @@ pub(crate) fn merge_use_trees(use_trees: Vec<UseTree>, merge_by: SharedPrefix) -
 pub(crate) fn flatten_use_trees(use_trees: Vec<UseTree>) -> Vec<UseTree> {
     use_trees
         .into_iter()
-        .flat_map(UseTree::flatten)
+        .flat_map(|tree| tree.flatten(Attributes::Preserve))
         .map(|mut tree| {
             // If a path ends in `::self`, rewrite it to `::{self}`.
             if let Some(UseSegment::Slf(..)) = tree.path.last() {
@@ -587,7 +593,7 @@ impl UseTree {
         }
     }
 
-    fn flatten(self) -> Vec<UseTree> {
+    fn flatten(self, preserve_attrs: Attributes) -> Vec<UseTree> {
         if self.path.is_empty() {
             return vec![self];
         }
@@ -601,15 +607,19 @@ impl UseTree {
                 let prefix = &self.path[..self.path.len() - 1];
                 let mut result = vec![];
                 for nested_use_tree in list {
-                    for flattend in &mut nested_use_tree.clone().flatten() {
+                    for flattend in &mut nested_use_tree.clone().flatten(preserve_attrs) {
                         let mut new_path = prefix.to_vec();
                         new_path.append(&mut flattend.path);
                         result.push(UseTree {
                             path: new_path,
                             span: self.span,
                             list_item: None,
                             visibility: self.visibility.clone(),
-                            attrs: None,
+                            // only retain attributes for `ImportGranularity::Item`
+                            attrs: match preserve_attrs {
+                                Attributes::Preserve => self.attrs.clone(),
+                                Attributes::Discard => None,
+                            },
                         });
                     }
                 }
@@ -1228,12 +1238,12 @@ mod test {
     #[test]
     fn test_use_tree_flatten() {
         assert_eq!(
-            parse_use_tree("a::b::{c, d, e, f}").flatten(),
+            parse_use_tree("a::b::{c, d, e, f}").flatten(Attributes::Discard),
             parse_use_trees!("a::b::c", "a::b::d", "a::b::e", "a::b::f",)
         );
 
         assert_eq!(
-            parse_use_tree("a::b::{c::{d, e, f}, g, h::{i, j, k}}").flatten(),
+            parse_use_tree("a::b::{c::{d, e, f}, g, h::{i, j, k}}").flatten(Attributes::Discard),
             parse_use_trees![
                 "a::b::c::d",
                 "a::b::c::e",

diff --git a/tests/source/issue-5030.rs b/tests/source/issue-5030.rs
@@ -0,0 +1,7 @@
+// rustfmt-imports_granularity: Item
+
+#[cfg(feature = "foo")]
+use std::collections::{
+    HashMap,
+    HashSet,
+};
diff --git a/tests/target/issue-5030.rs b/tests/target/issue-5030.rs
@@ -0,0 +1,6 @@
+// rustfmt-imports_granularity: Item
+
+#[cfg(feature = "foo")]
+use std::collections::HashMap;
+#[cfg(feature = "foo")]
+use std::collections::HashSet;