Support typing.Self in methods (#17689)

## Summary

Fixes: astral-sh/ty#159 

This PR adds support for using `Self` in methods.
When the type of an annotation is `TypingSelf` it is converted to a type
var based on:
https://typing.python.org/en/latest/spec/generics.html#self

I just skipped Protocols because it had more problems and the tests was
not useful.
Also I need to create a follow up PR that implicitly assumes `self`
argument has type `Self`.

In order to infer the type in the `in_type_expression` method I needed
to have scope id and semantic index available. I used the idea from
[this PR](https://github.com/astral-sh/ruff/pull/17589/files) to pass
additional context to this method.
Also I think in all places that `in_type_expression` is called we need
to have this context because `Self` can be there so I didn't split the
method into one version with context and one without.

## Test Plan

Added new tests from spec.

---------

Co-authored-by: Micha Reiser <micha@reiser.io>
Co-authored-by: Carl Meyer <carl@astral.sh>

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

@@ -0,0 +1,190 @@
+# Self
+
+`Self` is treated as if it were a `TypeVar` bound to the class it's being used on.
+
+`typing.Self` is only available in Python 3.11 and later.
+
+## Methods
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self
+
+class Shape:
+    def set_scale(self: Self, scale: float) -> Self:
+        reveal_type(self)  # revealed: Self
+        return self
+
+    def nested_type(self) -> list[Self]:
+        return [self]
+
+    def nested_func(self: Self) -> Self:
+        def inner() -> Self:
+            reveal_type(self)  # revealed: Self
+            return self
+        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
+
+reveal_type(Shape().nested_type())  # revealed: @Todo(specialized non-generic class)
+reveal_type(Shape().nested_func())  # revealed: Shape
+
+class Circle(Shape):
+    def set_scale(self: Self, scale: float) -> Self:
+        reveal_type(self)  # revealed: Self
+        return self
+
+class Outer:
+    class Inner:
+        def foo(self: Self) -> Self:
+            reveal_type(self)  # revealed: Self
+            return self
+```
+
+## Class Methods
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self, TypeVar
+
+class Shape:
+    def foo(self: Self) -> Self:
+        return self
+
+    @classmethod
+    def bar(cls: type[Self]) -> Self:
+        # TODO: type[Shape]
+        reveal_type(cls)  # revealed: @Todo(unsupported type[X] special form)
+        return cls()
+
+class Circle(Shape): ...
+
+reveal_type(Shape().foo())  # revealed: Shape
+# TODO: Shape
+reveal_type(Shape.bar())  # revealed: Unknown
+```
+
+## Attributes
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self
+
+class LinkedList:
+    value: int
+    next_node: Self
+
+    def next(self: Self) -> Self:
+        reveal_type(self.value)  # revealed: int
+        return self.next_node
+
+reveal_type(LinkedList().next())  # revealed: LinkedList
+```
+
+## Generic Classes
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self, Generic, TypeVar
+
+T = TypeVar("T")
+
+class Container(Generic[T]):
+    value: T
+    def set_value(self: Self, value: T) -> Self:
+        return self
+
+int_container: Container[int] = Container[int]()
+reveal_type(int_container)  # revealed: Container[int]
+reveal_type(int_container.set_value(1))  # revealed: Container[int]
+```
+
+## Protocols
+
+TODO: <https://typing.python.org/en/latest/spec/generics.html#use-in-protocols>
+
+## Annotations
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self
+
+class Shape:
+    def union(self: Self, other: Self | None):
+        reveal_type(other)  # revealed: Self | None
+        return self
+```
+
+## Invalid Usage
+
+`Self` cannot be used in the signature of a function or variable.
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+from typing import Self, Generic, TypeVar
+
+T = TypeVar("T")
+
+# error: [invalid-type-form]
+def x(s: Self): ...
+
+# error: [invalid-type-form]
+b: Self
+
+# TODO: "Self" cannot be used in a function with a `self` or `cls` parameter that has a type annotation other than "Self"
+class Foo:
+    # TODO: rejected Self because self has a different type
+    def has_existing_self_annotation(self: T) -> Self:
+        return self  # error: [invalid-return-type]
+
+    def return_concrete_type(self) -> Self:
+        # TODO: tell user to use "Foo" instead of "Self"
+        # error: [invalid-return-type]
+        return Foo()
+
+    @staticmethod
+    # TODO: reject because of staticmethod
+    def make() -> Self:
+        # error: [invalid-return-type]
+        return Foo()
+
+class Bar(Generic[T]):
+    foo: T
+    def bar(self) -> T:
+        return self.foo
+
+# error: [invalid-type-form]
+class Baz(Bar[Self]): ...
+
+class MyMetaclass(type):
+    # TODO: rejected
+    def __new__(cls) -> Self:
+        return super().__new__(cls)
+```

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

@@ -30,7 +30,7 @@ def i(callback: Callable[Concatenate[int, P], R_co], *args: P.args, **kwargs: P.
 
 class Foo:
     def method(self, x: Self):
-        reveal_type(x)  # revealed: @Todo(Support for `typing.Self`)
+        reveal_type(x)  # revealed: Self
 ```
 
 ## Type expressions

crates/ty_python_semantic/resources/mdtest/call/constructor.md

@@ -75,7 +75,8 @@ constructor from it.
 from typing_extensions import Self
 
 class Base:
-    def __new__(cls, x: int) -> Self: ...
+    def __new__(cls, x: int) -> Self:
+        return cls()
 
 class Foo(Base): ...
 

crates/ty_python_semantic/src/types.rs

@@ -1,3 +1,4 @@
+use infer::enclosing_class_symbol;
 use itertools::Either;
 
 use std::slice::Iter;
@@ -4678,6 +4679,7 @@ impl<'db> Type<'db> {
     pub fn in_type_expression(
         &self,
         db: &'db dyn Db,
+        scope_id: ScopeId,
     ) -> Result<Type<'db>, InvalidTypeExpressionError<'db>> {
         match self {
             // Special cases for `float` and `complex`
@@ -4762,7 +4764,40 @@ impl<'db> Type<'db> {
                 // TODO: Use an opt-in rule for a bare `Callable`
                 KnownInstanceType::Callable => Ok(Type::Callable(CallableType::unknown(db))),
 
-                KnownInstanceType::TypingSelf => Ok(todo_type!("Support for `typing.Self`")),
+                KnownInstanceType::TypingSelf => {
+                    let index = semantic_index(db, scope_id.file(db));
+                    let Some(class_ty) = enclosing_class_symbol(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 Some(instance) = class_ty.to_instance(db) else {
+                        debug_assert!(
+                            false,
+                            "enclosing_class_symbol must return type that can be instantiated"
+                        );
+                        return Ok(Type::unknown());
+                    };
+                    Ok(Type::TypeVar(TypeVarInstance::new(
+                        db,
+                        ast::name::Name::new("Self"),
+                        class_def,
+                        Some(TypeVarBoundOrConstraints::UpperBound(instance)),
+                        TypeVarVariance::Invariant,
+                        None,
+                        TypeVarKind::Legacy,
+                    )))
+                }
                 KnownInstanceType::TypeAlias => Ok(todo_type!("Support for `typing.TypeAlias`")),
                 KnownInstanceType::TypedDict => Ok(todo_type!("Support for `typing.TypedDict`")),
 
@@ -4829,7 +4864,7 @@ impl<'db> Type<'db> {
                 let mut builder = UnionBuilder::new(db);
                 let mut invalid_expressions = smallvec::SmallVec::default();
                 for element in union.elements(db) {
-                    match element.in_type_expression(db) {
+                    match element.in_type_expression(db, scope_id) {
                         Ok(type_expr) => builder = builder.add(type_expr),
                         Err(InvalidTypeExpressionError {
                             fallback_type,

crates/ty_python_semantic/src/types/infer.rs

@@ -316,6 +316,31 @@ 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.
+///
+/// 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>(
+    db: &'db dyn Db,
+    semantic: &SemanticIndex<'db>,
+    scope: ScopeId,
+) -> Option<Type<'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())
+            } else {
+                None
+            }
+        })
+}
+
 /// A region within which we can infer types.
 #[derive(Copy, Clone, Debug)]
 pub(crate) enum InferenceRegion<'db> {
@@ -4582,27 +4607,6 @@ impl<'db> TypeInferenceBuilder<'db> {
         Some(infer_definition_types(self.db(), definition).binding_type(definition))
     }
 
-    /// 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
-    fn enclosing_class_symbol(&self, scope: ScopeId) -> Option<Type<'db>> {
-        self.index
-            .ancestor_scopes(scope.file_scope_id(self.db()))
-            .find_map(|(_, ancestor_scope)| {
-                if let NodeWithScopeKind::Class(class) = ancestor_scope.node() {
-                    let definition = self.index.expect_single_definition(class.node());
-                    let result = infer_definition_types(self.db(), definition);
-
-                    Some(result.declaration_type(definition).inner_type())
-                } else {
-                    None
-                }
-            })
-    }
-
     fn infer_call_expression(
         &mut self,
         call_expression_node: &ast::Expr,
@@ -4911,9 +4915,11 @@ impl<'db> TypeInferenceBuilder<'db> {
                                             [] => {
                                                 let scope = self.scope();
 
-                                                let Some(enclosing_class) =
-                                                    self.enclosing_class_symbol(scope)
-                                                else {
+                                                let Some(enclosing_class) = enclosing_class_symbol(
+                                                    self.db(),
+                                                    self.index,
+                                                    scope,
+                                                ) else {
                                                     overload.set_return_type(Type::unknown());
                                                     BoundSuperError::UnavailableImplicitArguments
                                                         .report_diagnostic(
@@ -7311,7 +7317,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                             TypeAndQualifiers::new(Type::unknown(), TypeQualifiers::FINAL)
                         }
                         _ => name_expr_ty
-                            .in_type_expression(self.db())
+                            .in_type_expression(self.db(), self.scope())
                             .unwrap_or_else(|error| {
                                 error.into_fallback_type(
                                     &self.context,
@@ -7491,7 +7497,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::Name(name) => match name.ctx {
                 ast::ExprContext::Load => self
                     .infer_name_expression(name)
-                    .in_type_expression(self.db())
+                    .in_type_expression(self.db(), self.scope())
                     .unwrap_or_else(|error| {
                         error.into_fallback_type(
                             &self.context,
@@ -7508,7 +7514,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::Attribute(attribute_expression) => match attribute_expression.ctx {
                 ast::ExprContext::Load => self
                     .infer_attribute_expression(attribute_expression)
-                    .in_type_expression(self.db())
+                    .in_type_expression(self.db(), self.scope())
                     .unwrap_or_else(|error| {
                         error.into_fallback_type(
                             &self.context,
@@ -8010,7 +8016,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                             generic_context,
                         );
                         specialized_class
-                            .in_type_expression(self.db())
+                            .in_type_expression(self.db(), self.scope())
                             .unwrap_or(Type::unknown())
                     }
                     None => {

crates/ty_python_semantic/src/types/known_instance.rs

@@ -87,9 +87,11 @@ pub enum KnownInstanceType<'db> {
     /// (which can also be found as `typing_extensions.Callable` or as `collections.abc.Callable`)
     Callable,
 
+    /// The symbol `typing.Self`
+    TypingSelf,
+
     // Various special forms, special aliases and type qualifiers that we don't yet understand
     // (all currently inferred as TODO in most contexts):
-    TypingSelf,
     Final,
     ClassVar,
     Concatenate,