diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index a19daed9db..9a6190cb52 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -5,7 +5,7 @@ use ruff_db::parsed::{ParsedModuleRef, parsed_module}; use ruff_python_ast as ast; use ruff_python_parser::{Token, TokenAt, TokenKind}; use ruff_text_size::{Ranged, TextRange, TextSize}; -use ty_python_semantic::{Completion, SemanticModel}; +use ty_python_semantic::{Completion, NameKind, SemanticModel}; use crate::Db; use crate::find_node::covering_node; @@ -325,38 +325,7 @@ fn import_from_tokens(tokens: &[Token]) -> Option<&Token> { /// This has the effect of putting all dunder attributes after "normal" /// attributes, and all single-underscore attributes after dunder attributes. fn compare_suggestions(c1: &Completion, c2: &Completion) -> Ordering { - /// A helper type for sorting completions based only on name. - /// - /// This sorts "normal" names first, then dunder names and finally - /// single-underscore names. This matches the order of the variants defined for - /// this enum, which is in turn picked up by the derived trait implementation - /// for `Ord`. - #[derive(Clone, Copy, Eq, PartialEq, PartialOrd, Ord)] - enum Kind { - Normal, - Dunder, - Sunder, - } - - impl Kind { - fn classify(c: &Completion) -> Kind { - // Dunder needs a prefix and suffix double underscore. - // When there's only a prefix double underscore, this - // results in explicit name mangling. We let that be - // classified as-if they were single underscore names. - // - // Ref: - if c.name.starts_with("__") && c.name.ends_with("__") { - Kind::Dunder - } else if c.name.starts_with('_') { - Kind::Sunder - } else { - Kind::Normal - } - } - } - - let (kind1, kind2) = (Kind::classify(c1), Kind::classify(c2)); + let (kind1, kind2) = (NameKind::classify(&c1.name), NameKind::classify(&c2.name)); kind1.cmp(&kind2).then_with(|| c1.name.cmp(&c2.name)) } @@ -472,6 +441,11 @@ mod tests { ", ); test.assert_completions_include("filter"); + // Sunder items should be filtered out + test.assert_completions_do_not_include("_T"); + // Dunder attributes should not be stripped + test.assert_completions_include("__annotations__"); + // See `private_symbols_in_stub` for more comprehensive testing private of symbol filtering. } #[test] @@ -536,6 +510,112 @@ re. test.assert_completions_include("findall"); } + #[test] + fn private_symbols_in_stub() { + let test = CursorTest::builder() + .source( + "package/__init__.pyi", + r#"\ +from typing import TypeAlias, Literal, TypeVar, ParamSpec, TypeVarTuple, Protocol + +public_name = 1 +_private_name = 1 +__mangled_name = 1 +__dunder_name__ = 1 + +public_type_var = TypeVar("public_type_var") +_private_type_var = TypeVar("_private_type_var") +__mangled_type_var = TypeVar("__mangled_type_var") + +public_param_spec = ParamSpec("public_param_spec") +_private_param_spec = ParamSpec("_private_param_spec") + +public_type_var_tuple = TypeVarTuple("public_type_var_tuple") +_private_type_var_tuple = TypeVarTuple("_private_type_var_tuple") + +public_explicit_type_alias: TypeAlias = Literal[1] +_private_explicit_type_alias: TypeAlias = Literal[1] + +class PublicProtocol(Protocol): + def method(self) -> None: ... + +class _PrivateProtocol(Protocol): + def method(self) -> None: ... +"#, + ) + .source("main.py", "import package; package.") + .build(); + test.assert_completions_include("public_name"); + test.assert_completions_include("_private_name"); + test.assert_completions_include("__mangled_name"); + test.assert_completions_include("__dunder_name__"); + test.assert_completions_include("public_type_var"); + test.assert_completions_do_not_include("_private_type_var"); + test.assert_completions_do_not_include("__mangled_type_var"); + test.assert_completions_include("public_param_spec"); + test.assert_completions_do_not_include("_private_param_spec"); + test.assert_completions_include("public_type_var_tuple"); + test.assert_completions_do_not_include("_private_type_var_tuple"); + test.assert_completions_include("public_explicit_type_alias"); + test.assert_completions_include("_private_explicit_type_alias"); + test.assert_completions_include("PublicProtocol"); + test.assert_completions_do_not_include("_PrivateProtocol"); + } + + /// Unlike [`private_symbols_in_stub`], this test doesn't use a `.pyi` file so all of the names + /// are visible. + #[test] + fn private_symbols_in_module() { + let test = CursorTest::builder() + .source( + "package/__init__.py", + r#"\ +from typing import TypeAlias, Literal, TypeVar, ParamSpec, TypeVarTuple, Protocol + +public_name = 1 +_private_name = 1 +__mangled_name = 1 +__dunder_name__ = 1 + +public_type_var = TypeVar("public_type_var") +_private_type_var = TypeVar("_private_type_var") +__mangled_type_var = TypeVar("__mangled_type_var") + +public_param_spec = ParamSpec("public_param_spec") +_private_param_spec = ParamSpec("_private_param_spec") + +public_type_var_tuple = TypeVarTuple("public_type_var_tuple") +_private_type_var_tuple = TypeVarTuple("_private_type_var_tuple") + +public_explicit_type_alias: TypeAlias = Literal[1] +_private_explicit_type_alias: TypeAlias = Literal[1] + +class PublicProtocol(Protocol): + def method(self) -> None: ... + +class _PrivateProtocol(Protocol): + def method(self) -> None: ... +"#, + ) + .source("main.py", "import package; package.") + .build(); + test.assert_completions_include("public_name"); + test.assert_completions_include("_private_name"); + test.assert_completions_include("__mangled_name"); + test.assert_completions_include("__dunder_name__"); + test.assert_completions_include("public_type_var"); + test.assert_completions_include("_private_type_var"); + test.assert_completions_include("__mangled_type_var"); + test.assert_completions_include("public_param_spec"); + test.assert_completions_include("_private_param_spec"); + test.assert_completions_include("public_type_var_tuple"); + test.assert_completions_include("_private_type_var_tuple"); + test.assert_completions_include("public_explicit_type_alias"); + test.assert_completions_include("_private_explicit_type_alias"); + test.assert_completions_include("PublicProtocol"); + test.assert_completions_include("_PrivateProtocol"); + } + #[test] fn one_function_prefix() { let test = cursor_test( diff --git a/crates/ty_python_semantic/src/lib.rs b/crates/ty_python_semantic/src/lib.rs index 85a1a6ec5d..f0537886a9 100644 --- a/crates/ty_python_semantic/src/lib.rs +++ b/crates/ty_python_semantic/src/lib.rs @@ -15,7 +15,7 @@ pub use program::{ PythonVersionWithSource, SearchPathSettings, }; pub use python_platform::PythonPlatform; -pub use semantic_model::{Completion, HasType, SemanticModel}; +pub use semantic_model::{Completion, HasType, NameKind, SemanticModel}; pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin}; pub use util::diagnostics::add_inferred_python_version_hint_to_diagnostic; diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index cff4a3fe76..283549c821 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/crates/ty_python_semantic/src/semantic_model.rs @@ -68,12 +68,10 @@ impl<'db> SemanticModel<'db> { return vec![]; }; let ty = Type::module_literal(self.db, self.file, &module); + let builtin = module.is_known(KnownModule::Builtins); crate::types::all_members(self.db, ty) .into_iter() - .map(|name| Completion { - name, - builtin: module.is_known(KnownModule::Builtins), - }) + .map(|name| Completion { name, builtin }) .collect() } @@ -130,6 +128,39 @@ impl<'db> SemanticModel<'db> { } } +/// A classification of symbol names. +/// +/// The ordering here is used for sorting completions. +/// +/// This sorts "normal" names first, then dunder names and finally +/// single-underscore names. This matches the order of the variants defined for +/// this enum, which is in turn picked up by the derived trait implementation +/// for `Ord`. +#[derive(Clone, Copy, Eq, PartialEq, PartialOrd, Ord)] +pub enum NameKind { + Normal, + Dunder, + Sunder, +} + +impl NameKind { + pub fn classify(name: &Name) -> NameKind { + // Dunder needs a prefix and suffix double underscore. + // When there's only a prefix double underscore, this + // results in explicit name mangling. We let that be + // classified as-if they were single underscore names. + // + // Ref: + if name.starts_with("__") && name.ends_with("__") { + NameKind::Dunder + } else if name.starts_with('_') { + NameKind::Sunder + } else { + NameKind::Normal + } + } +} + /// A suggestion for code completion. #[derive(Clone, Debug)] pub struct Completion { diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index 7cf05fa57f..84675ce56f 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -1,10 +1,10 @@ -use crate::Db; -use crate::place::{imported_symbol, place_from_bindings, place_from_declarations}; +use crate::place::{Place, imported_symbol, place_from_bindings, place_from_declarations}; use crate::semantic_index::place::ScopeId; use crate::semantic_index::{ attribute_scopes, global_scope, imported_modules, place_table, semantic_index, use_def_map, }; -use crate::types::{ClassBase, ClassLiteral, KnownClass, Type}; +use crate::types::{ClassBase, ClassLiteral, KnownClass, KnownInstanceType, Type}; +use crate::{Db, NameKind}; use ruff_python_ast::name::Name; use rustc_hash::FxHashSet; @@ -144,13 +144,41 @@ impl AllMembers { let Some(symbol_name) = place_table.place_expr(symbol_id).as_name() else { continue; }; - if !imported_symbol(db, file, symbol_name, None) - .place - .is_unbound() - { - self.members - .insert(place_table.place_expr(symbol_id).expect_name().clone()); + let Place::Type(ty, _) = imported_symbol(db, file, symbol_name, None).place + else { + continue; + }; + + // Filter private symbols from stubs if they appear to be internal types + let is_stub_file = file.path(db).extension() == Some("pyi"); + let is_private_symbol = match NameKind::classify(symbol_name) { + NameKind::Dunder | NameKind::Normal => false, + NameKind::Sunder => true, + }; + if is_private_symbol && is_stub_file { + match ty { + Type::NominalInstance(instance) + if matches!( + instance.class.known(db), + Some( + KnownClass::TypeVar + | KnownClass::TypeVarTuple + | KnownClass::ParamSpec + ) + ) => + { + continue; + } + Type::ClassLiteral(class) if class.is_protocol(db) => continue, + Type::KnownInstance( + KnownInstanceType::TypeVar(_) | KnownInstanceType::TypeAliasType(_), + ) => continue, + _ => {} + } } + + self.members + .insert(place_table.place_expr(symbol_id).expect_name().clone()); } let module_name = module.name();