[pylint] Avoid suggesting set rewrites for non-hashable types

crates/ruff_linter/resources/test/fixtures/pylint/literal_membership.py

@@ -4,7 +4,12 @@
 1 in (
     1, 2, 3
 )
-
-# OK
 fruits = ["cherry", "grapes"]
 "cherry" in fruits
+_ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in ("a", "b")}
+
+# OK
+fruits in [[1, 2, 3], [4, 5, 6]]
+fruits in [1, 2, 3]
+1 in [[1, 2, 3], [4, 5, 6]]
+_ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in (["a", "b"], ["c", "d"])}

crates/ruff_linter/src/rules/pylint/rules/literal_membership.rs

@@ -1,6 +1,7 @@
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, CmpOp, Expr};
+use ruff_python_semantic::analyze::typing;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -25,7 +26,8 @@ use crate::checkers::ast::Checker;
 /// ## Fix safety
 /// This rule's fix is marked as unsafe, as the use of a `set` literal will
 /// error at runtime if the sequence contains unhashable elements (like lists
-/// or dictionaries).
+/// or dictionaries). While Ruff will attempt to infer the hashability of the
+/// elements, it may not always be able to do so.
 ///
 /// ## References
 /// - [What’s New In Python 3.2](https://docs.python.org/3/whatsnew/3.2.html#optimizations)
@@ -57,7 +59,40 @@ pub(crate) fn literal_membership(checker: &mut Checker, compare: &ast::ExprCompa
         return;
     };
 
-    if !matches!(right, Expr::List(_) | Expr::Tuple(_)) {
+    let elts = match right {
+        Expr::List(ast::ExprList { elts, .. }) => elts,
+        Expr::Tuple(ast::ExprTuple { elts, .. }) => elts,
+        _ => return,
+    };
+
+    // If `left`, or any of the elements in `right`, are known to _not_ be hashable, return.
+    if std::iter::once(compare.left.as_ref())
+        .chain(elts)
+        .any(|expr| match expr {
+            // Expressions that are known _not_ to be hashable.
+            Expr::List(_)
+            | Expr::Set(_)
+            | Expr::Dict(_)
+            | Expr::ListComp(_)
+            | Expr::SetComp(_)
+            | Expr::DictComp(_)
+            | Expr::GeneratorExp(_)
+            | Expr::Await(_)
+            | Expr::Yield(_)
+            | Expr::YieldFrom(_) => true,
+            // Expressions that can be _inferred_ not to be hashable.
+            Expr::Name(name) => {
+                let Some(id) = checker.semantic().resolve_name(name) else {
+                    return false;
+                };
+                let binding = checker.semantic().binding(id);
+                typing::is_list(binding, checker.semantic())
+                    || typing::is_dict(binding, checker.semantic())
+                    || typing::is_set(binding, checker.semantic())
+            }
+            _ => false,
+        })
+    {
         return;
     }
 

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR6201_literal_membership.py.snap

@@ -48,8 +48,8 @@ literal_membership.py:4:6: PLR6201 [*] Use a `set` literal when testing for memb
 5 | |     1, 2, 3
 6 | | )
   | |_^ PLR6201
-7 |   
-8 |   # OK
+7 |   fruits = ["cherry", "grapes"]
+8 |   "cherry" in fruits
   |
   = help: Convert to `set`
 
@@ -62,8 +62,29 @@ literal_membership.py:4:6: PLR6201 [*] Use a `set` literal when testing for memb
 5 5 |     1, 2, 3
 6   |-)
   6 |+}
-7 7 | 
-8 8 | # OK
-9 9 | fruits = ["cherry", "grapes"]
+7 7 | fruits = ["cherry", "grapes"]
+8 8 | "cherry" in fruits
+9 9 | _ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in ("a", "b")}
+
+literal_membership.py:9:70: PLR6201 [*] Use a `set` literal when testing for membership
+   |
+ 7 | fruits = ["cherry", "grapes"]
+ 8 | "cherry" in fruits
+ 9 | _ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in ("a", "b")}
+   |                                                                      ^^^^^^^^^^ PLR6201
+10 | 
+11 | # OK
+   |
+   = help: Convert to `set`
+
+ℹ Unsafe fix
+6  6  | )
+7  7  | fruits = ["cherry", "grapes"]
+8  8  | "cherry" in fruits
+9     |-_ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in ("a", "b")}
+   9  |+_ = {key: value for key, value in {"a": 1, "b": 2}.items() if key in {"a", "b"}}
+10 10 | 
+11 11 | # OK
+12 12 | fruits in [[1, 2, 3], [4, 5, 6]]
 
 

crates/ruff_python_semantic/src/analyze/typing.rs

@@ -426,17 +426,32 @@ fn check_type<T: TypeChecker>(binding: &Binding, semantic: &SemanticModel) -> bo
             // ```
             //
             // The type checker might know how to infer the type based on `init_expr`.
-            Some(Stmt::Assign(ast::StmtAssign { value, .. })) => {
-                T::match_initializer(value.as_ref(), semantic)
+            Some(Stmt::Assign(ast::StmtAssign { targets, value, .. })) => {
+                // TODO(charlie): Replace this with `find_binding_value`, which matches the values.
+                if targets
+                    .iter()
+                    .any(|target| target.range().contains_range(binding.range()))
+                {
+                    T::match_initializer(value.as_ref(), semantic)
+                } else {
+                    false
+                }
             }
 
             // ```python
             // x: annotation = some_expr
             // ```
             //
             // In this situation, we check only the annotation.
-            Some(Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. })) => {
-                T::match_annotation(annotation.as_ref(), semantic)
+            Some(Stmt::AnnAssign(ast::StmtAnnAssign {
+                target, annotation, ..
+            })) => {
+                // TODO(charlie): Replace this with `find_binding_value`, which matches the values.
+                if target.range().contains_range(binding.range()) {
+                    T::match_annotation(annotation.as_ref(), semantic)
+                } else {
+                    false
+                }
             }
             _ => false,
         },
@@ -466,8 +481,15 @@ fn check_type<T: TypeChecker>(binding: &Binding, semantic: &SemanticModel) -> bo
             // ```
             //
             // It's a typed declaration, type annotation is the only source of information.
-            Some(Stmt::AnnAssign(ast::StmtAnnAssign { annotation, .. })) => {
-                T::match_annotation(annotation.as_ref(), semantic)
+            Some(Stmt::AnnAssign(ast::StmtAnnAssign {
+                target, annotation, ..
+            })) => {
+                // TODO(charlie): Replace this with `find_binding_value`, which matches the values.
+                if target.range().contains_range(binding.range()) {
+                    T::match_annotation(annotation.as_ref(), semantic)
+                } else {
+                    false
+                }
             }
             _ => false,
         },