Enable automatic rewrites of typing.Deque and typing.DefaultDict (#4420)

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -2,8 +2,7 @@ use rustpython_parser::ast::{self, Constant, Expr, ExprKind, Operator};
 
 use ruff_python_ast::call_path::{from_unqualified_name, CallPath};
 use ruff_python_stdlib::typing::{
-    IMMUTABLE_GENERIC_TYPES, IMMUTABLE_TYPES, PEP_585_BUILTINS_ELIGIBLE, PEP_593_SUBSCRIPTS,
-    SUBSCRIPTS,
+    IMMUTABLE_GENERIC_TYPES, IMMUTABLE_TYPES, PEP_585_GENERICS, PEP_593_SUBSCRIPTS, SUBSCRIPTS,
 };
 
 use crate::context::Context;
@@ -62,23 +61,90 @@ pub fn match_annotated_subscript<'a>(
     })
 }
 
-/// Returns `true` if `Expr` represents a reference to a typing object with a
-/// PEP 585 built-in.
-pub fn is_pep585_builtin(expr: &Expr, context: &Context) -> bool {
-    context.resolve_call_path(expr).map_or(false, |call_path| {
-        PEP_585_BUILTINS_ELIGIBLE.contains(&call_path.as_slice())
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum ModuleMember {
+    /// A builtin symbol, like `"list"`.
+    BuiltIn(&'static str),
+    /// A module member, like `("collections", "deque")`.
+    Member(&'static str, &'static str),
+}
+
+impl std::fmt::Display for ModuleMember {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            ModuleMember::BuiltIn(name) => std::write!(f, "{name}"),
+            ModuleMember::Member(module, member) => std::write!(f, "{module}.{member}"),
+        }
+    }
+}
+
+/// Returns the PEP 585 standard library generic variant for a `typing` module reference, if such
+/// a variant exists.
+pub fn to_pep585_generic(expr: &Expr, context: &Context) -> Option<ModuleMember> {
+    context.resolve_call_path(expr).and_then(|call_path| {
+        let [module, name] = call_path.as_slice() else {
+            return None;
+        };
+        PEP_585_GENERICS
+            .iter()
+            .find_map(|((from_module, from_member), (to_module, to_member))| {
+                if module == from_module && name == from_member {
+                    if to_module.is_empty() {
+                        Some(ModuleMember::BuiltIn(to_member))
+                    } else {
+                        Some(ModuleMember::Member(to_module, to_member))
+                    }
+                } else {
+                    None
+                }
+            })
     })
 }
 
-/// Returns `true` if `Expr` represents a reference to a typing object with a
-/// PEP 603 built-in.
-pub fn is_pep604_builtin(expr: &Expr, context: &Context) -> bool {
-    context.resolve_call_path(expr).map_or(false, |call_path| {
-        context.match_typing_call_path(&call_path, "Optional")
-            || context.match_typing_call_path(&call_path, "Union")
-    })
+#[derive(Debug, Copy, Clone)]
+pub enum Pep604Operator {
+    /// The union operator, e.g., `Union[str, int]`, expressible as `str | int` after PEP 604.
+    Union,
+    /// The union operator, e.g., `Optional[str]`, expressible as `str | None` after PEP 604.
+    Optional,
 }
 
+/// Return the PEP 604 operator variant to which the given subscript [`Expr`] corresponds, if any.
+pub fn to_pep604_operator(value: &Expr, slice: &Expr, context: &Context) -> Option<Pep604Operator> {
+    /// Returns `true` if any argument in the slice is a string.
+    fn any_arg_is_str(slice: &Expr) -> bool {
+        match &slice.node {
+            ExprKind::Constant(ast::ExprConstant {
+                value: Constant::Str(_),
+                ..
+            }) => true,
+            ExprKind::Tuple(ast::ExprTuple { elts, .. }) => elts.iter().any(any_arg_is_str),
+            _ => false,
+        }
+    }
+
+    // If any of the _arguments_ are forward references, we can't use PEP 604.
+    // Ex) `Union["str", "int"]` can't be converted to `"str" | "int"`.
+    if any_arg_is_str(slice) {
+        return None;
+    }
+
+    context
+        .resolve_call_path(value)
+        .as_ref()
+        .and_then(|call_path| {
+            if context.match_typing_call_path(call_path, "Optional") {
+                Some(Pep604Operator::Optional)
+            } else if context.match_typing_call_path(call_path, "Union") {
+                Some(Pep604Operator::Union)
+            } else {
+                None
+            }
+        })
+}
+
+/// Return `true` if `Expr` represents a reference to a type annotation that resolves to an
+/// immutable type.
 pub fn is_immutable_annotation(context: &Context, expr: &Expr) -> bool {
     match &expr.node {
         ExprKind::Name(_) | ExprKind::Attribute(_) => {
@@ -144,6 +210,7 @@ const IMMUTABLE_FUNCS: &[&[&str]] = &[
     &["re", "compile"],
 ];
 
+/// Return `true` if `func` is a function that returns an immutable object.
 pub fn is_immutable_func(
     context: &Context,
     func: &Expr,