Store improvements

- Store games by source
- Added convenience dunder methods (contains, iter, len, getitem)
- Added store.get
- Changed `enable_in_pipeline` for `toggle_in_pipeline`

src/store/store.py

@@ -18,6 +18,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import logging
+from typing import MutableMapping, Generator, Any
 
 from src import shared
 from src.game import Game
@@ -31,24 +32,59 @@ class Store:
     managers: dict[type[Manager], Manager]
     pipeline_managers: set[Manager]
     pipelines: dict[str, Pipeline]
-    games: dict[str, Game]
+    source_games: MutableMapping[str, MutableMapping[str, Game]]
 
     def __init__(self) -> None:
         self.managers = {}
         self.pipeline_managers = set()
         self.pipelines = {}
-        self.games = {}
+        self.source_games = {}
+
+    def __contains__(self, obj: object) -> bool:
+        """Check if the game is present in the store with the `in` keyword"""
+        if not isinstance(obj, Game):
+            return False
+        if not (source_mapping := self.source_games.get(obj.source)):
+            return False
+        return obj.game_id in source_mapping
+
+    def __iter__(self) -> Generator[Game, None, None]:
+        """Iterate through the games in the store with `for ... in`"""
+        for _source_id, games_mapping in self.source_games.items():
+            for _game_id, game in games_mapping.items():
+                yield game
+
+    def __len__(self) -> int:
+        """Get the number of games in the store with the `len` builtin"""
+        return sum(len(source_mapping) for source_mapping in self.source_games)
+
+    def __getitem__(self, game_id: str) -> Game:
+        """Get a game by its id with `store["game_id_goes_here"]`"""
+        for game in iter(self):
+            if game.game_id == game_id:
+                return game
+        raise KeyError("Game not found in store")
+
+    def get(self, game_id: str, default: Any = None) -> Game | Any:
+        """Get a game by its ID, with a fallback if not found"""
+        try:
+            game = self[game_id]
+            return game
+        except KeyError:
+            return default
 
     def add_manager(self, manager: Manager, in_pipeline=True):
         """Add a manager to the store"""
         manager_type = type(manager)
         self.managers[manager_type] = manager
-        if in_pipeline:
-            self.enable_manager_in_pipelines(manager_type)
+        self.toggle_manager_in_pipelines(manager_type, in_pipeline)
 
-    def enable_manager_in_pipelines(self, manager_type: type[Manager]):
-        """Make a manager run in new pipelines"""
-        self.pipeline_managers.add(self.managers[manager_type])
+    def toggle_manager_in_pipelines(self, manager_type: type[Manager], enable: bool):
+        """Change if a manager should run in new pipelines"""
+        if enable:
+            self.pipeline_managers.add(self.managers[manager_type])
+        else:
+            self.pipeline_managers.discard(self.managers[manager_type])
 
     def cleanup_game(self, game: Game) -> None:
         """Remove a game's files"""
@@ -74,7 +110,7 @@ class Store:
             return None
 
         # Handle game duplicates
-        stored_game = self.games.get(game.game_id)
+        stored_game = self.get(game.game_id)
         if not stored_game:
             # New game, do as normal
             logging.debug("New store game %s (%s)", game.name, game.game_id)
@@ -96,11 +132,15 @@ class Store:
             for signal in manager.signals:
                 game.connect(signal, manager.execute_resilient_manager_logic)
 
+        # Add the game to the store
+        if not game.source in self.source_games:
+            self.source_games[game.source] = {}
+        self.source_games[game.source][game.game_id] = game
+
         # Run the pipeline for the game
         if not run_pipeline:
             return None
         pipeline = Pipeline(game, additional_data, self.pipeline_managers)
-        self.games[game.game_id] = game
         self.pipelines[game.game_id] = pipeline
         pipeline.advance()
         return pipeline