Merge pull request #3525 from aatle/sprite-stubs

Improve sprite stubs

Author: MyreMylar

buildconfig/stubs/pygame/sprite.pyi

@@ -27,9 +27,7 @@ from pygame.rect import FRect, Rect
 from pygame.surface import Surface
 from pygame.typing import Point, RectLike
 
-# define some useful protocols first, which sprite functions accept
-# sprite functions don't need all sprite attributes to be present in the
-# arguments passed, they only use a few which are marked in the below protocols
+# Some sprite functions only need objects with certain attributes, not always a sprite
 class _HasRect(Protocol):
     @property
     def rect(self) -> Optional[Union[FRect, Rect]]: ...
@@ -41,51 +39,10 @@ class _HasImageAndRect(_HasRect, Protocol):
 
 # mask in addition to rect
 class _HasMaskAndRect(_HasRect, Protocol):
-    mask: Mask
-
-# radius in addition to rect
-class _HasRadiusAndRect(_HasRect, Protocol):
-    radius: float
-
-# non-generic Group, used in Sprite
-_Group = AbstractGroup[Any]
-
-# protocol helps with structural subtyping for typevars in sprite group generics
-# and allows the use of any class with the required attributes and methods
-class _SupportsSprite(_HasImageAndRect, Protocol):
-    @property
-    def image(self) -> Optional[Surface]: ...
-    @image.setter
-    def image(self, value: Optional[Surface]) -> None: ...
     @property
-    def rect(self) -> Optional[Union[FRect, Rect]]: ...
-    @rect.setter
-    def rect(self, value: Optional[Union[FRect, Rect]]) -> None: ...
-    @property
-    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 update(self, *args: Any, **kwargs: Any) -> None: ...
-    def add(self, *groups: _Group) -> None: ...
-    def remove(self, *groups: _Group) -> None: ...
-    def kill(self) -> None: ...
-    def alive(self) -> bool: ...
-    def groups(self) -> list[_Group]: ...
-
-# also a protocol
-class _SupportsDirtySprite(_SupportsSprite, Protocol):
-    dirty: int
-    blendmode: int
-    source_rect: Union[FRect, Rect]
-    visible: int
-    _layer: int
-    def _set_visible(self, val: int) -> None: ...
-    def _get_visible(self) -> int: ...
+    def mask(self) -> Mask: ...
 
-# concrete sprite implementation class
-class Sprite(_SupportsSprite):
+class Sprite(_HasImageAndRect):
     @property
     def image(self) -> Optional[Surface]: ...
     @image.setter
@@ -98,52 +55,47 @@ 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: _GroupOrGroups[Any]) -> None: ...
+    def add_internal(self, group: AbstractGroup[Any]) -> None: ...
+    def remove_internal(self, group: AbstractGroup[Any]) -> 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: _GroupOrGroups[Any]) -> None: ...
+    def remove(self, *groups: _GroupOrGroups[Any]) -> None: ...
     def kill(self) -> None: ...
     def alive(self) -> bool: ...
-    def groups(self) -> list[AbstractGroup[_SupportsSprite]]: ...
+    def groups(self) -> list[AbstractGroup[Sprite]]: ...
 
-# concrete dirty sprite implementation class
-class DirtySprite(Sprite, _SupportsDirtySprite):
+class DirtySprite(Sprite):
     dirty: int
     blendmode: int
     source_rect: Union[FRect, Rect]
     visible: int
     _layer: int
-    def _set_visible(self, val: int) -> None: ...
-    def _get_visible(self) -> int: ...
 
-# typevar bound to Sprite, _SupportsSprite Protocol ensures sprite
-# subclass passed to group has image and rect attributes
-_TSprite = TypeVar("_TSprite", bound=_SupportsSprite)
-_TSprite2 = TypeVar("_TSprite2", bound=_SupportsSprite)
-_TDirtySprite = TypeVar("_TDirtySprite", bound=_SupportsDirtySprite)
+_SpriteT = TypeVar("_SpriteT", bound=Sprite)
+_SpriteT2 = TypeVar("_SpriteT2", bound=Sprite)
+_DirtySpriteT = TypeVar("_DirtySpriteT", bound=DirtySprite)
 
-# typevar for sprite or iterable of sprites, used in Group init, add and remove
-_SpriteOrIterable = Union[_TSprite, Iterable[_SpriteOrIterable[_TSprite]]]
+_GroupOrGroups = Union[AbstractGroup[_SpriteT], Iterable[_GroupOrGroups[_SpriteT]]]
+_SpriteOrSprites = Union[_SpriteT, Iterable[_SpriteOrSprites[_SpriteT]]]
 
