[ty] List available members for a given type (#18251)

This PR adds initial support for listing all attributes of
an object. It is exposed through a new `all_members`
routine in `ty_extensions`, which is in turn used to test
the functionality.

The purpose of listing all members is for code
completion. That is, given a `object.<CURSOR>`, we
would like to list all available attributes on
`object`.

crates/ty_python_semantic/src/semantic_index.rs

@@ -99,7 +99,9 @@ pub(crate) fn use_def_map<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<UseD
     index.use_def_map(scope.file_scope_id(db))
 }
 
-/// Returns all attribute assignments (and their method scope IDs) for a specific class body scope.
+/// Returns all attribute assignments (and their method scope IDs) with a symbol name matching
+/// the one given for a specific class body scope.
+///
 /// Only call this when doing type inference on the same file as `class_body_scope`, otherwise it
 /// introduces a direct dependency on that file's AST.
 pub(crate) fn attribute_assignments<'db, 's>(
@@ -109,6 +111,28 @@ pub(crate) fn attribute_assignments<'db, 's>(
 ) -> impl Iterator<Item = (BindingWithConstraintsIterator<'db, 'db>, FileScopeId)> + use<'s, 'db> {
     let file = class_body_scope.file(db);
     let index = semantic_index(db, file);
+
+    attribute_scopes(db, class_body_scope).filter_map(|function_scope_id| {
+        let attribute_table = index.instance_attribute_table(function_scope_id);
+        let symbol = attribute_table.symbol_id_by_name(name)?;
+        let use_def = &index.use_def_maps[function_scope_id];
+        Some((
+            use_def.instance_attribute_bindings(symbol),
+            function_scope_id,
+        ))
+    })
+}
+
+/// Returns all attribute assignments as scope IDs for a specific class body scope.
+///
+/// Only call this when doing type inference on the same file as `class_body_scope`, otherwise it
+/// introduces a direct dependency on that file's AST.
+pub(crate) fn attribute_scopes<'db, 's>(
+    db: &'db dyn Db,
+    class_body_scope: ScopeId<'db>,
+) -> impl Iterator<Item = FileScopeId> + use<'s, 'db> {
+    let file = class_body_scope.file(db);
+    let index = semantic_index(db, file);
     let class_scope_id = class_body_scope.file_scope_id(db);
 
     ChildrenIter::new(index, class_scope_id).filter_map(|(child_scope_id, scope)| {
@@ -124,13 +148,7 @@ pub(crate) fn attribute_assignments<'db, 's>(
             };
 
         function_scope.node().as_function()?;
-        let attribute_table = index.instance_attribute_table(function_scope_id);
-        let symbol = attribute_table.symbol_id_by_name(name)?;
-        let use_def = &index.use_def_maps[function_scope_id];
-        Some((
-            use_def.instance_attribute_bindings(symbol),
-            function_scope_id,
-        ))
+        Some(function_scope_id)
     })
 }
 
@@ -519,7 +537,7 @@ pub struct ChildrenIter<'a> {
 }
 
 impl<'a> ChildrenIter<'a> {
-    fn new(module_symbol_table: &'a SemanticIndex, parent: FileScopeId) -> Self {
+    pub(crate) fn new(module_symbol_table: &'a SemanticIndex, parent: FileScopeId) -> Self {
         let descendants = DescendantsIter::new(module_symbol_table, parent);
 
         Self {

crates/ty_python_semantic/src/types.rs

@@ -65,6 +65,7 @@ mod context;
 mod diagnostic;
 mod display;
 mod generics;
+mod ide_support;
 mod infer;
 mod instance;
 mod mro;
@@ -7662,6 +7663,8 @@ pub enum KnownFunction {
     GenericContext,
     /// `ty_extensions.dunder_all_names`
     DunderAllNames,
+    /// `ty_extensions.all_members`
+    AllMembers,
 }
 
 impl KnownFunction {
@@ -7721,7 +7724,8 @@ impl KnownFunction {
             | Self::IsSubtypeOf
             | Self::GenericContext
             | Self::DunderAllNames
-            | Self::StaticAssert => module.is_ty_extensions(),
+            | Self::StaticAssert
+            | Self::AllMembers => module.is_ty_extensions(),
         }
     }
 }
@@ -9390,7 +9394,8 @@ pub(crate) mod tests {
                 | KnownFunction::IsSingleValued
                 | KnownFunction::IsAssignableTo
                 | KnownFunction::IsEquivalentTo
-                | KnownFunction::IsGradualEquivalentTo => KnownModule::TyExtensions,
+                | KnownFunction::IsGradualEquivalentTo
+                | KnownFunction::AllMembers => KnownModule::TyExtensions,
             };
 
             let function_definition = known_module_symbol(&db, module, function_name)

crates/ty_python_semantic/src/types/call/bind.rs

@@ -3,6 +3,7 @@
 //! [signatures][crate::types::signatures], we have to handle the fact that the callable might be a
 //! union of types, each of which might contain multiple overloads.
 
+use itertools::Itertools;
 use smallvec::{SmallVec, smallvec};
 
 use super::{
@@ -22,7 +23,8 @@ use crate::types::signatures::{Parameter, ParameterForm};
 use crate::types::{
     BoundMethodType, DataclassParams, DataclassTransformerParams, FunctionDecorators, FunctionType,
     KnownClass, KnownFunction, KnownInstanceType, MethodWrapperKind, PropertyInstanceType,
-    SpecialFormType, TupleType, TypeMapping, UnionType, WrapperDescriptorKind, todo_type,
+    SpecialFormType, TupleType, TypeMapping, UnionType, WrapperDescriptorKind, ide_support,
+    todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, Severity, SubDiagnostic};
 use ruff_python_ast as ast;
@@ -656,6 +658,18 @@ impl<'db> Bindings<'db> {
                             }
                         }
 
+                        Some(KnownFunction::AllMembers) => {
+                            if let [Some(ty)] = overload.parameter_types() {
+                                overload.set_return_type(TupleType::from_elements(
+                                    db,
+                                    ide_support::all_members(db, *ty)
+                                        .into_iter()
+                                        .sorted()
+                                        .map(|member| Type::string_literal(db, &member)),
+                                ));
+                            }
+                        }
+
                         Some(KnownFunction::Len) => {
                             if let [Some(first_arg)] = overload.parameter_types() {
                                 if let Some(len_ty) = first_arg.len(db) {

crates/ty_python_semantic/src/types/ide_support.rs

@@ -0,0 +1,183 @@
+use crate::Db;
+use crate::semantic_index::symbol::ScopeId;
+use crate::semantic_index::{
+    attribute_scopes, global_scope, semantic_index, symbol_table, use_def_map,
+};
+use crate::symbol::{imported_symbol, symbol_from_bindings, symbol_from_declarations};
+use crate::types::{ClassBase, ClassLiteral, KnownClass, Type};
+use ruff_python_ast::name::Name;
+use rustc_hash::FxHashSet;
+
+struct AllMembers {
+    members: FxHashSet<Name>,
+}
+
+impl AllMembers {
+    fn of<'db>(db: &'db dyn Db, ty: Type<'db>) -> Self {
+        let mut all_members = Self {
+            members: FxHashSet::default(),
+        };
+        all_members.extend_with_type(db, ty);
+        all_members
+    }
+
+    fn extend_with_type<'db>(&mut self, db: &'db dyn Db, ty: Type<'db>) {
+        match ty {
+            Type::Union(union) => self.members.extend(
+                union
+                    .elements(db)
+                    .iter()
+                    .map(|ty| AllMembers::of(db, *ty).members)
+                    .reduce(|acc, members| acc.intersection(&members).cloned().collect())
+                    .unwrap_or_default(),
+            ),
+
+            Type::Intersection(intersection) => self.members.extend(
+                intersection
+                    .positive(db)
+                    .iter()
+                    .map(|ty| AllMembers::of(db, *ty).members)
+                    .reduce(|acc, members| acc.union(&members).cloned().collect())
+                    .unwrap_or_default(),
+            ),
+
+            Type::NominalInstance(instance) => {
+                let (class_literal, _specialization) = instance.class.class_literal(db);
+                self.extend_with_class_members(db, class_literal);
+                self.extend_with_instance_members(db, class_literal);
+            }
+
+            Type::ClassLiteral(class_literal) => {
+                self.extend_with_class_members(db, class_literal);
+
+                if let Type::ClassLiteral(meta_class_literal) = ty.to_meta_type(db) {
+                    self.extend_with_class_members(db, meta_class_literal);
+                }
+            }
+
+            Type::GenericAlias(generic_alias) => {
+                let class_literal = generic_alias.origin(db);
+                self.extend_with_class_members(db, class_literal);
+            }
+
+            Type::SubclassOf(subclass_of_type) => {
+                if let Some(class_literal) = subclass_of_type.subclass_of().into_class() {
+                    self.extend_with_class_members(db, class_literal.class_literal(db).0);
+                }
+            }
+
+            Type::Dynamic(_) | Type::Never | Type::AlwaysTruthy | Type::AlwaysFalsy => {}
+
+            Type::IntLiteral(_)
+            | Type::BooleanLiteral(_)
+            | Type::StringLiteral(_)
+            | Type::BytesLiteral(_)
+            | Type::LiteralString
+            | Type::Tuple(_)
+            | Type::PropertyInstance(_)
+            | Type::FunctionLiteral(_)
+            | Type::BoundMethod(_)
+            | Type::MethodWrapper(_)
+            | Type::WrapperDescriptor(_)
+            | Type::DataclassDecorator(_)
+            | Type::DataclassTransformer(_)
+            | Type::Callable(_)
+            | Type::ProtocolInstance(_)
+            | Type::SpecialForm(_)
+            | Type::KnownInstance(_)
+            | Type::TypeVar(_)
+            | Type::BoundSuper(_) => {
+                if let Type::ClassLiteral(class_literal) = ty.to_meta_type(db) {
+                    self.extend_with_class_members(db, class_literal);
+                }
+            }
+
+            Type::ModuleLiteral(literal) => {
+                self.extend_with_type(db, KnownClass::ModuleType.to_instance(db));
+
+                let Some(file) = literal.module(db).file() else {
+                    return;
+                };
+
+                let module_scope = global_scope(db, file);
+                let use_def_map = use_def_map(db, module_scope);
+                let symbol_table = symbol_table(db, module_scope);
+
+                for (symbol_id, _) in use_def_map.all_public_declarations() {
+                    let symbol_name = symbol_table.symbol(symbol_id).name();
+                    if !imported_symbol(db, file, symbol_name, None)
+                        .symbol
+                        .is_unbound()
+                    {
+                        self.members
+                            .insert(symbol_table.symbol(symbol_id).name().clone());
+                    }
+                }
+            }
+        }
+    }
+
+    fn extend_with_declarations_and_bindings(&mut self, db: &dyn Db, scope_id: ScopeId) {
+        let use_def_map = use_def_map(db, scope_id);
+        let symbol_table = symbol_table(db, scope_id);
+
+        for (symbol_id, declarations) in use_def_map.all_public_declarations() {
+            if symbol_from_declarations(db, declarations)
+                .is_ok_and(|result| !result.symbol.is_unbound())
+            {
+                self.members
+                    .insert(symbol_table.symbol(symbol_id).name().clone());
+            }
+        }
+
+        for (symbol_id, bindings) in use_def_map.all_public_bindings() {
+            if !symbol_from_bindings(db, bindings).is_unbound() {
+                self.members
+                    .insert(symbol_table.symbol(symbol_id).name().clone());
+            }
+        }
+    }
+
+    fn extend_with_class_members<'db>(
+        &mut self,
+        db: &'db dyn Db,
+        class_literal: ClassLiteral<'db>,
+    ) {
+        for parent in class_literal
+            .iter_mro(db, None)
+            .filter_map(ClassBase::into_class)
+            .map(|class| class.class_literal(db).0)
+        {
+            let parent_scope = parent.body_scope(db);
+            self.extend_with_declarations_and_bindings(db, parent_scope);
+        }
+    }
+
+    fn extend_with_instance_members<'db>(
+        &mut self,
+        db: &'db dyn Db,
+        class_literal: ClassLiteral<'db>,
+    ) {
+        for parent in class_literal
+            .iter_mro(db, None)
+            .filter_map(ClassBase::into_class)
+            .map(|class| class.class_literal(db).0)
+        {
+            let class_body_scope = parent.body_scope(db);
+            let file = class_body_scope.file(db);
+            let index = semantic_index(db, file);
+            for function_scope_id in attribute_scopes(db, class_body_scope) {
+                let attribute_table = index.instance_attribute_table(function_scope_id);
+                for symbol in attribute_table.symbols() {
+                    self.members.insert(symbol.name().clone());
+                }
+            }
+        }
+    }
+}
+
+/// List all members of a given type: anything that would be valid when accessed
+/// as an attribute on an object of the given type.
+pub(crate) fn all_members<'db>(db: &'db dyn Db, ty: Type<'db>) -> FxHashSet<Name> {
+    AllMembers::of(db, ty).members
+}

crates/ty_python_semantic/src/types/infer.rs

@@ -6719,6 +6719,37 @@ impl<'db> TypeInferenceBuilder<'db> {
         right: Type<'db>,
         range: TextRange,
     ) -> Result<Type<'db>, CompareUnsupportedError<'db>> {
+        let is_str_literal_in_tuple = |literal: Type<'db>, tuple: TupleType<'db>| {
+            // Protect against doing a lot of work for pathologically large
+            // tuples.
+            //
+            // Ref: https://github.com/astral-sh/ruff/pull/18251#discussion_r2115909311
+            if tuple.len(self.db()) > 1 << 12 {
+                return None;
+            }
+
+            let mut definitely_true = false;
+            let mut definitely_false = true;
+            for element in tuple.elements(self.db()) {
+                if element.is_string_literal() {
+                    if literal == *element {
+                        definitely_true = true;
+                        definitely_false = false;
+                    }
+                } else if !literal.is_disjoint_from(self.db(), *element) {
+                    definitely_false = false;
+                }
+            }
+
+            if definitely_true {
+                Some(true)
+            } else if definitely_false {
+                Some(false)
+            } else {
+                None
+            }
+        };
+
         // Note: identity (is, is not) for equal builtin types is unreliable and not part of the
         // language spec.
         // - `[ast::CompOp::Is]`: return `false` if unequal, `bool` if equal
@@ -6850,6 +6881,30 @@ impl<'db> TypeInferenceBuilder<'db> {
                     }
                 }
             }
+            (Type::StringLiteral(_), Type::Tuple(tuple)) if op == ast::CmpOp::In => {
+                if let Some(answer) = is_str_literal_in_tuple(left, tuple) {
+                    return Ok(Type::BooleanLiteral(answer));
+                }
+
+                self.infer_binary_type_comparison(
+                    KnownClass::Str.to_instance(self.db()),
+                    op,
+                    right,
+                    range,
+                )
+            }
+            (Type::StringLiteral(_), Type::Tuple(tuple)) if op == ast::CmpOp::NotIn => {
+                if let Some(answer) = is_str_literal_in_tuple(left, tuple) {
+                    return Ok(Type::BooleanLiteral(!answer));
+                }
+
+                self.infer_binary_type_comparison(
+                    KnownClass::Str.to_instance(self.db()),
+                    op,
+                    right,
+                    range,
+                )
+            }
             (Type::StringLiteral(_), _) => self.infer_binary_type_comparison(
                 KnownClass::Str.to_instance(self.db()),
                 op,