Rewrite xml.etree.cElementTree to xml.etree.ElementTree

README.md

@@ -671,6 +671,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP019 | TypingTextStrAlias | `typing.Text` is deprecated, use `str` | 🛠 |
 | UP020 | OpenAlias | Use builtin `open` | 🛠 |
 | UP021 | ReplaceUniversalNewlines | `universal_newlines` is deprecated, use `text` | 🛠 |
+| UP023 | RewriteCElementTree | `cElementTree` is deprecated, use `ElementTree` | 🛠 |
 
 ### pep8-naming (N)
 

resources/test/fixtures/pyupgrade/UP023.py

@@ -0,0 +1,31 @@
+# These two imports have something after cElementTree, so they should be fixed.
+from xml.etree.cElementTree import XML, Element, SubElement
+import xml.etree.cElementTree as ET
+
+# Weird spacing should not cause issues.
+from   xml.etree.cElementTree    import  XML
+import    xml.etree.cElementTree       as      ET
+
+# Multi line imports should also work fine.
+from xml.etree.cElementTree import (
+    XML,
+    Element,
+    SubElement,
+)
+if True:
+    import xml.etree.cElementTree as ET
+    from xml.etree import cElementTree as CET
+
+from xml.etree import cElementTree as ET
+
+import contextlib, xml.etree.cElementTree as ET
+
+# This should fix the second, but not the first invocation.
+import xml.etree.cElementTree, xml.etree.cElementTree as ET
+
+# The below items should NOT be changed.
+import xml.etree.cElementTree
+
+from .xml.etree.cElementTree import XML
+
+from xml.etree import cElementTree

ruff.schema.json

@@ -881,6 +881,7 @@
         "UP02",
         "UP020",
         "UP021",
+        "UP023",
         "W",
         "W2",
         "W29",

src/checkers/ast.rs

@@ -650,6 +650,9 @@ where
                         ));
                     }
                 }
