Merge pull request #3053 from Matiiss/matiiss-allow-sprite-group-subscripts

Add runtime support for `pygame.sprite.AbstractGroup` subscripts

Author: MyreMylar

buildconfig/stubs/pygame/sprite.pyi

@@ -1,3 +1,4 @@
+import types
 from collections.abc import Callable, Iterable, Iterator
 from typing import (
     Any,
@@ -144,6 +145,7 @@ _TDirtySprite = TypeVar("_TDirtySprite", bound=_DirtySpriteSupportsGroup)
 class AbstractGroup(Generic[_TSprite]):
     spritedict: dict[_TSprite, Optional[Union[FRect, Rect]]]
     lostsprites: list[Union[FRect, Rect]]
+    def __class_getitem__(cls, generic: Any) -> types.GenericAlias: ...
     def __init__(self) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[_TSprite]: ...

src_py/sprite.py

@@ -84,6 +84,7 @@
 # specific ones that aren't quite so general but fit into common
 # specialized cases.
 
+import types
 from warnings import warn
 from typing import Optional
 
@@ -371,6 +372,9 @@ class AbstractGroup:
 
     """
 
+    def __class_getitem__(cls, generic):
+        return types.GenericAlias(cls, generic)
+
     # protected identifier value to identify sprite groups, and avoid infinite recursion
     _spritegroup = True
 

test/sprite_test.py

@@ -1,6 +1,8 @@
 #################################### IMPORTS ###################################
 
 
+import types
+import typing
 import unittest
 
 import pygame
@@ -660,6 +662,16 @@ def update(self, *args, **kwargs):
         self.assertEqual(test_sprite.sink, [1, 2, 3])
         self.assertEqual(test_sprite.sink_kwargs, {"foo": 4, "bar": 5})
 
+    def test_type_subscript(self):
+        try:
+            group_generic_alias = sprite.Group[sprite.Sprite]
+        except TypeError as e:
+            self.fail(e)
+
+        self.assertIsInstance(group_generic_alias, types.GenericAlias)
+        self.assertIs(typing.get_origin(group_generic_alias), sprite.Group)
+        self.assertEqual(typing.get_args(group_generic_alias), (sprite.Sprite,))
+
 
 ################################################################################