Improve SpriteList typing (#1977)

* Add missing SpriteType annotations

* Add broader Iterable[Texture] annotation for preload_textures

* Remove extra whitespace after preload_textures

* Add None return annotation to enable_spatial_hashing

* Add None return annotation to SpriteList.sort

* Add None return annotation to SpriteList.shuffle

* Add None return annotation to SpriteList.reverse

* Add None return annotation to SpriteList.insert

* Add None return annotation to SpriteList.extend

* Add None return annotation to SpriteList.remove

* Add None return annotation to SpriteList.swap

* Add None return annotation to SpriteList.append

* Add None return annotation to SpriteList.clear

* Add SpriteType return annotation to SpriteList.__getitem__

* Add None return type to internal _update and _*_buffermethods on SpriteList

* Add None return type to on_update and update_animation

* Add None return annotation to _recalculate_spatial_hashes

* Revert __getitem__ return annotation for the moment

* Add None return type for SpriteList.__setitem__

* Add None return annotation to SpriteList._init_deferred

* Use @overload typing for SpriteList.__getitem__

* Annotate SpriteList's program attribute

* Import Program type and use it to annotate

* Annotate protected buffer props

* Ugly handling for buffer initialization on Optional[Buffer] protected attrs

* Revert "Use @overload typing for SpriteList.__getitem__"

This reverts commit 64c5ae9.

Author: pushfoo

arcade/sprite_list/sprite_list.py

@@ -33,7 +33,7 @@
     get_window,
     gl,
 )
-from arcade.gl import Texture2D
+from arcade.gl import Texture2D, Program
 from arcade.types import Color, RGBA255
 from arcade.gl.types import OpenGlFilter, BlendFunction, PyGLenum
 from arcade.gl.buffer import Buffer
@@ -104,7 +104,7 @@ def __init__(
         lazy: bool = False,
         visible: bool = True,
     ):
-        self.program = None
+        self.program: Optional[Program] = None
         self._atlas: Optional[TextureAtlas] = atlas
         self._initialized = False
         self._lazy = lazy
