From db34167542ac3e0567d7c7b1bce976a32b5922a3 Mon Sep 17 00:00:00 2001
From: Alex Waygood <alex.waygood@gmail.com>
Date: Thu, 8 May 2025 13:42:44 +0100
Subject: [PATCH] [ty] Support extending `__all__` with a literal tuple or set
 as well as a literal list

---
 .../resources/mdtest/import/dunder_all.md          |  4 +++-
 crates/ty_python_semantic/src/dunder_all.rs        | 14 ++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/import/dunder_all.md b/crates/ty_python_semantic/resources/mdtest/import/dunder_all.md
index 96654b769a3c6..b0b9d7c42702f 100644
--- a/crates/ty_python_semantic/resources/mdtest/import/dunder_all.md
+++ b/crates/ty_python_semantic/resources/mdtest/import/dunder_all.md
@@ -269,6 +269,8 @@ import subexporter
 
 __all__ = []
 __all__.extend(["C", "D"])
+__all__.extend(("E", "F"))
+__all__.extend({"G", "H"})
 __all__.extend(subexporter.__all__)
 
 class C: ...
@@ -281,7 +283,7 @@ class D: ...
 import exporter
 from ty_extensions import dunder_all_names
 
-# revealed: tuple[Literal["A"], Literal["B"], Literal["C"], Literal["D"]]
+# revealed: tuple[Literal["A"], Literal["B"], Literal["C"], Literal["D"], Literal["E"], Literal["F"], Literal["G"], Literal["H"]]
 reveal_type(dunder_all_names(exporter))
 ```
 
diff --git a/crates/ty_python_semantic/src/dunder_all.rs b/crates/ty_python_semantic/src/dunder_all.rs
index 73d1f034fb3bb..b02d5204bf223 100644
--- a/crates/ty_python_semantic/src/dunder_all.rs
+++ b/crates/ty_python_semantic/src/dunder_all.rs
@@ -94,14 +94,16 @@ impl<'db> DunderAllNamesCollector<'db> {
     }
 
     /// Extends the current set of names with the names from the given expression which can be
-    /// either a list of names or a module's `__all__` variable.
+    /// either a list/tuple/set of string-literal names or a module's `__all__` variable.
     ///
-    /// Returns `true` if the expression is a valid list or module `__all__`, `false` otherwise.
-    fn extend_from_list_or_module(&mut self, expr: &ast::Expr) -> bool {
+    /// Returns `true` if the expression is a valid list/tuple/set or module `__all__`, `false` otherwise.
+    fn extend(&mut self, expr: &ast::Expr) -> bool {
         match expr {
             // `__all__ += [...]`
             // `__all__.extend([...])`
-            ast::Expr::List(ast::ExprList { elts, .. }) => self.add_names(elts),
+            ast::Expr::List(ast::ExprList { elts, .. })
+            | ast::Expr::Tuple(ast::ExprTuple { elts, .. })
+            | ast::Expr::Set(ast::ExprSet { elts, .. }) => self.add_names(elts),
 
             // `__all__ += module.__all__`
             // `__all__.extend(module.__all__)`
@@ -155,7 +157,7 @@ impl<'db> DunderAllNamesCollector<'db> {
             // `__all__.extend([...])`
             // `__all__.extend(module.__all__)`
             "extend" => {
-                if !self.extend_from_list_or_module(argument) {
+                if !self.extend(argument) {
                     return false;
                 }
             }
@@ -330,7 +332,7 @@ impl<'db> StatementVisitor<'db> for DunderAllNamesCollector<'db> {
                 if !is_dunder_all(target) {
                     return;
                 }
-                if !self.extend_from_list_or_module(value) {
+                if !self.extend(value) {
                     self.invalid = true;
                 }
             }