🎨 Simplified SourceIterator-s

- Using generator functions
- Common generator init and next in base class
- Explicited that error handling should happen in generator

src/importer/sources/bottles_source.py

@@ -17,11 +17,6 @@ from src.utils.save_cover import resize_cover, save_cover
 
 class BottlesSourceIterator(SourceIterator):
     source: "BottlesSource"
-    generator: Generator = None
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.generator = self.generator_builder()
 
     def generator_builder(self) -> Optional[Game]:
         """Generator method producing games"""
@@ -55,13 +50,6 @@ class BottlesSourceIterator(SourceIterator):
             # Produce game
             yield game
 
-    def __next__(self) -> Optional[Game]:
-        try:
-            game = next(self.generator)
-        except StopIteration:
-            raise
-        return game
-
 
 class BottlesSource(Source):
     """Generic Bottles source"""

src/importer/sources/heroic_source.py

@@ -38,7 +38,7 @@ class HeroicSubSource(TypedDict):
 
 class HeroicSourceIterator(SourceIterator):
     source: "HeroicSource"
-    generator: Generator = None
+
     sub_sources: dict[str, HeroicSubSource] = {
         "sideload": {
             "service": "sideload",
@@ -89,9 +89,10 @@ class HeroicSourceIterator(SourceIterator):
 
         return Game(values, allow_side_effects=False)
 
-    def sub_sources_generator(self):
+    def generator_builder(self):
         """Generator method producing games from all the Heroic sub-sources"""
-        for _key, sub_source in self.sub_sources.items():
+
+        for sub_source in self.sub_sources.values():
             # Skip disabled sub-sources
             if not shared.schema.get_boolean("heroic-import-" + sub_source["service"]):
                 continue
@@ -112,17 +113,6 @@ class HeroicSourceIterator(SourceIterator):
                     continue
                 yield game
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.generator = self.sub_sources_generator()
-
-    def __next__(self) -> Optional[Game]:
-        try:
-            game = next(self.generator)
-        except StopIteration:
-            raise
-        return game
-
 
 class HeroicSource(Source):
     """Generic heroic games launcher source"""

src/importer/sources/lutris_source.py

@@ -1,6 +1,6 @@
 from sqlite3 import connect
 from time import time
-from typing import Optional
+from typing import Optional, Generator
 
 from src import shared
 from src.game import Game
@@ -11,11 +11,12 @@ from src.utils.save_cover import resize_cover, save_cover
 
 class LutrisSourceIterator(SourceIterator):
     source: "LutrisSource"
-    import_steam = False
+
-    db_connection = None
+    def generator_builder(self) -> Optional[Game]:
-    db_cursor = None
+        """Generator method producing games"""
-    db_location = None
+
-    db_games_request = """
+        # Query the database
+        request = """
             SELECT id, name, slug, runner, hidden
             FROM 'games'
             WHERE
@@ -26,26 +27,12 @@ class LutrisSourceIterator(SourceIterator):
                 AND (runner IS NOT "steam" OR :import_steam)
             ;
         """
-    db_request_params = None
+        params = {"import_steam": shared.schema.get_boolean("lutris-import-steam")}
-
+        connection = connect(self.source.location / "pga.db")
-    def __init__(self, *args, **kwargs) -> None:
+        cursor = connection.execute(request, params)
-        super().__init__(*args, **kwargs)
-        self.import_steam = shared.schema.get_boolean("lutris-import-steam")
-        self.db_location = self.source.location / "pga.db"
-        self.db_connection = connect(self.db_location)
-        self.db_request_params = {"import_steam": self.import_steam}
-        self.db_cursor = self.db_connection.execute(
-            self.db_games_request, self.db_request_params
-        )
-
-    def __next__(self) -> Optional[Game]:
-        row = None
-        try:
-            row = self.db_cursor.__next__()
-        except StopIteration as error:
-            self.db_connection.close()
-            raise error
 
+        # Create games from the DB results
+        for row in cursor:
             # Create game
             values = {
                 "version": shared.SPEC_VERSION,
@@ -66,7 +53,8 @@ class LutrisSourceIterator(SourceIterator):
             if image_path.exists():
                 save_cover(values["game_id"], resize_cover(image_path))
 
-        return game
+            # Produce game
+            yield game
 
 
 class LutrisSource(Source):

src/importer/sources/source.py

@@ -2,7 +2,7 @@ import sys
 from abc import abstractmethod
 from collections.abc import Iterable, Iterator
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Generator
 
 from src.game import Game
 
@@ -11,20 +11,28 @@ class SourceIterator(Iterator):
     """Data producer for a source of games"""
 
     source: "Source" = None
+    generator: Generator = None
 
     def __init__(self, source: "Source") -> None:
         super().__init__()
         self.source = source
+        self.generator = self.generator_builder()
 
     def __iter__(self) -> "SourceIterator":
         return self
 
-    @abstractmethod
     def __next__(self) -> Optional[Game]:
-        """Get the next generated game from the source.
+        return next(self.generator)
-        Raises StopIteration when exhausted.
+
-        May raise any other exception signifying an error on this specific game.
+    @abstractmethod
-        May return None when a game has been skipped without an error."""
+    def generator_builder(self) -> Generator[Optional[Game], None, None]:
+        """
+        Method that returns a generator that produces games
+        * Should be implemented as a generator method
+        * May yield `None` when an iteration hasn't produced a game
+        * In charge of handling per-game errors
+        * Returns when exhausted
+        """
 
 
 class Source(Iterable):

src/importer/sources/steam_source.py

@@ -1,7 +1,7 @@
 import re
 from pathlib import Path
 from time import time
-from typing import Iterator, Optional
+from typing import Iterable, Optional, Generator
 
 from src import shared
 from src.game import Game
@@ -22,58 +22,54 @@ from src.utils.steam import SteamHelper, SteamInvalidManifestError
 
 class SteamSourceIterator(SourceIterator):
     source: "SteamSource"
-    manifests: set = None
-    manifests_iterator: Iterator[Path] = None
-    installed_state_mask: int = 4
-    appid_cache: set = None
 
-    def __init__(self, *args, **kwargs) -> None:
+    def get_manifest_dirs(self) -> Iterable[Path]:
-        super().__init__(*args, **kwargs)
+        """Get dirs that contain steam app manifests"""
-
-        self.appid_cache = set()
-        self.manifests = set()
-
-        # Get dirs that contain steam app manifests
         libraryfolders_path = self.source.location / "steamapps" / "libraryfolders.vdf"
         with open(libraryfolders_path, "r") as file:
             contents = file.read()
-        steamapps_dirs = [
+        return [
             Path(path) / "steamapps"
             for path in re.findall('"path"\s+"(.*)"\n', contents, re.IGNORECASE)
         ]
 
-        # Get app manifests
+    def get_manifests(self) -> Iterable[Path]:
-        for steamapps_dir in steamapps_dirs:
+        """Get app manifests"""
+        manifests = set()
+        for steamapps_dir in self.get_manifest_dirs():
             if not steamapps_dir.is_dir():
                 continue
-            self.manifests.update(
+            manifests.update(
                 [
                     manifest
                     for manifest in steamapps_dir.glob("appmanifest_*.acf")
                     if manifest.is_file()
                 ]
             )
+        return manifests
 
-        self.manifests_iterator = iter(self.manifests)
+    def generator_builder(self) -> Optional[Game]:
-
+        """Generator method producing games"""
-    def __next__(self) -> Optional[Game]:
+        appid_cache = set()
+        manifests = self.get_manifests()
+        for manifest in manifests:
             # Get metadata from manifest
-        manifest_path = next(self.manifests_iterator)
             steam = SteamHelper()
             try:
-            local_data = steam.get_manifest_data(manifest_path)
+                local_data = steam.get_manifest_data(manifest)
             except (OSError, SteamInvalidManifestError):
-            return None
+                continue
 
             # Skip non installed games
-        if not int(local_data["stateflags"]) & self.installed_state_mask:
+            INSTALLED_MASK: int = 4
-            return None
+            if not int(local_data["stateflags"]) & INSTALLED_MASK:
+                continue
 
             # Skip duplicate appids
             appid = local_data["appid"]
-        if appid in self.appid_cache:
+            if appid in appid_cache:
-            return None
+                continue
-        self.appid_cache.add(appid)
+            appid_cache.add(appid)
 
             # Build game from local data
             values = {
@@ -87,16 +83,17 @@ class SteamSourceIterator(SourceIterator):
             game = Game(values, allow_side_effects=False)
 
             # Add official cover image
-        cover_path = (
+            image_path = (
                 self.source.location
                 / "appcache"
                 / "librarycache"
                 / f"{appid}_library_600x900.jpg"
             )
-        if cover_path.is_file():
+            if image_path.is_file():
-            save_cover(game.game_id, resize_cover(cover_path))
+                save_cover(game.game_id, resize_cover(image_path))
 
-        return game
+            # Produce game
+            yield game
 
 
 class SteamSource(Source):