@@ -137,22 +137,25 @@ def __init__(
         # Index buffer
         self._sprite_index_data = array("i", [0] * self._idx_capacity)
 
-        self._sprite_pos_buf = None
-        self._sprite_size_buf = None
-        self._sprite_angle_buf = None
-        self._sprite_color_buf = None
-        self._sprite_texture_buf = None
+        # Define and annotate storage space for buffers
+        self._sprite_pos_buf: Optional[Buffer] = None
+        self._sprite_size_buf: Optional[Buffer] = None
+        self._sprite_angle_buf: Optional[Buffer] = None
+        self._sprite_color_buf: Optional[Buffer] = None
+        self._sprite_texture_buf: Optional[Buffer] = None
+
         # Index buffer
-        self._sprite_index_buf = None
+        self._sprite_index_buf: Optional[Buffer] = None
+
+        self._geometry: Optional[Geometry] = None
 
-        self._geometry = None
         # Flags for signaling if a buffer needs to be written to the opengl buffer
-        self._sprite_pos_changed = False
-        self._sprite_size_changed = False
-        self._sprite_angle_changed = False
-        self._sprite_color_changed = False
-        self._sprite_texture_changed = False
-        self._sprite_index_changed = False
+        self._sprite_pos_changed: bool = False
+        self._sprite_size_changed: bool = False
+        self._sprite_angle_changed: bool = False
+        self._sprite_color_changed: bool = False
+        self._sprite_texture_changed: bool = False
+        self._sprite_index_changed: bool = False
 
         # Used in collision detection optimization
         from .spatial_hash import SpatialHash
@@ -179,7 +182,7 @@ def __init__(
         except RuntimeError:
             pass
 
-    def _init_deferred(self):
+    def _init_deferred(self) -> None:
         """
         Since spritelist can be created before the window we need to defer initialization.
         It also makes us able to support lazy loading.
@@ -252,7 +255,7 @@ def __iter__(self) -> Iterator[SpriteType]:
     def __getitem__(self, i):
         return self.sprite_list[i]
 
-    def __setitem__(self, index: int, sprite: SpriteType):
+    def __setitem__(self, index: int, sprite: SpriteType) -> None:
         """Replace a sprite at a specific index"""
         # print(f"{id(self)} : {id(sprite)} __setitem__({index})")
 
@@ -525,7 +528,7 @@ def index(self, sprite: SpriteType) -> int:
         """
         return self.sprite_list.index(sprite)
 
-    def clear(self, deep: bool = True):
+    def clear(self, deep: bool = True) -> None:
         """
         Remove all the sprites resetting the spritelist
         to it's initial state.
@@ -588,7 +591,7 @@ def pop(self, index: int = -1) -> SpriteType:
         self.remove(sprite)
         return sprite
 
-    def append(self, sprite: SpriteType):
+    def append(self, sprite: SpriteType) -> None:
         """
         Add a new sprite to the list.
 
@@ -624,7 +627,7 @@ def append(self, sprite: SpriteType):
                 raise ValueError("Sprite must have a texture when added to a SpriteList")
             self._atlas.add(sprite.texture)  # type: ignore
 
-    def swap(self, index_1: int, index_2: int):
+    def swap(self, index_1: int, index_2: int) -> None:
         """
         Swap two sprites by index
         :param index_1: Item index to swap
@@ -644,7 +647,7 @@ def swap(self, index_1: int, index_2: int):
         self._sprite_index_data[i1] = slot_2
         self._sprite_index_data[i2] = slot_1
 
-    def remove(self, sprite: SpriteType):
+    def remove(self, sprite: SpriteType) -> None:
         """
         Remove a specific sprite from the list.
         :param sprite: Item to remove from the list
@@ -676,7 +679,7 @@ def remove(self, sprite: SpriteType):
         if self.spatial_hash is not None:
             self.spatial_hash.remove(sprite)
 
-    def extend(self, sprites: Union[Iterable[SpriteType], "SpriteList"]):
+    def extend(self, sprites: Union[Iterable[SpriteType], SpriteList[SpriteType]]) -> None:
         """
         Extends the current list with the given iterable
 
@@ -685,7 +688,7 @@ def extend(self, sprites: Union[Iterable[SpriteType], "SpriteList"]):
         for sprite in sprites:
             self.append(sprite)
 
-    def insert(self, index: int, sprite: SpriteType):
+    def insert(self, index: int, sprite: SpriteType) -> None:
         """
         Inserts a sprite at a given index.
 
@@ -716,7 +719,7 @@ def insert(self, index: int, sprite: SpriteType):
         if self.spatial_hash is not None:
             self.spatial_hash.add(sprite)
 
-    def reverse(self):
+    def reverse(self) -> None:
         """
         Reverses the current list in-place
         """
@@ -732,7 +735,7 @@ def reverse(self):
 
         self._sprite_index_changed = True
 
-    def shuffle(self):
+    def shuffle(self) -> None:
         """
         Shuffles the current list in-place
         """
@@ -759,7 +762,7 @@ def shuffle(self):
 
         self._sprite_index_changed = True
 
-    def sort(self, *, key: Callable, reverse: bool = False):
+    def sort(self, *, key: Callable, reverse: bool = False) -> None:
         """
         Sort the spritelist in place using ``<`` comparison between sprites.
         This function is similar to python's :py:meth:`list.sort`.
@@ -799,7 +802,7 @@ def disable_spatial_hashing(self) -> None:
         """
         self.spatial_hash = None
 
-    def enable_spatial_hashing(self, spatial_hash_cell_size: int = 128):
+    def enable_spatial_hashing(self, spatial_hash_cell_size: int = 128) -> None:
         """Turn on spatial hashing."""
         if self.spatial_hash is None or self.spatial_hash.cell_size != spatial_hash_cell_size:
             LOG.debug("Enabled spatial hashing with cell size %s", spatial_hash_cell_size)
@@ -809,7 +812,7 @@ def enable_spatial_hashing(self, spatial_hash_cell_size: int = 128):
         else:
             LOG.debug("Spatial hashing is already enabled with size %s", spatial_hash_cell_size)
 
-    def _recalculate_spatial_hashes(self):
+    def _recalculate_spatial_hashes(self) -> None:
         if self.spatial_hash is None:
             from .spatial_hash import SpatialHash
             self.spatial_hash = SpatialHash(cell_size=self._spatial_hash_cell_size)
@@ -825,14 +828,14 @@ def update(self) -> None:
         for sprite in self.sprite_list:
             sprite.update()
 
-    def on_update(self, delta_time: float = 1 / 60):
+    def on_update(self, delta_time: float = 1 / 60) -> None:
         """
         Update the sprite. Similar to update, but also takes a delta-time.
         """
         for sprite in self.sprite_list:
             sprite.on_update(delta_time)
 
-    def update_animation(self, delta_time: float = 1 / 60):
+    def update_animation(self, delta_time: float = 1 / 60) -> None:
         """
         Call the update_animation in every sprite in the sprite list.
         """
@@ -870,7 +873,7 @@ def move(self, change_x: float, change_y: float) -> None:
             sprite.center_x += change_x
             sprite.center_y += change_y
 
-    def preload_textures(self, texture_list: List["Texture"]) -> None:
+    def preload_textures(self, texture_list: Iterable["Texture"]) -> None:
         """
         Preload a set of textures that will be used for sprites in this
         sprite list.
@@ -885,7 +888,6 @@ def preload_textures(self, texture_list: List["Texture"]) -> None:
             self._atlas.add( # type: ignore
                 texture)
 
-
     def write_sprite_buffers_to_gpu(self) -> None:
         """
         Ensure buffers are resized and fresh sprite data
@@ -903,7 +905,7 @@ def write_sprite_buffers_to_gpu(self) -> None:
         """
         self._write_sprite_buffers_to_gpu()
 
-    def _write_sprite_buffers_to_gpu(self):
+    def _write_sprite_buffers_to_gpu(self) -> None:
         LOG.debug(
             "[%s] SpriteList._write_sprite_buffers_to_gpu: pos=%s, size=%s, angle=%s, color=%s tex=%s idx=%s",
             id(self),
@@ -1043,7 +1045,7 @@ def draw_hit_boxes(self, color: RGBA255 = (0, 0, 0, 255), line_thickness: float
         for sprite in self.sprite_list:
             sprite.draw_hit_box(color, line_thickness)
 
-    def _normalize_index_buffer(self):
+    def _normalize_index_buffer(self) -> None:
         """
         Removes unused slots in the index buffer.
         The other buffers don't need this because they re-use slots.
@@ -1062,7 +1064,7 @@ def _normalize_index_buffer(self):
         # NOTE: Right now the index buffer is always normalized
         pass
 
-    def _grow_sprite_buffers(self):
+    def _grow_sprite_buffers(self) -> None:
         """Double the internal buffer sizes"""
         # Resize sprite buffers if needed
         if self._sprite_buffer_slots <= self._buf_capacity:
@@ -1087,19 +1089,20 @@ def _grow_sprite_buffers(self):
         self._sprite_texture_data.extend([0] * extend_by)
 
         if self._initialized:
-            self._sprite_pos_buf.orphan(double=True)
-            self._sprite_size_buf.orphan(double=True)
-            self._sprite_angle_buf.orphan(double=True)
-            self._sprite_color_buf.orphan(double=True)
-            self._sprite_texture_buf.orphan(double=True)
+            # Proper initialization implies these buffers are allocated
+            self._sprite_pos_buf.orphan(double=True) # type: ignore
+            self._sprite_size_buf.orphan(double=True) # type: ignore
+            self._sprite_angle_buf.orphan(double=True) # type: ignore
+            self._sprite_color_buf.orphan(double=True) # type: ignore
+            self._sprite_texture_buf.orphan(double=True) # type: ignore
 
         self._sprite_pos_changed = True
         self._sprite_size_changed = True
         self._sprite_angle_changed = True
         self._sprite_color_changed = True
         self._sprite_texture_changed = True
 
-    def _grow_index_buffer(self):
+    def _grow_index_buffer(self) -> None:
         # Extend the index buffer capacity if needed
         if self._sprite_index_slots <= self._idx_capacity:
             return
@@ -1127,7 +1130,7 @@ def _grow_index_buffer(self):
 
         self._sprite_index_changed = True
 
-    def _update_all(self, sprite: SpriteType):
+    def _update_all(self, sprite: SpriteType) -> None:
         """
         Update all sprite data. This is faster when adding and moving sprites.
         This duplicate code, but reduces call overhead, dict lookups etc.
@@ -1168,7 +1171,7 @@ def _update_all(self, sprite: SpriteType):
         self._sprite_texture_data[slot] = tex_slot
         self._sprite_texture_changed = True
 
-    def _update_texture(self, sprite) -> None:
+    def _update_texture(self, sprite: SpriteType) -> None:
         """Make sure we update the texture for this sprite for the next batch
         drawing"""
         # We cannot interact with texture atlases unless the context
@@ -1274,7 +1277,7 @@ def _update_size(self, sprite: SpriteType) -> None:
         self._sprite_size_data[slot * 2 + 1] = sprite._height
         self._sprite_size_changed = True
 
-    def _update_width(self, sprite: SpriteType):
+    def _update_width(self, sprite: SpriteType) -> None:
         """
         Called by the Sprite class to update the size/scale in this sprite.
         Necessary for batch drawing of items.
@@ -1285,7 +1288,7 @@ def _update_width(self, sprite: SpriteType):
         self._sprite_size_data[slot * 2] = sprite._width
         self._sprite_size_changed = True
 
-    def _update_height(self, sprite: SpriteType):
+    def _update_height(self, sprite: SpriteType) -> None:
         """
         Called by the Sprite class to update the size/scale in this sprite.
         Necessary for batch drawing of items.
@@ -1296,7 +1299,7 @@ def _update_height(self, sprite: SpriteType):
         self._sprite_size_data[slot * 2 + 1] = sprite._height
         self._sprite_size_changed = True
 
-    def _update_angle(self, sprite: SpriteType):
+    def _update_angle(self, sprite: SpriteType) -> None:
         """
         Called by the Sprite class to update the angle in this sprite.
         Necessary for batch drawing of items.