+                if self.settings.enabled.contains(&CheckCode::UP023) {
+                    pyupgrade::plugins::replace_c_element_tree(self, stmt);
+                }
 
                 for alias in names {
                     if alias.node.name.contains('.') && alias.node.asname.is_none() {
@@ -819,6 +822,9 @@ where
             } => {
                 // Track `import from` statements, to ensure that we can correctly attribute
                 // references like `from typing import Union`.
+                if self.settings.enabled.contains(&CheckCode::UP023) {
+                    pyupgrade::plugins::replace_c_element_tree(self, stmt);
+                }
                 if level.map(|level| level == 0).unwrap_or(true) {
                     if let Some(module) = module {
                         self.from_imports
@@ -1552,9 +1558,6 @@ where
                     pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
                 }
 
-                if self.settings.enabled.contains(&CheckCode::UP019) {
-                    pyupgrade::plugins::typing_text_str_alias(self, expr);
-                }
                 if self.settings.enabled.contains(&CheckCode::UP016) {
                     pyupgrade::plugins::remove_six_compat(self, expr);
                 }
@@ -1564,7 +1567,9 @@ where
                 {
                     pyupgrade::plugins::datetime_utc_alias(self, expr);
                 }
-
+                if self.settings.enabled.contains(&CheckCode::UP019) {
+                    pyupgrade::plugins::typing_text_str_alias(self, expr);
+                }
                 if self.settings.enabled.contains(&CheckCode::YTT202) {
                     flake8_2020::plugins::name_or_attribute(self, expr);
                 }

src/checks.rs

@@ -229,6 +229,7 @@ pub enum CheckCode {
     UP019,
     UP020,
     UP021,
+    UP023,
     // pydocstyle
     D100,
     D101,
@@ -844,6 +845,7 @@ pub enum CheckKind {
     NativeLiterals,
     OpenAlias,
     ReplaceUniversalNewlines,
+    RewriteCElementTree,
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1223,6 +1225,7 @@ impl CheckCode {
             CheckCode::UP019 => CheckKind::TypingTextStrAlias,
             CheckCode::UP020 => CheckKind::OpenAlias,
             CheckCode::UP021 => CheckKind::ReplaceUniversalNewlines,
+            CheckCode::UP023 => CheckKind::RewriteCElementTree,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1647,6 +1650,7 @@ impl CheckCode {
             CheckCode::UP019 => CheckCategory::Pyupgrade,
             CheckCode::UP020 => CheckCategory::Pyupgrade,
             CheckCode::UP021 => CheckCategory::Pyupgrade,
+            CheckCode::UP023 => CheckCategory::Pyupgrade,
             CheckCode::W292 => CheckCategory::Pycodestyle,
             CheckCode::W605 => CheckCategory::Pycodestyle,
             CheckCode::YTT101 => CheckCategory::Flake82020,
@@ -1862,6 +1866,7 @@ impl CheckKind {
             CheckKind::TypingTextStrAlias => &CheckCode::UP019,
             CheckKind::OpenAlias => &CheckCode::UP020,
             CheckKind::ReplaceUniversalNewlines => &CheckCode::UP021,
+            CheckKind::RewriteCElementTree => &CheckCode::UP023,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(..) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(..) => &CheckCode::D410,
@@ -2595,6 +2600,9 @@ impl CheckKind {
             CheckKind::ReplaceUniversalNewlines => {
                 "`universal_newlines` is deprecated, use `text`".to_string()
             }
+            CheckKind::RewriteCElementTree => {
+                "`cElementTree` is deprecated, use `ElementTree`".to_string()
+            }
             CheckKind::ConvertNamedTupleFunctionalToClass(name) => {
                 format!("Convert `{name}` from `NamedTuple` functional to class syntax")
             }
@@ -3040,6 +3048,7 @@ impl CheckKind {
                 | CheckKind::OpenAlias
                 | CheckKind::NewLineAfterLastParagraph
                 | CheckKind::ReplaceUniversalNewlines
+                | CheckKind::RewriteCElementTree
                 | CheckKind::NewLineAfterSectionName(..)
                 | CheckKind::NoBlankLineAfterFunction(..)
                 | CheckKind::NoBlankLineBeforeClass(..)

src/checks_gen.rs

@@ -532,6 +532,7 @@ pub enum CheckCodePrefix {
     UP02,
     UP020,
     UP021,
+    UP023,
     W,
     W2,
     W29,
@@ -759,6 +760,7 @@ impl CheckCodePrefix {
                 CheckCode::UP019,
                 CheckCode::UP020,
                 CheckCode::UP021,
+                CheckCode::UP023,
                 CheckCode::D100,
                 CheckCode::D101,
                 CheckCode::D102,
@@ -2415,6 +2417,7 @@ impl CheckCodePrefix {
                     CheckCode::UP019,
                     CheckCode::UP020,
                     CheckCode::UP021,
+                    CheckCode::UP023,
                 ]
             }
             CheckCodePrefix::U0 => {
@@ -2445,6 +2448,7 @@ impl CheckCodePrefix {
                     CheckCode::UP019,
                     CheckCode::UP020,
                     CheckCode::UP021,
+                    CheckCode::UP023,
                 ]
             }
             CheckCodePrefix::U00 => {
@@ -2659,6 +2663,7 @@ impl CheckCodePrefix {
                 CheckCode::UP019,
                 CheckCode::UP020,
                 CheckCode::UP021,
+                CheckCode::UP023,
             ],
             CheckCodePrefix::UP0 => vec![
                 CheckCode::UP001,
@@ -2681,6 +2686,7 @@ impl CheckCodePrefix {
                 CheckCode::UP019,
                 CheckCode::UP020,
                 CheckCode::UP021,
+                CheckCode::UP023,
             ],
             CheckCodePrefix::UP00 => vec![
                 CheckCode::UP001,
@@ -2722,9 +2728,10 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP017 => vec![CheckCode::UP017],
             CheckCodePrefix::UP018 => vec![CheckCode::UP018],
             CheckCodePrefix::UP019 => vec![CheckCode::UP019],
-            CheckCodePrefix::UP02 => vec![CheckCode::UP020, CheckCode::UP021],
+            CheckCodePrefix::UP02 => vec![CheckCode::UP020, CheckCode::UP021, CheckCode::UP023],
             CheckCodePrefix::UP020 => vec![CheckCode::UP020],
             CheckCodePrefix::UP021 => vec![CheckCode::UP021],
+            CheckCodePrefix::UP023 => vec![CheckCode::UP023],
             CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
             CheckCodePrefix::W2 => vec![CheckCode::W292],
             CheckCodePrefix::W29 => vec![CheckCode::W292],
@@ -3288,6 +3295,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP02 => SuffixLength::Two,
             CheckCodePrefix::UP020 => SuffixLength::Three,
             CheckCodePrefix::UP021 => SuffixLength::Three,
+            CheckCodePrefix::UP023 => SuffixLength::Three,
             CheckCodePrefix::W => SuffixLength::Zero,
             CheckCodePrefix::W2 => SuffixLength::One,
             CheckCodePrefix::W29 => SuffixLength::Two,

src/pyupgrade/mod.rs

@@ -40,6 +40,7 @@ mod tests {
     #[test_case(CheckCode::UP018, Path::new("UP018.py"); "UP018")]
     #[test_case(CheckCode::UP019, Path::new("UP019.py"); "UP019")]
     #[test_case(CheckCode::UP021, Path::new("UP021.py"); "UP021")]
+    #[test_case(CheckCode::UP023, Path::new("UP023.py"); "UP023")]
     fn checks(check_code: CheckCode, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
         let mut checks = test_path(

src/pyupgrade/plugins/mod.rs

@@ -7,6 +7,7 @@ pub use open_alias::open_alias;
 pub use redundant_open_modes::redundant_open_modes;
 pub use remove_six_compat::remove_six_compat;
 pub use replace_universal_newlines::replace_universal_newlines;
+pub use rewrite_c_element_tree::replace_c_element_tree;
 pub use super_call_with_parameters::super_call_with_parameters;
 pub use type_of_primitive::type_of_primitive;
 pub use typing_text_str_alias::typing_text_str_alias;
@@ -27,6 +28,7 @@ mod open_alias;
 mod redundant_open_modes;
 mod remove_six_compat;
 mod replace_universal_newlines;
+mod rewrite_c_element_tree;
 mod super_call_with_parameters;
 mod type_of_primitive;
 mod typing_text_str_alias;

src/pyupgrade/plugins/rewrite_c_element_tree.rs

@@ -0,0 +1,57 @@
+use rustpython_ast::{Located, Stmt, StmtKind};
+
+use crate::ast::types::Range;
+use crate::autofix::Fix;
+use crate::checkers::ast::Checker;
+use crate::checks::{Check, CheckKind};
+
+fn add_check_for_node<T>(checker: &mut Checker, node: &Located<T>) {
+    let mut check = Check::new(CheckKind::RewriteCElementTree, Range::from_located(node));
+    if checker.patch(check.kind.code()) {
+        let contents = checker
+            .locator
+            .slice_source_code_range(&Range::from_located(node));
+        check.amend(Fix::replacement(
+            contents.replacen("cElementTree", "ElementTree", 1),
+            node.location,
+            node.end_location.unwrap(),
+        ));
+    }
+    checker.add_check(check);
+}
+
+/// UP023
+pub fn replace_c_element_tree(checker: &mut Checker, stmt: &Stmt) {
+    match &stmt.node {
+        StmtKind::Import { names } => {
+            // Ex) `import xml.etree.cElementTree as ET`
+            for name in names {
+                if name.node.name == "xml.etree.cElementTree" && name.node.asname.is_some() {
+                    add_check_for_node(checker, name);
+                }
+            }
+        }
+        StmtKind::ImportFrom {
+            module,
+            names,
+            level,
+        } => {
+            if level.map_or(false, |level| level > 0) {
+                // Ex) `import .xml.etree.cElementTree as ET`
+            } else if let Some(module) = module {
+                if module == "xml.etree.cElementTree" {
+                    // Ex) `from xml.etree.cElementTree import XML`
+                    add_check_for_node(checker, stmt);
+                } else if module == "xml.etree" {
+                    // Ex) `from xml.etree import cElementTree as ET`
+                    for name in names {
+                        if name.node.name == "cElementTree" && name.node.asname.is_some() {
+                            add_check_for_node(checker, name);
+                        }
+                    }
+                }
+            }
+        }
+        _ => unreachable!("Expected StmtKind::Import | StmtKind::ImportFrom"),
+    }
+}