sprites are now generic to the groups they are contained in

Author: nachomaiz

buildconfig/stubs/pygame/sprite.pyi

@@ -26,8 +26,8 @@ from pygame.rect import FRect, Rect
 from pygame.surface import Surface
 from pygame.typing import Point, RectLike
 
-# non-generic Group, used in Sprite
-_Group = AbstractGroup[_SupportsSprite]
+# generic Group, used in Sprite
+_TGroup = TypeVar("_TGroup", bound=AbstractGroup)
 
 # define some useful protocols first, which sprite functions accept
 # sprite functions don't need all sprite attributes to be present in the
@@ -50,7 +50,8 @@ class _HasRadiusAndRect(_HasRect, Protocol):
     radius: float
 
 # protocol helps with structural subtyping for typevars in sprite group generics
-class _SupportsSprite(_HasImageAndRect, Protocol):
+class _SupportsSprite(Generic[_TGroup], _HasImageAndRect, Protocol):
+    def __class_getitem__(cls, item: type[AbstractGroup], /) -> types.GenericAlias: ...
     @property
     def image(self) -> Optional[Surface]: ...
     @image.setter
@@ -63,17 +64,17 @@ class _SupportsSprite(_HasImageAndRect, Protocol):
     def layer(self) -> int: ...
     @layer.setter
     def layer(self, value: int) -> None: ...
-    def add_internal(self, group: _Group) -> None: ...
-    def remove_internal(self, group: _Group) -> None: ...
+    def add_internal(self, group: _TGroup) -> None: ...
+    def remove_internal(self, group: _TGroup) -> None: ...
     def update(self, *args: Any, **kwargs: Any) -> None: ...
-    def add(self, *groups: _Group) -> None: ...
-    def remove(self, *groups: _Group) -> None: ...
+    def add(self, *groups: _TGroup) -> None: ...
+    def remove(self, *groups: _TGroup) -> None: ...
     def kill(self) -> None: ...
     def alive(self) -> bool: ...
-    def groups(self) -> list[_Group]: ...
+    def groups(self) -> list[_TGroup]: ...
 
 # also a protocol
-class _SupportsDirtySprite(_SupportsSprite, Protocol):
+class _SupportsDirtySprite(_SupportsSprite[_TGroup], Protocol):
     dirty: int
     blendmode: int
     source_rect: Union[FRect, Rect]
@@ -83,7 +84,7 @@ class _SupportsDirtySprite(_SupportsSprite, Protocol):
     def _get_visible(self) -> int: ...
 
 # concrete sprite implementation class
-class Sprite(_SupportsSprite):
+class Sprite(_SupportsSprite[_TGroup]):
     @property
     def image(self) -> Optional[Surface]: ...
     @image.setter
@@ -96,18 +97,18 @@ class Sprite(_SupportsSprite):
     def layer(self) -> int: ...
     @layer.setter
     def layer(self, value: int) -> None: ...
-    def __init__(self, *groups: _Group) -> None: ...
-    def add_internal(self, group: _Group) -> None: ...
-    def remove_internal(self, group: _Group) -> None: ...
+    def __init__(self, *groups: _TGroup) -> None: ...
+    def add_internal(self, group: _TGroup) -> None: ...
+    def remove_internal(self, group: _TGroup) -> None: ...
     def update(self, *args: Any, **kwargs: Any) -> None: ...
-    def add(self, *groups: _Group) -> None: ...
-    def remove(self, *groups: _Group) -> None: ...
+    def add(self, *groups: _TGroup) -> None: ...
+    def remove(self, *groups: _TGroup) -> None: ...
     def kill(self) -> None: ...
     def alive(self) -> bool: ...
-    def groups(self) -> list[_Group]: ...
+    def groups(self) -> list[_TGroup]: ...
 
 # concrete dirty sprite implementation class
-class DirtySprite(_SupportsDirtySprite):
+class DirtySprite(Sprite[_TGroup], _SupportsDirtySprite[_TGroup]):
     dirty: int
     blendmode: int
     source_rect: Union[FRect, Rect]

src_py/sprite.py

@@ -109,6 +109,8 @@ class Sprite:
 
     """
 
+    __class_getitem__ = classmethod(types.GenericAlias)
+
     def __init__(self, *groups):
         self.__g = {}  # The groups the sprite is in
         self.__image: Optional[pygame.surface.Surface] = None