Implement U016: Remove six compatibility code (#1013)

README.md

@@ -618,6 +618,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP013 | ConvertTypedDictFunctionalToClass | Convert `...` from `TypedDict` functional to class syntax | 🛠 |
 | UP014 | ConvertNamedTupleFunctionalToClass | Convert `...` from `NamedTuple` functional to class syntax | 🛠 |
 | UP015 | RedundantOpenModes | Unnecessary open mode parameters | 🛠 |
+| UP016 | RemoveSixCompat | Unnecessary `six` compatibility usage | 🛠 |
 
 ### pep8-naming (N)
 

resources/test/fixtures/pyupgrade/UP016.py

@@ -0,0 +1,87 @@
+# Replace names by built-in names, whether namespaced or not
+# https://github.com/search?q=%22from+six+import%22&type=code
+import six
+from six.moves import map  # No need
+from six import text_type
+
+six.text_type  # str
+six.binary_type  # bytes
+six.class_types  # (type,)
+six.string_types  # (str,)
+six.integer_types  # (int,)
+six.unichr  # chr
+six.iterbytes  # iter
+six.print_(...)  # print(...)
+six.exec_(c, g, l)  # exec(c, g, l)
+six.advance_iterator(it)  # next(it)
+six.next(it)  # next(it)
+six.callable(x)  # callable(x)
+six.moves.range(x)  # range(x)
+six.moves.xrange(x)  # range(x)
+isinstance(..., six.class_types)  # isinstance(..., type)
+issubclass(..., six.integer_types)  # issubclass(..., int)
+isinstance(..., six.string_types)  # isinstance(..., str)
+
+# Replace call on arg by method call on arg
+six.iteritems(dct)  # dct.items()
+six.iterkeys(dct)  # dct.keys()
+six.itervalues(dct)  # dct.values()
+six.viewitems(dct)  # dct.items()
+six.viewkeys(dct)  # dct.keys()
+six.viewvalues(dct)  # dct.values()
+six.assertCountEqual(self, a1, a2)  # self.assertCountEqual(a1, a2)
+six.assertRaisesRegex(self, e, r, fn)  # self.assertRaisesRegex(e, r, fn)
+six.assertRegex(self, s, r)  # self.assertRegex(s, r)
+
+# Replace call on arg by arg attribute
+six.get_method_function(meth)  # meth.__func__
+six.get_method_self(meth)  # meth.__self__
+six.get_function_closure(fn)  # fn.__closure__
+six.get_function_code(fn)  # fn.__code__
+six.get_function_defaults(fn)  # fn.__defaults__
+six.get_function_globals(fn)  # fn.__globals__
+
+# Replace by string literal
+six.b("...")  # b'...'
+six.u("...")  # '...'
+six.ensure_binary("...")  # b'...'
+six.ensure_str("...")  # '...'
+six.ensure_text("...")  # '...'
+six.b(string)  # no change
+
+# Replace by simple expression
+six.get_unbound_function(meth)  # meth
+six.create_unbound_method(fn, cls)  # fn
+
+# Raise exception
+six.raise_from(exc, exc_from)  # raise exc from exc_from
+six.reraise(tp, exc, tb)  # raise exc.with_traceback(tb)
+six.reraise(*sys.exc_info())  # raise
+
+# Int / Bytes conversion
+six.byte2int(bs)  # bs[0]
+six.indexbytes(bs, i)  # bs[i]
+six.int2byte(i)  # bytes((i, ))
+
+# Special cases for next calls
+next(six.iteritems(dct))  # next(iter(dct.items()))
+next(six.iterkeys(dct))  # next(iter(dct.keys()))
+next(six.itervalues(dct))  # next(iter(dct.values()))
+
+# TODO: To implement
+
+
+# Rewrite classes
+@six.python_2_unicode_compatible  # Remove
+class C(six.Iterator):
+    pass  # class C: pass
+
+
+class C(six.with_metaclass(M, B)):
+    pass  # class C(B, metaclass=M): pass
+
+
+# class C(B, metaclass=M): pass
+@six.add_metaclass(M)
+class C(B):
+    pass

src/ast/helpers.rs

@@ -11,6 +11,16 @@ use rustpython_parser::lexer::Tok;
 use crate::ast::types::Range;
 use crate::SourceCodeLocator;
 
+/// Create an `Expr` with default location from an `ExprKind`.
+pub fn create_expr(node: ExprKind) -> Expr {
+    Expr::new(Location::default(), Location::default(), node)
+}
+
+/// Create a `Stmt` with a default location from a `StmtKind`.
+pub fn create_stmt(node: StmtKind) -> Stmt {
+    Stmt::new(Location::default(), Location::default(), node)
+}
+
 fn collect_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
     match &expr.node {
         ExprKind::Call { func, .. } => {

src/check_ast.rs

@@ -1487,6 +1487,10 @@ where
                     pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
                 }
 
+                if self.settings.enabled.contains(&CheckCode::UP016) {
+                    pyupgrade::plugins::remove_six_compat(self, expr);
+                }
+
                 if self.settings.enabled.contains(&CheckCode::YTT202) {
                     flake8_2020::plugins::name_or_attribute(self, expr);
                 }
@@ -1575,6 +1579,9 @@ where
                 if self.settings.enabled.contains(&CheckCode::UP012) {
                     pyupgrade::plugins::unnecessary_encode_utf8(self, expr, func, args, keywords);
                 }
+                if self.settings.enabled.contains(&CheckCode::UP016) {
+                    pyupgrade::plugins::remove_six_compat(self, expr);
+                }
 
                 // flake8-super
                 if self.settings.enabled.contains(&CheckCode::UP008) {

src/checks.rs

@@ -223,6 +223,7 @@ pub enum CheckCode {
     UP013,
     UP014,
     UP015,
+    UP016,
     // pydocstyle
     D100,
     D101,
@@ -792,6 +793,7 @@ pub enum CheckKind {
     ConvertTypedDictFunctionalToClass(String),
     ConvertNamedTupleFunctionalToClass(String),
     RedundantOpenModes,
+    RemoveSixCompat,
     // pydocstyle
     BlankLineAfterLastSection(String),
     BlankLineAfterSection(String),
@@ -1146,6 +1148,7 @@ impl CheckCode {
             CheckCode::UP013 => CheckKind::ConvertTypedDictFunctionalToClass("...".to_string()),
             CheckCode::UP014 => CheckKind::ConvertNamedTupleFunctionalToClass("...".to_string()),
             CheckCode::UP015 => CheckKind::RedundantOpenModes,
+            CheckCode::UP016 => CheckKind::RemoveSixCompat,
             // pydocstyle
             CheckCode::D100 => CheckKind::PublicModule,
             CheckCode::D101 => CheckKind::PublicClass,
@@ -1531,6 +1534,7 @@ impl CheckCode {
             CheckCode::UP013 => CheckCategory::Pyupgrade,
             CheckCode::UP014 => CheckCategory::Pyupgrade,
             CheckCode::UP015 => CheckCategory::Pyupgrade,
+            CheckCode::UP016 => CheckCategory::Pyupgrade,
             CheckCode::W292 => CheckCategory::Pycodestyle,
             CheckCode::W605 => CheckCategory::Pycodestyle,
             CheckCode::YTT101 => CheckCategory::Flake82020,
@@ -1739,6 +1743,7 @@ impl CheckKind {
             CheckKind::ConvertTypedDictFunctionalToClass(_) => &CheckCode::UP013,
             CheckKind::ConvertNamedTupleFunctionalToClass(_) => &CheckCode::UP014,
             CheckKind::RedundantOpenModes => &CheckCode::UP015,
+            CheckKind::RemoveSixCompat => &CheckCode::UP016,
             // pydocstyle
             CheckKind::BlankLineAfterLastSection(_) => &CheckCode::D413,
             CheckKind::BlankLineAfterSection(_) => &CheckCode::D410,
@@ -2443,6 +2448,7 @@ impl CheckKind {
             }
             CheckKind::UnnecessaryEncodeUTF8 => "Unnecessary call to `encode` as UTF-8".to_string(),
             CheckKind::RedundantOpenModes => "Unnecessary open mode parameters".to_string(),
+            CheckKind::RemoveSixCompat => "Unnecessary `six` compatibility usage".to_string(),
             CheckKind::ConvertTypedDictFunctionalToClass(name) => {
                 format!("Convert `{name}` from `TypedDict` functional to class syntax")
             }
@@ -2775,6 +2781,7 @@ impl CheckKind {
                 | CheckKind::CommentedOutCode
                 | CheckKind::ConvertNamedTupleFunctionalToClass(..)
                 | CheckKind::ConvertTypedDictFunctionalToClass(..)
+                | CheckKind::RemoveSixCompat
                 | CheckKind::DashedUnderlineAfterSection(..)
                 | CheckKind::DeprecatedUnittestAlias(..)
                 | CheckKind::DoNotAssertFalse
@@ -2891,6 +2898,7 @@ pub static CODE_REDIRECTS: Lazy<FxHashMap<&'static str, CheckCode>> = Lazy::new(
         ("U013", CheckCode::UP013),
         ("U014", CheckCode::UP014),
         ("U015", CheckCode::UP015),
+        ("U016", CheckCode::UP016),
         // TODO(charlie): Remove by 2023-02-01.
         ("I252", CheckCode::TID252),
         ("M001", CheckCode::RUF100),

src/checks_gen.rs

@@ -434,6 +434,7 @@ pub enum CheckCodePrefix {
     U013,
     U014,
     U015,
+    U016,
     UP,
     UP0,
     UP00,
@@ -452,6 +453,7 @@ pub enum CheckCodePrefix {
     UP013,
     UP014,
     UP015,
+    UP016,
     W,
     W2,
     W29,
@@ -1610,6 +1612,7 @@ impl CheckCodePrefix {
                     CheckCode::UP013,
                     CheckCode::UP014,
                     CheckCode::UP015,
+                    CheckCode::UP016,
                 ]
             }
             CheckCodePrefix::U0 => {
@@ -1634,6 +1637,7 @@ impl CheckCodePrefix {
                     CheckCode::UP013,
                     CheckCode::UP014,
                     CheckCode::UP015,
+                    CheckCode::UP016,
                 ]
             }
             CheckCodePrefix::U00 => {
@@ -1740,6 +1744,7 @@ impl CheckCodePrefix {
                     CheckCode::UP013,
                     CheckCode::UP014,
                     CheckCode::UP015,
+                    CheckCode::UP016,
                 ]
             }
             CheckCodePrefix::U010 => {
@@ -1796,6 +1801,15 @@ impl CheckCodePrefix {
                 );
                 vec![CheckCode::UP015]
             }
+            CheckCodePrefix::U016 => {
+                eprintln!(
+                    "{}{} {}",
+                    "warning".yellow().bold(),
+                    ":".bold(),
+                    "`U016` has been remapped to `UP016`".bold()
+                );
+                vec![CheckCode::UP016]
+            }
             CheckCodePrefix::UP => vec![
                 CheckCode::UP001,
                 CheckCode::UP003,
@@ -1811,6 +1825,7 @@ impl CheckCodePrefix {
                 CheckCode::UP013,
                 CheckCode::UP014,
                 CheckCode::UP015,
+                CheckCode::UP016,
             ],
             CheckCodePrefix::UP0 => vec![
                 CheckCode::UP001,
@@ -1827,6 +1842,7 @@ impl CheckCodePrefix {
                 CheckCode::UP013,
                 CheckCode::UP014,
                 CheckCode::UP015,
+                CheckCode::UP016,
             ],
             CheckCodePrefix::UP00 => vec![
                 CheckCode::UP001,
@@ -1853,6 +1869,7 @@ impl CheckCodePrefix {
                 CheckCode::UP013,
                 CheckCode::UP014,
                 CheckCode::UP015,
+                CheckCode::UP016,
             ],
             CheckCodePrefix::UP010 => vec![CheckCode::UP010],
             CheckCodePrefix::UP011 => vec![CheckCode::UP011],
@@ -1860,6 +1877,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP013 => vec![CheckCode::UP013],
             CheckCodePrefix::UP014 => vec![CheckCode::UP014],
             CheckCodePrefix::UP015 => vec![CheckCode::UP015],
+            CheckCodePrefix::UP016 => vec![CheckCode::UP016],
             CheckCodePrefix::W => vec![CheckCode::W292, CheckCode::W605],
             CheckCodePrefix::W2 => vec![CheckCode::W292],
             CheckCodePrefix::W29 => vec![CheckCode::W292],
@@ -2337,6 +2355,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::U013 => SuffixLength::Three,
             CheckCodePrefix::U014 => SuffixLength::Three,
             CheckCodePrefix::U015 => SuffixLength::Three,
+            CheckCodePrefix::U016 => SuffixLength::Three,
             CheckCodePrefix::UP => SuffixLength::Zero,
             CheckCodePrefix::UP0 => SuffixLength::One,
             CheckCodePrefix::UP00 => SuffixLength::Two,
@@ -2355,6 +2374,7 @@ impl CheckCodePrefix {
             CheckCodePrefix::UP013 => SuffixLength::Three,
             CheckCodePrefix::UP014 => SuffixLength::Three,
             CheckCodePrefix::UP015 => SuffixLength::Three,
+            CheckCodePrefix::UP016 => SuffixLength::Three,
             CheckCodePrefix::W => SuffixLength::Zero,
             CheckCodePrefix::W2 => SuffixLength::One,
             CheckCodePrefix::W29 => SuffixLength::Two,

src/pyupgrade/mod.rs

@@ -36,6 +36,7 @@ mod tests {
     #[test_case(CheckCode::UP013, Path::new("UP013.py"); "UP013")]
     #[test_case(CheckCode::UP014, Path::new("UP014.py"); "UP014")]
     #[test_case(CheckCode::UP015, Path::new("UP015.py"); "UP015")]
+    #[test_case(CheckCode::UP016, Path::new("UP016.py"); "UP016")]
     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

@@ -2,6 +2,7 @@ pub use convert_named_tuple_functional_to_class::convert_named_tuple_functional_
 pub use convert_typed_dict_functional_to_class::convert_typed_dict_functional_to_class;
 pub use deprecated_unittest_alias::deprecated_unittest_alias;
 pub use redundant_open_modes::redundant_open_modes;
+pub use remove_six_compat::remove_six_compat;
 pub use super_call_with_parameters::super_call_with_parameters;
 pub use type_of_primitive::type_of_primitive;
 pub use unnecessary_encode_utf8::unnecessary_encode_utf8;
@@ -16,6 +17,7 @@ mod convert_named_tuple_functional_to_class;
 mod convert_typed_dict_functional_to_class;
 mod deprecated_unittest_alias;
 mod redundant_open_modes;
+mod remove_six_compat;
 mod super_call_with_parameters;
 mod type_of_primitive;
 mod unnecessary_encode_utf8;

src/pyupgrade/plugins/remove_six_compat.rs

@@ -0,0 +1,427 @@
+use anyhow::{bail, Result};
+use log::error;
+use rustpython_ast::{Constant, Expr, ExprContext, ExprKind, Keyword, StmtKind};
+
+use crate::ast::helpers::{collect_call_paths, create_expr, create_stmt, dealias_call_path};
+use crate::ast::types::Range;
+use crate::autofix::Fix;
+use crate::check_ast::Checker;
+use crate::checks::{Check, CheckCode, CheckKind};
+use crate::code_gen::SourceGenerator;
+use crate::SourceCodeLocator;
+
+/// Return `true` if the `Expr` is a reference to `${module}.${any}`.
+fn is_module_member(call_path: &[&str], module: &str) -> bool {
+    call_path
+        .first()
+        .map_or(false, |module_name| *module_name == module)
+}
+
+fn map_name(name: &str, expr: &Expr, patch: bool) -> Option<Check> {
+    let replacement = match name {
+        "text_type" => Some("str"),
+        "binary_type" => Some("bytes"),
+        "class_types" => Some("(type,)"),
+        "string_types" => Some("(str,)"),
+        "integer_types" => Some("(int,)"),
+        "unichr" => Some("chr"),
+        "iterbytes" => Some("iter"),
+        "print_" => Some("print"),
+        "exec_" => Some("exec"),
+        "advance_iterator" => Some("next"),
+        "next" => Some("next"),
+        "range" => Some("range"),  // TODO: six.moves
+        "xrange" => Some("range"), // TODO: six.moves
+        "callable" => Some("callable"),
+        _ => None,
+    };
+    if let Some(replacement) = replacement {
+        let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
+        if patch {
+            check.amend(Fix::replacement(
+                replacement.to_string(),
+                expr.location,
+                expr.end_location.unwrap(),
+            ));
+        }
+        Some(check)
+    } else {
+        None
+    }
+}
+
+fn replace_by_str_literal(
+    arg: &Expr,
+    binary: bool,
+    expr: &Expr,
+    patch: bool,
+    locator: &SourceCodeLocator,
+) -> Option<Check> {
+    match &arg.node {
+        ExprKind::Constant { .. } => {
+            let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
+            if patch {
+                let content = format!(
+                    "{}{}",
+                    if binary { "b" } else { "" },
+                    locator.slice_source_code_range(&Range {
+                        location: arg.location,
+                        end_location: arg.end_location.unwrap(),
+                    })
+                );
+                check.amend(Fix::replacement(
+                    content,
+                    expr.location,
+                    expr.end_location.unwrap(),
+                ));
+            };
+            Some(check)
+        }
+        _ => None,
+    }
+}
+
+// `func(arg)` => `arg.attr`
+fn replace_call_on_arg_by_arg_attribute(
+    attr: &str,
+    arg: &Expr,
+    expr: &Expr,
+    patch: bool,
+) -> Result<Check> {
+    let attribute = ExprKind::Attribute {
+        value: Box::new(arg.clone()),
+        attr: attr.to_string(),
+        ctx: ExprContext::Load,
+    };
+    replace_by_expr_kind(attribute, expr, patch)
+}
+
+// `func(arg, **args)` => `arg.method(**args)`
+fn replace_call_on_arg_by_arg_method_call(
+    method_name: &str,
+    args: &[Expr],
+    expr: &Expr,
+    patch: bool,
+) -> Result<Option<Check>> {
+    if args.is_empty() {
+        bail!("Expected at least one argument");
+    }
+    if let ([arg], other_args) = args.split_at(1) {
+        let call = ExprKind::Call {
+            func: Box::new(create_expr(ExprKind::Attribute {
+                value: Box::new(arg.clone()),
+                attr: method_name.to_string(),
+                ctx: ExprContext::Load,
+            })),
+            args: other_args
+                .iter()
+                .map(|arg| create_expr(arg.node.clone()))
+                .collect(),
+            keywords: vec![],
+        };
+        let expr = replace_by_expr_kind(call, expr, patch)?;
+        Ok(Some(expr))
+    } else {
+        Ok(None)
+    }
+}
+
+// `expr` => `Expr(expr_kind)`
+fn replace_by_expr_kind(node: ExprKind, expr: &Expr, patch: bool) -> Result<Check> {
+    let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
+    if patch {
+        let mut generator = SourceGenerator::new();
+        generator.unparse_expr(&create_expr(node), 0);
+        let content = generator.generate()?;
+        check.amend(Fix::replacement(
+            content,
+            expr.location,
+            expr.end_location.unwrap(),
+        ));
+    }
+    Ok(check)
+}
+
+fn replace_by_stmt_kind(node: StmtKind, expr: &Expr, patch: bool) -> Result<Check> {
+    let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
+    if patch {
+        let mut generator = SourceGenerator::new();
+        generator.unparse_stmt(&create_stmt(node));
+        let content = generator.generate()?;
+        check.amend(Fix::replacement(
+            content,
+            expr.location,
+            expr.end_location.unwrap(),
+        ));
+    }
+    Ok(check)
+}
+
+// => `raise exc from cause`
+fn replace_by_raise_from(
+    exc: Option<ExprKind>,
+    cause: Option<ExprKind>,
+    expr: &Expr,
+    patch: bool,
+) -> Result<Check> {
+    let stmt_kind = StmtKind::Raise {
+        exc: exc.map(|exc| Box::new(create_expr(exc))),
+        cause: cause.map(|cause| Box::new(create_expr(cause))),
+    };
+    replace_by_stmt_kind(stmt_kind, expr, patch)
+}
+
+fn replace_by_index_on_arg(
+    arg: &Expr,
+    index: &ExprKind,
+    expr: &Expr,
+    patch: bool,
+) -> Result<Check> {
+    let index = ExprKind::Subscript {
+        value: Box::new(create_expr(arg.node.clone())),
+        slice: Box::new(create_expr(index.clone())),
+        ctx: ExprContext::Load,
+    };
+    replace_by_expr_kind(index, expr, patch)
+}
+
+fn handle_reraise(args: &[Expr], expr: &Expr, patch: bool) -> Result<Option<Check>> {
+    if let [_, exc, tb] = args {
+        let check = replace_by_raise_from(
+            Some(ExprKind::Call {
+                func: Box::new(create_expr(ExprKind::Attribute {
+                    value: Box::new(create_expr(exc.node.clone())),
+                    attr: "with_traceback".to_string(),
+                    ctx: ExprContext::Load,
+                })),
+                args: vec![create_expr(tb.node.clone())],
+                keywords: vec![],
+            }),
+            None,
+            expr,
+            patch,
+        )?;
+        Ok(Some(check))
+    } else if let [arg] = args {
+        if let ExprKind::Starred { value, .. } = &arg.node {
+            if let ExprKind::Call { func, .. } = &value.node {
+                if let ExprKind::Attribute { value, attr, .. } = &func.node {
+                    if let ExprKind::Name { id, .. } = &value.node {
+                        if id == "sys" && attr == "exc_info" {
+                            let check = replace_by_raise_from(None, None, expr, patch)?;
+                            return Ok(Some(check));
+                        };
+                    };
+                };
+            };
+        };
+        Ok(None)
+    } else {
+        Ok(None)
+    }
+}
+
+fn handle_func(
+    func: &Expr,
+    args: &[Expr],
+    keywords: &[Keyword],
+    expr: &Expr,
+    patch: bool,
+    locator: &SourceCodeLocator,
+) -> Result<Option<Check>> {
+    let func_name = match &func.node {
+        ExprKind::Attribute { attr, .. } => attr,
+        ExprKind::Name { id, .. } => id,
+        _ => bail!("Unexpected func: {:?}", func),
+    };
+    let check = match (func_name.as_str(), args, keywords) {
+        ("b", [arg], []) => replace_by_str_literal(arg, true, expr, patch, locator),
+        ("ensure_binary", [arg], []) => replace_by_str_literal(arg, true, expr, patch, locator),
+        ("u", [arg], []) => replace_by_str_literal(arg, false, expr, patch, locator),
+        ("ensure_str", [arg], []) => replace_by_str_literal(arg, false, expr, patch, locator),
+        ("ensure_text", [arg], []) => replace_by_str_literal(arg, false, expr, patch, locator),
+        ("iteritems", args, []) => {
+            replace_call_on_arg_by_arg_method_call("items", args, expr, patch)?
+        }
+        ("viewitems", args, []) => {
+            replace_call_on_arg_by_arg_method_call("items", args, expr, patch)?
+        }
+        ("iterkeys", args, []) => {
+            replace_call_on_arg_by_arg_method_call("keys", args, expr, patch)?
+        }
+        ("viewkeys", args, []) => {
+            replace_call_on_arg_by_arg_method_call("keys", args, expr, patch)?
+        }
+        ("itervalues", args, []) => {
+            replace_call_on_arg_by_arg_method_call("values", args, expr, patch)?
+        }
+        ("viewvalues", args, []) => {
+            replace_call_on_arg_by_arg_method_call("values", args, expr, patch)?
+        }
+        ("get_method_function", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__func__", arg, expr, patch,
+        )?),
+        ("get_method_self", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__self__", arg, expr, patch,
+        )?),
+        ("get_function_closure", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__closure__",
+            arg,
+            expr,
+            patch,
+        )?),
+        ("get_function_code", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__code__", arg, expr, patch,
+        )?),
+        ("get_function_defaults", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__defaults__",
+            arg,
+            expr,
+            patch,
+        )?),
+        ("get_function_globals", [arg], []) => Some(replace_call_on_arg_by_arg_attribute(
+            "__globals__",
+            arg,
+            expr,
+            patch,
+        )?),
+        ("create_unbound_method", [arg, _], _) => {
+            Some(replace_by_expr_kind(arg.node.clone(), expr, patch)?)
+        }
+        ("get_unbound_function", [arg], []) => {
+            Some(replace_by_expr_kind(arg.node.clone(), expr, patch)?)
+        }
+        ("assertCountEqual", args, []) => {
+            replace_call_on_arg_by_arg_method_call("assertCountEqual", args, expr, patch)?
+        }
+        ("assertRaisesRegex", args, []) => {
+            replace_call_on_arg_by_arg_method_call("assertRaisesRegex", args, expr, patch)?
+        }
+        ("assertRegex", args, []) => {
+            replace_call_on_arg_by_arg_method_call("assertRegex", args, expr, patch)?
+        }
+        ("raise_from", [exc, cause], []) => Some(replace_by_raise_from(
+            Some(exc.node.clone()),
+            Some(cause.node.clone()),
+            expr,
+            patch,
+        )?),
+        ("reraise", args, []) => handle_reraise(args, expr, patch)?,
+        ("byte2int", [arg], []) => Some(replace_by_index_on_arg(
+            arg,
+            &ExprKind::Constant {
+                value: Constant::Int(0.into()),
+                kind: None,
+            },
+            expr,
+            patch,
+        )?),
+        ("indexbytes", [arg, index], []) => {
+            Some(replace_by_index_on_arg(arg, &index.node, expr, patch)?)
+        }
+        ("int2byte", [arg], []) => Some(replace_by_expr_kind(
+            ExprKind::Call {
+                func: Box::new(create_expr(ExprKind::Name {
+                    id: "bytes".to_string(),
+                    ctx: ExprContext::Load,
+                })),
+                args: vec![create_expr(ExprKind::Tuple {
+                    elts: vec![create_expr(arg.node.clone())],
+                    ctx: ExprContext::Load,
+                })],
+                keywords: vec![],
+            },
+            expr,
+            patch,
+        )?),
+        _ => None,
+    };
+    Ok(check)
+}
+
+fn handle_next_on_six_dict(expr: &Expr, patch: bool, checker: &Checker) -> Result<Option<Check>> {
+    let ExprKind::Call { func, args, .. } = &expr.node else {
+        return Ok(None);
+    };
+    let ExprKind::Name { id, .. } = &func.node else {
+        return Ok(None);
+    };
+    if id != "next" {
+        return Ok(None);
+    }
+    let [arg] = &args[..] else { return Ok(None); };
+    let call_path = dealias_call_path(collect_call_paths(arg), &checker.import_aliases);
+    if !is_module_member(&call_path, "six") {
+        return Ok(None);
+    }
+    let ExprKind::Call { func, args, .. } = &arg.node else {return Ok(None);};
+    let ExprKind::Attribute { attr, .. } = &func.node else {return Ok(None);};
+    let [dict_arg] = &args[..] else {return Ok(None);};
+    let method_name = match attr.as_str() {
+        "iteritems" => "items",
+        "iterkeys" => "keys",
+        "itervalues" => "values",
+        _ => return Ok(None),
+    };
+    match replace_by_expr_kind(
+        ExprKind::Call {
+            func: Box::new(create_expr(ExprKind::Name {
+                id: "iter".to_string(),
+                ctx: ExprContext::Load,
+            })),
+            args: vec![create_expr(ExprKind::Call {
+                func: Box::new(create_expr(ExprKind::Attribute {
+                    value: Box::new(dict_arg.clone()),
+                    attr: method_name.to_string(),
+                    ctx: ExprContext::Load,
+                })),
+                args: vec![],
+                keywords: vec![],
+            })],
+            keywords: vec![],
+        },
+        arg,
+        patch,
+    ) {
+        Ok(check) => Ok(Some(check)),
+        Err(err) => Err(err),
+    }
+}
+
+/// UP016
+pub fn remove_six_compat(checker: &mut Checker, expr: &Expr) {
+    match handle_next_on_six_dict(expr, checker.patch(&CheckCode::UP016), checker) {
+        Ok(Some(check)) => {
+            checker.add_check(check);
+            return;
+        }
+        Ok(None) => (),
+        Err(err) => {
+            error!("Error while removing `six` reference: {}", err);
+            return;
+        }
+    };
+    let call_path = dealias_call_path(collect_call_paths(expr), &checker.import_aliases);
+    if is_module_member(&call_path, "six") {
+        let patch = checker.patch(&CheckCode::UP016);
+        let check = match &expr.node {
+            ExprKind::Call {
+                func,
+                args,
+                keywords,
+            } => match handle_func(func, args, keywords, expr, patch, checker.locator) {
+                Ok(check) => check,
+                Err(err) => {
+                    error!("Failed to remove `six` reference: {err}");
+                    return;
+                }
+            },
+            ExprKind::Attribute { attr, .. } => map_name(attr.as_str(), expr, patch),
+            ExprKind::Name { id, .. } => map_name(id.as_str(), expr, patch),
+            _ => return,
+        };
+        if let Some(check) = check {
+            checker.add_check(check);
+        }
+    }
+}

src/pyupgrade/snapshots/ruff__pyupgrade__tests__UP016_UP016.py.snap

@@ -0,0 +1,770 @@
+---
+source: src/pyupgrade/mod.rs
+expression: checks
+---
+- kind: RemoveSixCompat
+  location:
+    row: 7
+    column: 0
+  end_location:
+    row: 7
+    column: 13
+  fix:
+    content: str
+    location:
+      row: 7
+      column: 0
+    end_location:
+      row: 7
+      column: 13
+- kind: RemoveSixCompat
+  location:
+    row: 8
+    column: 0
+  end_location:
+    row: 8
+    column: 15
+  fix:
+    content: bytes
+    location:
+      row: 8
+      column: 0
+    end_location:
+      row: 8
+      column: 15
+- kind: RemoveSixCompat
+  location:
+    row: 9
+    column: 0
+  end_location:
+    row: 9
+    column: 15
+  fix:
+    content: "(type,)"
+    location:
+      row: 9
+      column: 0
+    end_location:
+      row: 9
+      column: 15
+- kind: RemoveSixCompat
+  location:
+    row: 10
+    column: 0
+  end_location:
+    row: 10
+    column: 16
+  fix:
+    content: "(str,)"
+    location:
+      row: 10
+      column: 0
+    end_location:
+      row: 10
+      column: 16
+- kind: RemoveSixCompat
+  location:
+    row: 11
+    column: 0
+  end_location:
+    row: 11
+    column: 17
+  fix:
+    content: "(int,)"
+    location:
+      row: 11
+      column: 0
+    end_location:
+      row: 11
+      column: 17
+- kind: RemoveSixCompat
+  location:
+    row: 12
+    column: 0
+  end_location:
+    row: 12
+    column: 10
+  fix:
+    content: chr
+    location:
+      row: 12
+      column: 0
+    end_location:
+      row: 12
+      column: 10
+- kind: RemoveSixCompat
+  location:
+    row: 13
+    column: 0
+  end_location:
+    row: 13
+    column: 13
+  fix:
+    content: iter
+    location:
+      row: 13
+      column: 0
+    end_location:
+      row: 13
+      column: 13
+- kind: RemoveSixCompat
+  location:
+    row: 14
+    column: 0
+  end_location:
+    row: 14
+    column: 10
+  fix:
+    content: print
+    location:
+      row: 14
+      column: 0
+    end_location:
+      row: 14
+      column: 10
+- kind: RemoveSixCompat
+  location:
+    row: 15
+    column: 0
+  end_location:
+    row: 15
+    column: 9
+  fix:
+    content: exec
+    location:
+      row: 15
+      column: 0
+    end_location:
+      row: 15
+      column: 9
+- kind: RemoveSixCompat
+  location:
+    row: 16
+    column: 0
+  end_location:
+    row: 16
+    column: 20
+  fix:
+    content: next
+    location:
+      row: 16
+      column: 0
+    end_location:
+      row: 16
+      column: 20
+- kind: RemoveSixCompat
+  location:
+    row: 17
+    column: 0
+  end_location:
+    row: 17
+    column: 8
+  fix:
+    content: next
+    location:
+      row: 17
+      column: 0
+    end_location:
+      row: 17
+      column: 8
+- kind: RemoveSixCompat
+  location:
+    row: 18
+    column: 0
+  end_location:
+    row: 18
+    column: 12
+  fix:
+    content: callable
+    location:
+      row: 18
+      column: 0
+    end_location:
+      row: 18
+      column: 12
+- kind: RemoveSixCompat
+  location:
+    row: 19
+    column: 0
+  end_location:
+    row: 19
+    column: 15
+  fix:
+    content: range
+    location:
+      row: 19
+      column: 0
+    end_location:
+      row: 19
+      column: 15
+- kind: RemoveSixCompat
+  location:
+    row: 20
+    column: 0
+  end_location:
+    row: 20
+    column: 16
+  fix:
+    content: range
+    location:
+      row: 20
+      column: 0
+    end_location:
+      row: 20
+      column: 16
+- kind: RemoveSixCompat
+  location:
+    row: 21
+    column: 16
+  end_location:
+    row: 21
+    column: 31
+  fix:
+    content: "(type,)"
+    location:
+      row: 21
+      column: 16
+    end_location:
+      row: 21
+      column: 31
+- kind: RemoveSixCompat
+  location:
+    row: 22
+    column: 16
+  end_location:
+    row: 22
+    column: 33
+  fix:
+    content: "(int,)"
+    location:
+      row: 22
+      column: 16
+    end_location:
+      row: 22
+      column: 33
+- kind: RemoveSixCompat
+  location:
+    row: 23
+    column: 16
+  end_location:
+    row: 23
+    column: 32
+  fix:
+    content: "(str,)"
+    location:
+      row: 23
+      column: 16
+    end_location:
+      row: 23
+      column: 32
+- kind: RemoveSixCompat
+  location:
+    row: 26
+    column: 0
+  end_location:
+    row: 26
+    column: 18
+  fix:
+    content: dct.items()
+    location:
+      row: 26
+      column: 0
+    end_location:
+      row: 26
+      column: 18
+- kind: RemoveSixCompat
+  location:
+    row: 27
+    column: 0
+  end_location:
+    row: 27
+    column: 17
+  fix:
+    content: dct.keys()
+    location:
+      row: 27
+      column: 0
+    end_location:
+      row: 27
+      column: 17
+- kind: RemoveSixCompat
+  location:
+    row: 28
+    column: 0
+  end_location:
+    row: 28
+    column: 19
+  fix:
+    content: dct.values()
+    location:
+      row: 28
+      column: 0
+    end_location:
+      row: 28
+      column: 19
+- kind: RemoveSixCompat
+  location:
+    row: 29
+    column: 0
+  end_location:
+    row: 29
+    column: 18
+  fix:
+    content: dct.items()
+    location:
+      row: 29
+      column: 0
+    end_location:
+      row: 29
+      column: 18
+- kind: RemoveSixCompat
+  location:
+    row: 30
+    column: 0
+  end_location:
+    row: 30
+    column: 17
+  fix:
+    content: dct.keys()
+    location:
+      row: 30
+      column: 0
+    end_location:
+      row: 30
+      column: 17
+- kind: RemoveSixCompat
+  location:
+    row: 31
+    column: 0
+  end_location:
+    row: 31
+    column: 19
+  fix:
+    content: dct.values()
+    location:
+      row: 31
+      column: 0
+    end_location:
+      row: 31
+      column: 19
+- kind: RemoveSixCompat
+  location:
+    row: 32
+    column: 0
+  end_location:
+    row: 32
+    column: 34
+  fix:
+    content: "self.assertCountEqual(a1, a2)"
+    location:
+      row: 32
+      column: 0
+    end_location:
+      row: 32
+      column: 34
+- kind: RemoveSixCompat
+  location:
+    row: 33
+    column: 0
+  end_location:
+    row: 33
+    column: 37
+  fix:
+    content: "self.assertRaisesRegex(e, r, fn)"
+    location:
+      row: 33
+      column: 0
+    end_location:
+      row: 33
+      column: 37
+- kind: RemoveSixCompat
+  location:
+    row: 34
+    column: 0
+  end_location:
+    row: 34
+    column: 27
+  fix:
+    content: "self.assertRegex(s, r)"
+    location:
+      row: 34
+      column: 0
+    end_location:
+      row: 34
+      column: 27
+- kind: RemoveSixCompat
+  location:
+    row: 37
+    column: 0
+  end_location:
+    row: 37
+    column: 29
+  fix:
+    content: meth.__func__
+    location:
+      row: 37
+      column: 0
+    end_location:
+      row: 37
+      column: 29
+- kind: RemoveSixCompat
+  location:
+    row: 38
+    column: 0
+  end_location:
+    row: 38
+    column: 25
+  fix:
+    content: meth.__self__
+    location:
+      row: 38
+      column: 0
+    end_location:
+      row: 38
+      column: 25
+- kind: RemoveSixCompat
+  location:
+    row: 39
+    column: 0
+  end_location:
+    row: 39
+    column: 28
+  fix:
+    content: fn.__closure__
+    location:
+      row: 39
+      column: 0
+    end_location:
+      row: 39
+      column: 28
+- kind: RemoveSixCompat
+  location:
+    row: 40
+    column: 0
+  end_location:
+    row: 40
+    column: 25
+  fix:
+    content: fn.__code__
+    location:
+      row: 40
+      column: 0
+    end_location:
+      row: 40
+      column: 25
+- kind: RemoveSixCompat
+  location:
+    row: 41
+    column: 0
+  end_location:
+    row: 41
+    column: 29
+  fix:
+    content: fn.__defaults__
+    location:
+      row: 41
+      column: 0
+    end_location:
+      row: 41
+      column: 29
+- kind: RemoveSixCompat
+  location:
+    row: 42
+    column: 0
+  end_location:
+    row: 42
+    column: 28
+  fix:
+    content: fn.__globals__
+    location:
+      row: 42
+      column: 0
+    end_location:
+      row: 42
+      column: 28
+- kind: RemoveSixCompat
+  location:
+    row: 45
+    column: 0
+  end_location:
+    row: 45
+    column: 12
+  fix:
+    content: "b\"...\""
+    location:
+      row: 45
+      column: 0
+    end_location:
+      row: 45
+      column: 12
+- kind: RemoveSixCompat
+  location:
+    row: 46
+    column: 0
+  end_location:
+    row: 46
+    column: 12
+  fix:
+    content: "\"...\""
+    location:
+      row: 46
+      column: 0
+    end_location:
+      row: 46
+      column: 12
+- kind: RemoveSixCompat
+  location:
+    row: 47
+    column: 0
+  end_location:
+    row: 47
+    column: 24
+  fix:
+    content: "b\"...\""
+    location:
+      row: 47
+      column: 0
+    end_location:
+      row: 47
+      column: 24
+- kind: RemoveSixCompat
+  location:
+    row: 48
+    column: 0
+  end_location:
+    row: 48
+    column: 21
+  fix:
+    content: "\"...\""
+    location:
+      row: 48
+      column: 0
+    end_location:
+      row: 48
+      column: 21
+- kind: RemoveSixCompat
+  location:
+    row: 49
+    column: 0
+  end_location:
+    row: 49
+    column: 22
+  fix:
+    content: "\"...\""
+    location:
+      row: 49
+      column: 0
+    end_location:
+      row: 49
+      column: 22
+- kind: RemoveSixCompat
+  location:
+    row: 53
+    column: 0
+  end_location:
+    row: 53
+    column: 30
+  fix:
+    content: meth
+    location:
+      row: 53
+      column: 0
+    end_location:
+      row: 53
+      column: 30
+- kind: RemoveSixCompat
+  location:
+    row: 54
+    column: 0
+  end_location:
+    row: 54
+    column: 34
+  fix:
+    content: fn
+    location:
+      row: 54
+      column: 0
+    end_location:
+      row: 54
+      column: 34
+- kind: RemoveSixCompat
+  location:
+    row: 57
+    column: 0
+  end_location:
+    row: 57
+    column: 29
+  fix:
+    content: raise exc from exc_from
+    location:
+      row: 57
+      column: 0
+    end_location:
+      row: 57
+      column: 29
+- kind: RemoveSixCompat
+  location:
+    row: 58
+    column: 0
+  end_location:
+    row: 58
+    column: 24
+  fix:
+    content: raise exc.with_traceback(tb)
+    location:
+      row: 58
+      column: 0
+    end_location:
+      row: 58
+      column: 24
+- kind: RemoveSixCompat
+  location:
+    row: 59
+    column: 0
+  end_location:
+    row: 59
+    column: 28
+  fix:
+    content: raise
+    location:
+      row: 59
+      column: 0
+    end_location:
+      row: 59
+      column: 28
+- kind: RemoveSixCompat
+  location:
+    row: 62
+    column: 0
+  end_location:
+    row: 62
+    column: 16
+  fix:
+    content: "bs[0]"
+    location:
+      row: 62
+      column: 0
+    end_location:
+      row: 62
+      column: 16
+- kind: RemoveSixCompat
+  location:
+    row: 63
+    column: 0
+  end_location:
+    row: 63
+    column: 21
+  fix:
+    content: "bs[i]"
+    location:
+      row: 63
+      column: 0
+    end_location:
+      row: 63
+      column: 21
+- kind: RemoveSixCompat
+  location:
+    row: 64
+    column: 0
+  end_location:
+    row: 64
+    column: 15
+  fix:
+    content: "bytes((i,))"
+    location:
+      row: 64
+      column: 0
+    end_location:
+      row: 64
+      column: 15
+- kind: RemoveSixCompat
+  location:
+    row: 67
+    column: 5
+  end_location:
+    row: 67
+    column: 23
+  fix:
+    content: iter(dct.items())
+    location:
+      row: 67
+      column: 5
+    end_location:
+      row: 67
+      column: 23
+- kind: RemoveSixCompat
+  location:
+    row: 67
+    column: 5
+  end_location:
+    row: 67
+    column: 23
+  fix:
+    content: dct.items()
+    location:
+      row: 67
+      column: 5
+    end_location:
+      row: 67
+      column: 23
+- kind: RemoveSixCompat
+  location:
+    row: 68
+    column: 5
+  end_location:
+    row: 68
+    column: 22
+  fix:
+    content: iter(dct.keys())
+    location:
+      row: 68
+      column: 5
+    end_location:
+      row: 68
+      column: 22
+- kind: RemoveSixCompat
+  location:
+    row: 68
+    column: 5
+  end_location:
+    row: 68
+    column: 22
+  fix:
+    content: dct.keys()
+    location:
+      row: 68
+      column: 5
+    end_location:
+      row: 68
+      column: 22
+- kind: RemoveSixCompat
+  location:
+    row: 69
+    column: 5
+  end_location:
+    row: 69
+    column: 24
+  fix:
+    content: iter(dct.values())
+    location:
+      row: 69
+      column: 5
+    end_location:
+      row: 69
+      column: 24
+- kind: RemoveSixCompat
+  location:
+    row: 69
+    column: 5
+  end_location:
+    row: 69
+    column: 24
+  fix:
+    content: dct.values()
+    location:
+      row: 69
+      column: 5
+    end_location:
+      row: 69
+      column: 24
+