-class AbstractGroup(Generic[_TSprite]):
-    spritedict: dict[_TSprite, Optional[Union[FRect, Rect]]]
+class AbstractGroup(Generic[_SpriteT]):
+    spritedict: dict[_SpriteT, Optional[Union[FRect, Rect]]]
     lostsprites: list[Union[FRect, Rect]]
     def __class_getitem__(cls, item: Any, /) -> types.GenericAlias: ...
     def __init__(self) -> None: ...
     def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[_TSprite]: ...
+    def __iter__(self) -> Iterator[_SpriteT]: ...
     def __bool__(self) -> bool: ...
     def __contains__(self, item: Any) -> bool: ...
-    def add_internal(self, sprite: _TSprite, layer: None = None) -> None: ...
-    def remove_internal(self, sprite: _TSprite) -> None: ...
-    def has_internal(self, sprite: _TSprite) -> bool: ...
+    def add_internal(self, sprite: _SpriteT, layer: Optional[int] = None) -> None: ...
+    def remove_internal(self, sprite: _SpriteT) -> None: ...
+    def has_internal(self, sprite: _SpriteT) -> bool: ...
     def copy(self) -> Self: ...
-    def sprites(self) -> list[_TSprite]: ...
-    def add(self, *sprites: _SpriteOrIterable[_TSprite]) -> None: ...
-    def remove(self, *sprites: _SpriteOrIterable[_TSprite]) -> None: ...
-    def has(self, *sprites: _SpriteOrIterable[_TSprite]) -> bool: ...
+    def sprites(self) -> list[_SpriteT]: ...
+    def add(self, *sprites: _SpriteOrSprites[_SpriteT]) -> None: ...
+    def remove(self, *sprites: _SpriteOrSprites[_SpriteT]) -> None: ...
+    def has(self, *sprites: _SpriteOrSprites[_SpriteT]) -> bool: ...
     def update(self, *args: Any, **kwargs: Any) -> None: ...
     def draw(
         self, surface: Surface, bgd: Optional[Surface] = None, special_flags: int = 0
@@ -155,41 +107,39 @@ class AbstractGroup(Generic[_TSprite]):
     ) -> None: ...
     def empty(self) -> None: ...
 
-class Group(AbstractGroup[_TSprite]):
-    def __init__(self, *sprites: _SpriteOrIterable[_TSprite]) -> None: ...
+class Group(AbstractGroup[_SpriteT]):
+    def __init__(self, *sprites: _SpriteOrSprites[_SpriteT]) -> None: ...
 
-# these are aliased in the code too
+# These deprecated types are just aliases in the code too
 @deprecated("Use `pygame.sprite.Group` instead")
-class RenderPlain(Group[_TSprite]): ...
+class RenderPlain(Group[_SpriteT]): ...
 
 @deprecated("Use `pygame.sprite.Group` instead")
-class RenderClear(Group[_TSprite]): ...
+class RenderClear(Group[_SpriteT]): ...
 
-class RenderUpdates(Group[_TSprite]): ...
+class RenderUpdates(Group[_SpriteT]): ...
 
 @deprecated("Use `pygame.sprite.RenderUpdates` instead")
-class OrderedUpdates(RenderUpdates[_TSprite]): ...
-
-class LayeredUpdates(AbstractGroup[_TSprite]):
-    def __init__(
-        self, *sprites: _SpriteOrIterable[_TSprite], **kwargs: Any
-    ) -> None: ...
-    def add(self, *sprites: _SpriteOrIterable[_TSprite], **kwargs: Any) -> None: ...
-    def get_sprites_at(self, pos: Point) -> list[_TSprite]: ...
-    def get_sprite(self, idx: int) -> _TSprite: ...
-    def remove_sprites_of_layer(self, layer_nr: int) -> list[_TSprite]: ...
+class OrderedUpdates(RenderUpdates[_SpriteT]): ...
+
+class LayeredUpdates(AbstractGroup[_SpriteT]):
+    def __init__(self, *sprites: _SpriteOrSprites[_SpriteT], **kwargs: Any) -> None: ...
+    def add(self, *sprites: _SpriteOrSprites[_SpriteT], **kwargs: Any) -> None: ...
+    def get_sprites_at(self, pos: Point) -> list[_SpriteT]: ...
+    def get_sprite(self, idx: int) -> _SpriteT: ...
+    def remove_sprites_of_layer(self, layer_nr: int) -> list[_SpriteT]: ...
     def layers(self) -> list[int]: ...
-    def change_layer(self, sprite: _TSprite, new_layer: int) -> None: ...
-    def get_layer_of_sprite(self, sprite: _TSprite) -> int: ...
+    def change_layer(self, sprite: _SpriteT, new_layer: int) -> None: ...
+    def get_layer_of_sprite(self, sprite: _SpriteT) -> int: ...
     def get_top_layer(self) -> int: ...
     def get_bottom_layer(self) -> int: ...
