From 97b8d0f61f992ccfcd0a572690d7e7bdada6ae10 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Sun, 7 Apr 2024 00:16:06 +0100
Subject: [PATCH] [`flake8-slots`] Flag subclasses of call-based
 `typing.NamedTuple`s as well as subclasses of `collections.namedtuple()`
 (`SLOT002`) (#10808)

---
 .../test/fixtures/flake8_slots/SLOT002.py     |  8 +++
 .../rules/no_slots_in_namedtuple_subclass.rs  | 66 +++++++++++++------
 ...ake8_slots__tests__SLOT002_SLOT002.py.snap |  7 +-
 3 files changed, 60 insertions(+), 21 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_slots/SLOT002.py b/crates/ruff_linter/resources/test/fixtures/flake8_slots/SLOT002.py
index 59a9ba3f5579bf..11f2782b77ac25 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_slots/SLOT002.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_slots/SLOT002.py
@@ -6,6 +6,14 @@ class Bad(namedtuple("foo", ["str", "int"])):  # SLOT002
     pass
 
 
+class UnusualButStillBad(NamedTuple("foo", [("x", int, "y", int)])):  # SLOT002
+    pass
+
+
+class UnusualButOkay(NamedTuple("foo", [("x", int, "y", int)])):
+    __slots__ = ()
+
+
 class Good(namedtuple("foo", ["str", "int"])):  # OK
     __slots__ = ("foo",)
 
diff --git a/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_namedtuple_subclass.rs b/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_namedtuple_subclass.rs
index 187e0bcdf9062a..fdbf87efcc74d2 100644
--- a/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_namedtuple_subclass.rs
+++ b/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_namedtuple_subclass.rs
@@ -1,17 +1,16 @@
-use ruff_python_ast as ast;
-use ruff_python_ast::{Arguments, Expr, StmtClassDef};
+use std::fmt;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::identifier::Identifier;
-use ruff_python_ast::Stmt;
+use ruff_python_ast::{self as ast, identifier::Identifier, Arguments, Expr, Stmt, StmtClassDef};
+use ruff_python_semantic::SemanticModel;
 
 use crate::checkers::ast::Checker;
 use crate::rules::flake8_slots::rules::helpers::has_slots;
 
 /// ## What it does
-/// Checks for subclasses of `collections.namedtuple` that lack a `__slots__`
-/// definition.
+/// Checks for subclasses of `collections.namedtuple` or `typing.NamedTuple`
+/// that lack a `__slots__` definition.
 ///
 /// ## Why is this bad?
 /// In Python, the `__slots__` attribute allows you to explicitly define the
@@ -48,12 +47,28 @@ use crate::rules::flake8_slots::rules::helpers::has_slots;
 /// ## References
 /// - [Python documentation: `__slots__`](https://docs.python.org/3/reference/datamodel.html#slots)
 #[violation]
-pub struct NoSlotsInNamedtupleSubclass;
+pub struct NoSlotsInNamedtupleSubclass(NamedTupleKind);
 
 impl Violation for NoSlotsInNamedtupleSubclass {
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Subclasses of `collections.namedtuple()` should define `__slots__`")
+        let NoSlotsInNamedtupleSubclass(namedtuple_kind) = self;
+        format!("Subclasses of {namedtuple_kind} should define `__slots__`")
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+enum NamedTupleKind {
+    Collections,
+    Typing,
+}
+
+impl fmt::Display for NamedTupleKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Self::Collections => "`collections.namedtuple()`",
+            Self::Typing => "call-based `typing.NamedTuple()`",
+        })
     }
 }
 
@@ -67,22 +82,33 @@ pub(crate) fn no_slots_in_namedtuple_subclass(
         return;
     };
 
-    if bases.iter().any(|base| {
-        let Expr::Call(ast::ExprCall { func, .. }) = base else {
-            return false;
-        };
-        checker
-            .semantic()
-            .resolve_qualified_name(func)
-            .is_some_and(|qualified_name| {
-                matches!(qualified_name.segments(), ["collections", "namedtuple"])
-            })
-    }) {
+    if let Some(namedtuple_kind) = namedtuple_base(bases, checker.semantic()) {
         if !has_slots(&class.body) {
             checker.diagnostics.push(Diagnostic::new(
-                NoSlotsInNamedtupleSubclass,
+                NoSlotsInNamedtupleSubclass(namedtuple_kind),
                 stmt.identifier(),
             ));
         }
     }
 }
+
+/// If the class has a call-based namedtuple in its bases,
+/// return the kind of namedtuple it is
+/// (either `collections.namedtuple()`, or `typing.NamedTuple()`).
+/// Else, return `None`.
+fn namedtuple_base(bases: &[Expr], semantic: &SemanticModel) -> Option<NamedTupleKind> {
+    for base in bases {
+        let Expr::Call(ast::ExprCall { func, .. }) = base else {
+            continue;
+        };
+        let Some(qualified_name) = semantic.resolve_qualified_name(func) else {
+            continue;
+        };
+        match qualified_name.segments() {
+            ["collections", "namedtuple"] => return Some(NamedTupleKind::Collections),
+            ["typing", "NamedTuple"] => return Some(NamedTupleKind::Typing),
+            _ => continue,
+        }
+    }
+    None
+}
diff --git a/crates/ruff_linter/src/rules/flake8_slots/snapshots/ruff_linter__rules__flake8_slots__tests__SLOT002_SLOT002.py.snap b/crates/ruff_linter/src/rules/flake8_slots/snapshots/ruff_linter__rules__flake8_slots__tests__SLOT002_SLOT002.py.snap
index 4ee1ad407bf3df..d7670abd56fe8c 100644
--- a/crates/ruff_linter/src/rules/flake8_slots/snapshots/ruff_linter__rules__flake8_slots__tests__SLOT002_SLOT002.py.snap
+++ b/crates/ruff_linter/src/rules/flake8_slots/snapshots/ruff_linter__rules__flake8_slots__tests__SLOT002_SLOT002.py.snap
@@ -8,4 +8,9 @@ SLOT002.py:5:7: SLOT002 Subclasses of `collections.namedtuple()` should define `
 6 |     pass
   |
 
-
+SLOT002.py:9:7: SLOT002 Subclasses of call-based `typing.NamedTuple()` should define `__slots__`
+   |
+ 9 | class UnusualButStillBad(NamedTuple("foo", [("x", int, "y", int)])):  # SLOT002
+   |       ^^^^^^^^^^^^^^^^^^ SLOT002
+10 |     pass
+   |