Assume type of self is typing.Self

Author: Glyphack

crates/ty_python_semantic/resources/mdtest/annotations/self.md

@@ -29,12 +29,10 @@ class Shape:
         return inner()
 
     def implicit_self(self) -> Self:
-        # TODO: first argument in a method should be considered as "typing.Self"
         reveal_type(self)  # revealed: Unknown
         return self
 
-# TODO: should be `list[Shape]`
-reveal_type(Shape().nested_type())  # revealed: list[Self]
+reveal_type(Shape().nested_type())  # revealed: list[Shape]
 
 reveal_type(Shape().nested_func())  # revealed: Shape
 

crates/ty_python_semantic/src/types.rs

@@ -1,4 +1,4 @@
-use infer::enclosing_class_symbol;
+use infer::enclosing_class_definition;
 use itertools::Either;
 
 use std::slice::Iter;
@@ -4798,21 +4798,17 @@ impl<'db> Type<'db> {
 
                 KnownInstanceType::TypingSelf => {
                     let index = semantic_index(db, scope_id.file(db));
-                    let Some(class_ty) = enclosing_class_symbol(db, index, scope_id) else {
+                    let Some(class_def) = enclosing_class_definition(db, index, scope_id) else {
                         return Err(InvalidTypeExpressionError {
                             fallback_type: Type::unknown(),
                             invalid_expressions: smallvec::smallvec![
                                 InvalidTypeExpression::InvalidType(*self)
                             ],
                         });
                     };
-                    let Some(TypeDefinition::Class(class_def)) = class_ty.definition(db) else {
-                        debug_assert!(
-                            false,
-                            "enclosing_class_symbol must return a type with class definition"
-                        );
-                        return Ok(Type::unknown());
-                    };
+                    let class_ty = infer_definition_types(db, class_def)
+                        .declaration_type(class_def)
+                        .inner_type();
                     let Some(instance) = class_ty.to_instance(db) else {
                         debug_assert!(
                             false,

crates/ty_python_semantic/src/types/infer.rs

@@ -317,31 +317,45 @@ pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> U
     unpacker.finish()
 }
 
-/// Returns the type of the nearest enclosing class for the given scope.
+/// Returns the definition of the nearest enclosing class for the given scope.
 ///
 /// This function walks up the ancestor scopes starting from the given scope,
 /// and finds the closest class definition.
 ///
 /// Returns `None` if no enclosing class is found.a
-pub(crate) fn enclosing_class_symbol<'db>(
+pub(crate) fn enclosing_class_definition<'db>(
     db: &'db dyn Db,
     semantic: &SemanticIndex<'db>,
     scope: ScopeId,
-) -> Option<Type<'db>> {
+) -> Option<Definition<'db>> {
     semantic
         .ancestor_scopes(scope.file_scope_id(db))
         .find_map(|(_, ancestor_scope)| {
             if let NodeWithScopeKind::Class(class) = ancestor_scope.node() {
                 let definition = semantic.expect_single_definition(class.node());
-                let result = infer_definition_types(db, definition);
-
-                Some(result.declaration_type(definition).inner_type())
+                Some(definition)
             } else {
                 None
             }
         })
 }
 
+/// Returns the type of the nearest enclosing class for the given scope.
+///
+/// This function walks up the ancestor scopes starting from the given scope,
+/// and finds the closest class definition.
+///
+/// Returns `None` if no enclosing class is found.a
+pub(crate) fn enclosing_class_type<'db>(
+    db: &'db dyn Db,
+    semantic: &SemanticIndex<'db>,
+    scope: ScopeId,
+) -> Option<Type<'db>> {
+    let definition = enclosing_class_definition(db, semantic, scope)?;
+    let result = infer_definition_types(db, definition);
+    return Some(result.declaration_type(definition).inner_type());
+}
+
 /// A region within which we can infer types.
 #[derive(Copy, Clone, Debug)]
 pub(crate) enum InferenceRegion<'db> {
@@ -4975,7 +4989,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                                             [] => {
                                                 let scope = self.scope();
 
-                                                let Some(enclosing_class) = enclosing_class_symbol(
+                                                let Some(enclosing_class) = enclosing_class_type(
                                                     self.db(),
                                                     self.index,
                                                     scope,

crates/ty_python_semantic/src/types/signatures.rs

@@ -18,7 +18,7 @@ use smallvec::{smallvec, SmallVec};
 use super::{definition_expression_type, DynamicType, Type};
 use crate::semantic_index::definition::Definition;
 use crate::types::generics::{GenericContext, Specialization, TypeMapping};
-use crate::types::{todo_type, ClassLiteral, TypeVarInstance};
+use crate::types::{todo_type, ClassLiteral, KnownInstanceType, TypeVarInstance};
 use crate::{Db, FxOrderSet};
 use ruff_python_ast::{self as ast, name::Name};
 
@@ -1032,16 +1032,32 @@ impl<'db> Parameters<'db> {
                 },
             )
         });
-        let positional_or_keyword = args.iter().map(|arg| {
-            Parameter::from_node_and_kind(
-                db,
-                definition,
-                &arg.parameter,
-                ParameterKind::PositionalOrKeyword {
-                    name: arg.parameter.name.id.clone(),
-                    default_type: default_type(arg),
-                },
-            )
+        let positional_or_keyword = args.iter().enumerate().map(|(index, arg)| {
+            // TODO: Check for classmethod decorator
+            if index == 0 && arg.name() == "self" {
+                Parameter {
+                    annotated_type: Some(
+                        Type::KnownInstance(KnownInstanceType::TypingSelf)
+                            .in_type_expression(db, definition.scope(db))
+                            .unwrap(),
+                    ),
+                    kind: ParameterKind::PositionalOrKeyword {
+                        name: arg.parameter.name.id.clone(),
+                        default_type: default_type(arg),
+                    },
+                    form: ParameterForm::Value,
+                }
+            } else {
+                Parameter::from_node_and_kind(
+                    db,
+                    definition,
+                    &arg.parameter,
+                    ParameterKind::PositionalOrKeyword {
+                        name: arg.parameter.name.id.clone(),
+                        default_type: default_type(arg),
+                    },
+                )
+            }
         });
         let variadic = vararg.as_ref().map(|arg| {
             Parameter::from_node_and_kind(