GameObjectCollection.add: Behave correctly with inheritance

ppb/scenes.py

@@ -48,7 +48,9 @@ def add(self, game_object: Hashable, tags: Iterable[Hashable]=()) -> None:
         if isinstance(tags, (str, bytes)):
             raise TypeError("You passed a string instead of an iterable, this probably isn't what you intended.\n\nTry making it a tuple.")
         self.all.add(game_object)
-        self.kinds[type(game_object)].add(game_object)
+
+        for kind in type(game_object).mro():
+            self.kinds[kind].add(game_object)
         for tag in tags:
             self.tags[tag].add(game_object)
 
@@ -91,7 +93,8 @@ def remove(self, game_object: Hashable) -> None:
             container.remove(myObject)
         """
         self.all.remove(game_object)
-        self.kinds[type(game_object)].remove(game_object)
+        for kind in type(game_object).mro():
+            self.kinds[kind].remove(game_object)
         for s in self.tags.values():
             s.discard(game_object)
 
@@ -120,6 +123,14 @@ def __contains__(self, item: Hashable) -> bool:
     def __iter__(self) -> Iterator:
         return (x for x in self.game_objects)
 
+    @property
+    def kinds(self):
+        return self.game_objects.kinds
+
+    @property
+    def tags(self):
+        return self.game_objects.tags
+
     @property
     def main_camera(self) -> Camera:
         return next(self.game_objects.get(tag="main_camera"))

tests/test_scenes.py

@@ -17,6 +17,10 @@ class TestPlayer:
     pass
 
 
+class TestSubclassPlayer(TestPlayer):
+    pass
+
+
 class TestSprite:
     pass
 
@@ -26,9 +30,15 @@ def containers():
     yield BaseScene(Mock())
 
 
-@fixture()
-def player():
-    return TestPlayer()
+def players():
+    yield TestPlayer()
+    yield TestSubclassPlayer()
+
+
+def players_and_containers():
+    for player in players():
+        for container in containers():
+            yield player, container
 
 
 @fixture()
@@ -42,7 +52,7 @@ def scene():
     return BaseScene(engine)
 
 
-@mark.parametrize("container", containers())
+@mark.parametrize("player, container", players_and_containers())
 def test_add_methods(container, player, enemies):
     container.add(player)
     for group, sprite in zip(("red", "blue"), enemies):
@@ -52,7 +62,7 @@ def test_add_methods(container, player, enemies):
         assert enemy in container
 
 
-@mark.parametrize("container", containers())
+@mark.parametrize("player, container", players_and_containers())
 def test_get_methods(container, player, enemies):
 
     sprite = TestSprite()
@@ -61,43 +71,26 @@ def test_get_methods(container, player, enemies):
     container.add(enemies[1], ["red"])
     container.add(sprite)
 
-    enemy_set = set(container.get(kind=TestEnemy))
-    assert len(enemy_set) == 2
-    for enemy in enemies:
-        assert enemy in enemy_set
-
-    player_set = set(container.get(kind=TestPlayer))
-    assert len(player_set) == 1
-    assert player in player_set
+    assert set(container.get(kind=TestEnemy)) == set(enemies)
+    assert set(container.get(kind=TestPlayer)) == {player}
+    assert set(container.get(kind=TestSprite)) == {sprite}
 
-    sprite_set = set(container.get(kind=TestSprite))
-    assert len(sprite_set) == 1
-    assert sprite in sprite_set
+    assert set(container.get(tag="red")) == {player, enemies[1]}
 
-    red_set = set(container.get(tag="red"))
-    assert len(red_set) == 2
-    assert player in red_set
-    assert enemies[1] in red_set
-    assert enemies[0] not in red_set
-
-    null_set = set(container.get(tag="this doesn't exist"))
-    assert len(null_set) == 0
-    assert player not in null_set
-    assert enemies[0] not in null_set
-    assert enemies[1] not in null_set
+    assert set(container.get(tag="this doesn't exist")) == set()
 
     with raises(TypeError):
         container.get()
 
 
-@mark.parametrize("container", containers())
+@mark.parametrize("player, container", players_and_containers())
 def test_get_with_string_tags(container, player):
     """Test that addings a string instead of an array-like throws."""
     with raises(TypeError):
         container.add(player, "player")
 
 
-@mark.parametrize("container", containers())
+@mark.parametrize("player, container", players_and_containers())
 def test_remove_methods(container, player, enemies):
     container.add(player, ["test"])
     container.add(enemies[0], ["test"])
@@ -109,15 +102,19 @@ def test_remove_methods(container, player, enemies):
     container.remove(player)
 
     assert player not in container
-    assert player not in container.get(kind=TestPlayer)
-    assert player not in container.get(tag="test")
+    for kind in container.kinds:
+        assert player not in container.get(kind=kind)
+    for tag in container.tags:
+        assert player not in container.get(tag=tag)
+
     assert enemies[0] in container
     assert enemies[0] in container.get(tag="test")
     assert enemies[1] in container
 
 
-@mark.parametrize("container", [GameObjectCollection()])
-def test_collection_methods(container, player, enemies):
+@mark.parametrize("player", players())
+def test_collection_methods(player, enemies):
+    container = GameObjectCollection()
     container.add(player)
     container.add(enemies[0])