Filter private symbols from stubs if they are internal types (#19121)

This implements filtering of private symbols from stub files based on
type information as discussed in
#19102. It extends the previous
implementation to apply to all stub files, instead of just the
`builtins` module, and uses type information to retain private names
that are may be relevant at runtime.

Author: zanieb

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: <https://docs.python.org/3/reference/lexical_analysis.html#reserved-classes-of-identifiers>
-            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.<CURSOR>
         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.<CURSOR>")
+            .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.<CURSOR>")
+            .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(

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;
 

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: <https://docs.python.org/3/reference/lexical_analysis.html#reserved-classes-of-identifiers>
+        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 {

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();