Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor: Add asyncio support #563

Merged
merged 2 commits into from
May 11, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Refactor: Make Walker methods async
  • Loading branch information
glensc committed May 11, 2024
commit 32088167c3819b69296d2ff8118e4dbe1de443b2
72 changes: 40 additions & 32 deletions plextraktsync/plan/Walker.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from plextraktsync.trakt.TraktItem import TraktItem

if TYPE_CHECKING:
from typing import Any, Generator, Iterable
from typing import Any, AsyncGenerator, AsyncIterable, Generator, Iterable

from plexapi.video import Episode

Expand Down Expand Up @@ -69,7 +69,7 @@ def print_plan(self, print):
if self.plan.episodes:
print(f"Sync Episodes: {[x.title for x in self.plan.episodes]}")

def get_plex_movies(self) -> Generator[PlexLibraryItem, Any, None]:
async def get_plex_movies(self) -> Generator[PlexLibraryItem, Any, None]:
"""
Iterate over movie sections unless specific movie is requested
"""
Expand All @@ -80,45 +80,48 @@ def get_plex_movies(self) -> Generator[PlexLibraryItem, Any, None]:
else:
return

yield from movies
async for m in movies:
yield m

def find_movies(self) -> Generator[Media, Any, None]:
for plex in self.get_plex_movies():
async def find_movies(self) -> Generator[Media, Any, None]:
async for plex in self.get_plex_movies():
movie = self.mf.resolve_any(plex)
if not movie:
continue
yield movie

def get_plex_shows(self) -> Generator[PlexLibraryItem, Any, None]:
async def get_plex_shows(self) -> AsyncGenerator[PlexLibraryItem, Any, None]:
if self.plan.shows:
shows = self.media_from_items("show", self.plan.shows)
it = self.media_from_items("show", self.plan.shows)
elif self.plan.show_sections:
shows = self.media_from_sections(self.plan.show_sections)
it = self.media_from_sections(self.plan.show_sections)
else:
return

yield from shows
async for m in it:
yield m

def find_episodes(self):
async def find_episodes(self):
if self.plan.episodes:
yield from self.get_plex_episodes(self.plan.episodes)
async for m in self.get_plex_episodes(self.plan.episodes):
yield m

# Preload plex shows
plex_shows: dict[int, PlexLibraryItem] = {}
self.logger.info("Preload shows data")
for show in self.get_plex_shows():
async for show in self.get_plex_shows():
plex_shows[show.key] = show
self.logger.info(f"Preloaded shows data ({len(plex_shows)} shows)")

# Preload matches for shows
show_cache: dict[int, Media] = {}
self.logger.info("Preload shows matches")
it = self.progressbar(plex_shows.items(), desc="Processing show matches")
for show_id, ps in it:
async for show_id, ps in it:
show_cache[show_id] = self.mf.resolve_any(ps)
self.logger.info(f"Preloaded shows matches ({len(show_cache)} shows)")

for ep in self.episodes_from_sections(self.plan.show_sections):
async for ep in self.episodes_from_sections(self.plan.show_sections):
show_id = ep.show_id
ep.show = plex_shows[show_id]
show = show_cache[show_id] if show_id in show_cache else None
Expand All @@ -130,14 +133,15 @@ def find_episodes(self):
show_cache[show_id] = m.show
yield m

def walk_shows(self, shows: set[Media], title="Processing Shows"):
async def walk_shows(self, shows: set[Media], title="Processing Shows"):
if not shows:
return
yield from self.progressbar(shows, desc=title)
async for show in self.progressbar(shows, desc=title):
yield show

