Add support for typing.Self in classmethods and generic classes.

For #1283.

PiperOrigin-RevId: 574971273

pytype/abstract/_function_base.py

@@ -210,18 +210,24 @@ def __init__(self, callself, underlying):
     inst = abstract_utils.get_atomic_value(
         self._callself, default=self.ctx.convert.unsolvable)
     if self.underlying.should_set_self_annot():
-      self._self_annot = abstract_utils.get_generic_type(inst)
-      # TODO(b/224600845): Support typing.Self in generic classes.
-      if not self._self_annot and self._has_self_type_param():
-        # TODO(b/224600845): Support classmethods.
-        self._self_annot = self.ctx.convert.lookup_value("typing", "Self")
+      self._self_annot = self._get_self_annot(inst)
     if isinstance(inst, _instance_base.SimpleValue):
       self.alias_map = inst.instance_type_parameters.aliases
     elif isinstance(inst, _typing.TypeParameterInstance):
       self.alias_map = inst.instance.instance_type_parameters.aliases
     else:
       self.alias_map = None
 
+  def _get_self_annot(self, callself):
+    if not self._has_self_type_param():
+      return abstract_utils.get_generic_type(callself)
+    self_type = self.ctx.convert.lookup_value("typing", "Self")
+    if "classmethod" in self.underlying.decorators:
+      return _classes.ParameterizedClass(
+          self.ctx.convert.type_type, {abstract_utils.T: self_type}, self.ctx)
+    else:
+      return self_type
+
   def _has_self_type_param(self):
     if not isinstance(self.underlying, SignedFunction):
       return False

pytype/tests/test_typing_self.py

@@ -72,6 +72,45 @@ class B(A):
       assert_type(B().f(), B)
     """)
 
+  def test_classmethod(self):
+    self.CheckWithErrors("""
+      from typing_extensions import Self  # not-supported-yet
+      class A:
+        @classmethod
+        def build(cls) -> Self:
+          return cls()
+      class B(A):
+        pass
+      assert_type(A.build(), A)
+      assert_type(B.build(), B)
+    """)
+
+  def test_new(self):
+    self.CheckWithErrors("""
+      from typing_extensions import Self  # not-supported-yet
+      class A:
+        def __new__(cls) -> Self:
+          return super().__new__(cls)
+      class B(A):
+        pass
+      assert_type(A(), A)
+      assert_type(B(), B)
+    """)
+
+  def test_generic_class(self):
+    self.CheckWithErrors("""
+      from typing import Generic, TypeVar
+      from typing_extensions import Self  # not-supported-yet
+      T = TypeVar('T')
+      class A(Generic[T]):
+        def copy(self) -> Self:
+          return self
+      class B(A[T]):
+        pass
+      assert_type(A[int]().copy(), A[int])
+      assert_type(B[str]().copy(), B[str])
+    """)
+
 
 class SelfPyiTest(test_base.BaseTest):
   """Tests for typing.Self usage in type stubs."""
@@ -164,6 +203,23 @@ class C(foo.A.B):
         assert_type(C().f(), C)
       """)
 
+  def test_generic_class(self):
+    with self.DepTree([("foo.pyi", """
+      from typing import Generic, Self, TypeVar
+      T = TypeVar('T')
+      class A(Generic[T]):
+        def copy(self) -> Self: ...
+    """)]):
+      self.Check("""
+        import foo
+        from typing import TypeVar
+        T = TypeVar('T')
+        class B(foo.A[T]):
+          pass
+        assert_type(foo.A[int]().copy(), foo.A[int])
+        assert_type(B[str]().copy(), B[str])
+      """)
+
 
 class SelfReingestTest(test_base.BaseTest):
   """Tests for outputting typing.Self to a stub and reading the stub back in."""
@@ -232,6 +288,72 @@ class C(foo.A.B):
         assert_type(C().f(), C)
       """)
 
+  @test_utils.skipBeforePy((3, 11), "typing.Self is new in 3.11")
+  def test_import_from_typing(self):
+    with self.DepTree([("foo.py", """
+      from typing import Self  # pytype: disable=not-supported-yet
+      class A:
+        def f(self) -> Self:
+          return self
+    """)]):
+      self.Check("""
+        import foo
+        class B(foo.A):
+          pass
+        assert_type(foo.A().f(), foo.A)
+        assert_type(B().f(), B)
+      """)
+
+  def test_classmethod(self):
+    with self.DepTree([("foo.py", """
+      from typing_extensions import Self  # pytype: disable=not-supported-yet
+      class A:
+        @classmethod
+        def build(cls) -> Self:
+          return cls()
+    """)]):
+      self.Check("""
+        import foo
+        class B(foo.A):
+          pass
+        assert_type(foo.A.build(), foo.A)
+        assert_type(B.build(), B)
+      """)
+
+  def test_new(self):
+    with self.DepTree([("foo.py", """
+      from typing_extensions import Self  # pytype: disable=not-supported-yet
+      class A:
+        def __new__(cls) -> Self:
+          return super().__new__(cls)
+    """)]):
+      self.Check("""
+        import foo
+        class B(foo.A):
+          pass
+        assert_type(foo.A(), foo.A)
+        assert_type(B(), B)
+      """)
+
+  def test_generic_class(self):
+    with self.DepTree([("foo.py", """
+      from typing import Generic, TypeVar
+      from typing_extensions import Self  # pytype: disable=not-supported-yet
+      T = TypeVar('T')
+      class A(Generic[T]):
+        def copy(self) -> Self:
+          return self
+    """)]):
+      self.Check("""
+        import foo
+        from typing import TypeVar
+        T = TypeVar('T')
+        class B(foo.A[T]):
+          pass
+        assert_type(foo.A[int]().copy(), foo.A[int])
+        assert_type(B[str]().copy(), B[str])
+      """)
+
 
 if __name__ == "__main__":
   test_base.main()