-    def move_to_front(self, sprite: _TSprite) -> None: ...
-    def move_to_back(self, sprite: _TSprite) -> None: ...
-    def get_top_sprite(self) -> _TSprite: ...
-    def get_sprites_from_layer(self, layer: int) -> list[_TSprite]: ...
+    def move_to_front(self, sprite: _SpriteT) -> None: ...
+    def move_to_back(self, sprite: _SpriteT) -> None: ...
+    def get_top_sprite(self) -> _SpriteT: ...
+    def get_sprites_from_layer(self, layer: int) -> list[_SpriteT]: ...
     def switch_layer(self, layer1_nr: int, layer2_nr: int) -> None: ...
 
-class LayeredDirty(LayeredUpdates[_TDirtySprite]):
+class LayeredDirty(LayeredUpdates[_DirtySpriteT]):
     def draw(
         self,
         surface: Surface,
@@ -207,20 +157,19 @@ class LayeredDirty(LayeredUpdates[_TDirtySprite]):
     )
     def set_timing_treshold(self, time_ms: SupportsFloat) -> None: ...
 
-class GroupSingle(AbstractGroup[_TSprite]):
-    sprite: _TSprite
-    def __init__(self, sprite: Optional[_TSprite] = None) -> None: ...
+class GroupSingle(AbstractGroup[_SpriteT]):
+    sprite: Optional[_SpriteT]
+    def __init__(self, sprite: Optional[_SpriteT] = None) -> None: ...
 
-# argument to collide_rect must have rect attribute
 def collide_rect(left: _HasRect, right: _HasRect) -> bool: ...
 
 class collide_rect_ratio:
     ratio: float
     def __init__(self, ratio: float) -> None: ...
     def __call__(self, left: _HasRect, right: _HasRect) -> bool: ...
 
-# must have rect attribute, may optionally have radius attribute
-_SupportsCollideCircle = Union[_HasRect, _HasRadiusAndRect]
+# Must have rect attribute, may optionally have radius attribute
+_SupportsCollideCircle = _HasRect
 
 def collide_circle(
     left: _SupportsCollideCircle, right: _SupportsCollideCircle
@@ -233,32 +182,31 @@ class collide_circle_ratio:
         self, left: _SupportsCollideCircle, right: _SupportsCollideCircle
     ) -> bool: ...
 
-# argument to collide_mask must either have mask or have image attribute, in
+# Arguments to collide_mask must either have mask or have image attribute, in
 # addition to mandatorily having a rect attribute
 _SupportsCollideMask = Union[_HasImageAndRect, _HasMaskAndRect]
 
 def collide_mask(
     left: _SupportsCollideMask, right: _SupportsCollideMask
 ) -> Optional[tuple[int, int]]: ...
 
-# _HasRect typevar for sprite collide functions
-_THasRect = TypeVar("_THasRect", bound=_HasRect)
+_HasRectT = TypeVar("_HasRectT", bound=_HasRect)
 
 def spritecollide(
-    sprite: _THasRect,
-    group: AbstractGroup[_TSprite],
+    sprite: _HasRectT,
+    group: AbstractGroup[_SpriteT],
     dokill: bool,
-    collided: Optional[Callable[[_THasRect, _TSprite], Any]] = None,
-) -> list[_TSprite]: ...
+    collided: Optional[Callable[[_HasRectT, _SpriteT], bool]] = None,
+) -> list[_SpriteT]: ...
 def groupcollide(
-    groupa: AbstractGroup[_TSprite],
-    groupb: AbstractGroup[_TSprite2],
+    groupa: AbstractGroup[_SpriteT],
+    groupb: AbstractGroup[_SpriteT2],
     dokilla: bool,
     dokillb: bool,
-    collided: Optional[Callable[[_TSprite, _TSprite2], Any]] = None,
-) -> dict[_TSprite, list[_TSprite2]]: ...
+    collided: Optional[Callable[[_SpriteT, _SpriteT2], bool]] = None,
+) -> dict[_SpriteT, list[_SpriteT2]]: ...
 def spritecollideany(
-    sprite: _THasRect,
-    group: AbstractGroup[_TSprite],
-    collided: Optional[Callable[[_THasRect, _TSprite], Any]] = None,
-) -> Optional[_TSprite]: ...
+    sprite: _HasRectT,
+    group: AbstractGroup[_SpriteT],
+    collided: Optional[Callable[[_HasRectT, _SpriteT], bool]] = None,
+) -> Optional[_SpriteT]: ...