def get_plex_episodes(self, episodes: list[Episode]) -> Generator[Media, Any, None]:
async def get_plex_episodes(self, episodes: list[Episode]) -> Generator[Media, Any, None]:
it = self.progressbar(episodes, desc="Processing episodes")
for pe in it:
async for pe in it:
guid = PlexGuid(pe.grandparentGuid, "show")
show = self.mf.resolve_guid(guid)
if not show:
Expand All @@ -149,32 +153,34 @@ def get_plex_episodes(self, episodes: list[Episode]) -> Generator[Media, Any, No
me.show = show
yield me

def media_from_sections(self, sections: list[PlexLibrarySection]) -> Generator[PlexLibraryItem, Any, None]:
async def media_from_sections(self, sections: list[PlexLibrarySection]) -> AsyncGenerator[PlexLibraryItem, Any, None]:
for section in sections:
with measure_time(f"{section.title_link} processed", extra={"markup": True}):
self.set_window_title(f"Processing {section.title}")
it = self.progressbar(
section.pager(),
desc=f"Processing {section.title_link}",
)
yield from it
async for m in it:
yield m

def episodes_from_sections(self, sections: list[PlexLibrarySection]) -> Generator[PlexLibraryItem, Any, None]:
async def episodes_from_sections(self, sections: list[PlexLibrarySection]) -> Generator[PlexLibraryItem, Any, None]:
for section in sections:
with measure_time(f"{section.title_link} processed", extra={"markup": True}):
self.set_window_title(f"Processing {section.title}")
it = self.progressbar(
section.pager("episode"),
desc=f"Processing {section.title_link}",
)
yield from it
async for m in it:
yield m

def media_from_items(self, libtype: str, items: list) -> Generator[PlexLibraryItem, Any, None]:
async def media_from_items(self, libtype: str, items: list) -> AsyncGenerator[PlexLibraryItem, Any, None]:
it = self.progressbar(items, desc=f"Processing {libtype}s")
for m in it:
async for m in it:
yield PlexLibraryItem(m, plex=self.plex)

def episode_from_show(self, show: Media) -> Generator[Media, Any, None]:
async def episode_from_show(self, show: Media) -> Generator[Media, Any, None]:
for pe in show.plex.episodes():
me = self.mf.resolve_any(pe, show)
if not me:
Expand All @@ -183,24 +189,26 @@ def episode_from_show(self, show: Media) -> Generator[Media, Any, None]:
me.show = show
yield me

def progressbar(self, iterable: Iterable, **kwargs):
async def progressbar(self, iterable: AsyncIterable | Iterable, **kwargs):
if self._progressbar:
pb = self._progressbar(iterable, **kwargs)
with pb as it:
yield from it
async for m in it:
yield m
else:
yield from iterable
async for m in iterable:
yield m

def media_from_traktlist(self, items: Iterable, title="Trakt watchlist") -> Generator[Media, Any, None]:
async def media_from_traktlist(self, items: AsyncIterable, title="Trakt watchlist") -> Generator[Media, Any, None]:
it = self.progressbar(items, desc=f"Processing {title}")
for media in it:
async for media in it:
tm = TraktItem(media)
m = self.mf.resolve_trakt(tm)
yield m

def media_from_plexlist(self, items: Iterable) -> Generator[Media, Any, None]:
async def media_from_plexlist(self, items: AsyncIterable) -> Generator[Media, Any, None]:
it = self.progressbar(items, desc="Processing Plex watchlist")
for media in it:
async for media in it:
pm = PlexLibraryItem(media, plex=self.plex)
m = self.mf.resolve_any(pm)
if not m:
Expand Down
4 changes: 2 additions & 2 deletions plextraktsync/sync/Sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,10 @@ async def sync(self, walker: Walker, dry_run=False):
pm.hook.init(sync=self, pm=pm, is_partial=is_partial, dry_run=dry_run)

if self.config.need_library_walk:
for movie in walker.find_movies():
async for movie in walker.find_movies():
await pm.ahook.walk_movie(movie=movie, dry_run=dry_run)

for episode in walker.find_episodes():
async for episode in walker.find_episodes():
await pm.ahook.walk_episode(episode=episode, dry_run=dry_run)

await pm.ahook.fini(walker=walker, dry_run=dry_run)
2 changes: 1 addition & 1 deletion plextraktsync/sync/SyncRatingsPlugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def init(self):

@hookimpl
async def fini(self, walker: Walker, dry_run: bool):
for show in walker.walk_shows(self.shows, title="Syncing show ratings"):
async for show in walker.walk_shows(self.shows, title="Syncing show ratings"):
self.sync_ratings(show, dry_run=dry_run)

@hookimpl
Expand Down