[ty] Support extending __all__ from an imported module even when the module is not an ExprName node (#17947)

crates/ty_python_semantic/src/dunder_all.rs

@@ -6,11 +6,10 @@ use ruff_python_ast::name::Name;
 use ruff_python_ast::statement_visitor::{walk_stmt, StatementVisitor};
 use ruff_python_ast::{self as ast};
 
-use crate::semantic_index::ast_ids::{HasScopedExpressionId, HasScopedUseId};
+use crate::semantic_index::ast_ids::HasScopedExpressionId;
 use crate::semantic_index::symbol::ScopeId;
 use crate::semantic_index::{global_scope, semantic_index, SemanticIndex};
-use crate::symbol::{symbol_from_bindings, Boundness, Symbol};
-use crate::types::{infer_expression_types, Truthiness};
+use crate::types::{infer_expression_types, Truthiness, Type};
 use crate::{resolve_module, Db, ModuleName};
 
 #[allow(clippy::ref_option)]
@@ -111,18 +110,8 @@ impl<'db> DunderAllNamesCollector<'db> {
                 if attr != "__all__" {
                     return false;
                 }
-                let Some(name_node) = value.as_name_expr() else {
-                    return false;
-                };
-                let Symbol::Type(ty, Boundness::Bound) = symbol_from_bindings(
-                    self.db,
-                    self.index
-                        .use_def_map(self.scope.file_scope_id(self.db))
-                        .bindings_at_use(name_node.scoped_use_id(self.db, self.scope)),
-                ) else {
-                    return false;
-                };
-                let Some(module_literal) = ty.into_module_literal() else {
+                let Type::ModuleLiteral(module_literal) = self.standalone_expression_type(value)
+                else {
                     return false;
                 };
                 let Some(module_dunder_all_names) =
@@ -198,14 +187,21 @@ impl<'db> DunderAllNamesCollector<'db> {
         dunder_all_names(self.db, module.file())
     }
 
+    /// Infer the type of a standalone expression.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `expr` was not marked as a standalone expression during semantic indexing.
+    fn standalone_expression_type(&self, expr: &ast::Expr) -> Type<'db> {
+        infer_expression_types(self.db, self.index.expression(expr))
+            .expression_type(expr.scoped_expression_id(self.db, self.scope))
+    }
+
     /// Evaluate the given expression and return its truthiness.
     ///
     /// Returns [`None`] if the expression type doesn't implement `__bool__` correctly.
     fn evaluate_test_expr(&self, expr: &ast::Expr) -> Option<Truthiness> {
-        infer_expression_types(self.db, self.index.expression(expr))
-            .expression_type(expr.scoped_expression_id(self.db, self.scope))
-            .try_bool(self.db)
-            .ok()
+        self.standalone_expression_type(expr).try_bool(self.db).ok()
     }
 
     /// Add valid names to the set.

crates/ty_python_semantic/src/semantic_index.rs

@@ -387,6 +387,14 @@ impl<'db> SemanticIndex<'db> {
             .copied()
     }
 
+    pub(crate) fn is_standalone_expression(
+        &self,
+        expression_key: impl Into<ExpressionNodeKey>,
+    ) -> bool {
+        self.expressions_by_node
+            .contains_key(&expression_key.into())
+    }
+
     /// Returns the id of the scope that `node` creates.
     /// This is different from [`definition::Definition::scope`] which
     /// returns the scope in which that definition is defined in.

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -1479,23 +1479,37 @@ where
                 aug_assign @ ast::StmtAugAssign {
                     range: _,
                     target,
-                    op: _,
+                    op,
                     value,
                 },
             ) => {
                 debug_assert_eq!(&self.current_assignments, &[]);
                 self.visit_expr(value);
 
-                // See https://docs.python.org/3/library/ast.html#ast.AugAssign
-                if matches!(
-                    **target,
-                    ast::Expr::Attribute(_) | ast::Expr::Subscript(_) | ast::Expr::Name(_)
-                ) {
-                    self.push_assignment(aug_assign.into());
-                    self.visit_expr(target);
-                    self.pop_assignment();
-                } else {
-                    self.visit_expr(target);
+                match &**target {
+                    ast::Expr::Name(ast::ExprName { id, .. })
+                        if id == "__all__" && op.is_add() && self.in_module_scope() =>
+                    {
+                        if let ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) =
+                            &**value
+                        {
+                            if attr == "__all__" {
+                                self.add_standalone_expression(value);
+                            }
+                        }
+
+                        self.push_assignment(aug_assign.into());
+                        self.visit_expr(target);
+                        self.pop_assignment();
+                    }
+                    ast::Expr::Name(_) | ast::Expr::Attribute(_) | ast::Expr::Subscript(_) => {
+                        self.push_assignment(aug_assign.into());
+                        self.visit_expr(target);
+                        self.pop_assignment();
+                    }
+                    _ => {
+                        self.visit_expr(target);
+                    }
                 }
             }
             ast::Stmt::If(node) => {
@@ -1934,6 +1948,12 @@ where
                 }
                 walk_stmt(self, stmt);
             }
+            ast::Stmt::Expr(ast::StmtExpr { value, range: _ }) if self.in_module_scope() => {
+                if let Some(expr) = dunder_all_extend_argument(value) {
+                    self.add_standalone_expression(expr);
+                }
+                self.visit_expr(value);
+            }
             _ => {
                 walk_stmt(self, stmt);
             }
@@ -2623,3 +2643,43 @@ impl<'a> Unpackable<'a> {
         }
     }
 }
+
+/// Returns the single argument to `__all__.extend()`, if it is a call to `__all__.extend()`
+/// where it looks like the argument might be a `submodule.__all__` expression.
+/// Else, returns `None`.
+fn dunder_all_extend_argument(value: &ast::Expr) -> Option<&ast::Expr> {
+    let ast::ExprCall {
+        func,
+        arguments:
+            ast::Arguments {
+                args,
+                keywords,
+                range: _,
+            },
+        ..
+    } = value.as_call_expr()?;
+
+    let ast::ExprAttribute { value, attr, .. } = func.as_attribute_expr()?;
+
+    let ast::ExprName { id, .. } = value.as_name_expr()?;
+
+    if id != "__all__" {
+        return None;
+    }
+
+    if attr != "extend" {
+        return None;
+    }
+
+    if !keywords.is_empty() {
+        return None;
+    }
+
+    let [single_argument] = &**args else {
+        return None;
+    };
+
+    let ast::ExprAttribute { value, attr, .. } = single_argument.as_attribute_expr()?;
+
+    (attr == "__all__").then_some(value)
+}

crates/ty_python_semantic/src/types/infer.rs

@@ -5505,7 +5505,12 @@ impl<'db> TypeInferenceBuilder<'db> {
             ctx: _,
         } = attribute;
 
-        let value_type = self.infer_expression(value);
+        let value_type = if self.index.is_standalone_expression(&**value) {
+            self.infer_standalone_expression(value)
+        } else {
+            self.infer_expression(value)
+        };
+
         let db = self.db();
 
         value_type