unwind/unwind/web_models.py

85 lines
2.5 KiB
Python

from dataclasses import dataclass
from typing import Container, Iterable
from . import imdb, models
URL = str
Score100 = int # [0, 100]
@dataclass
class Rating:
canonical_title: str
imdb_score: Score100 | None
imdb_votes: int | None
media_type: str
movie_imdb_id: str
original_title: str | None
release_year: int
user_id: str | None
user_score: Score100 | None
@classmethod
def from_movie(cls, movie: models.Movie, *, rating: models.Rating | None = None):
return cls(
canonical_title=movie.title,
imdb_score=movie.imdb_score,
imdb_votes=movie.imdb_votes,
media_type=movie.media_type,
movie_imdb_id=movie.imdb_id,
original_title=movie.original_title,
release_year=movie.release_year,
user_id=str(rating.user_id) if rating else None,
user_score=rating.score if rating else None,
)
@dataclass
class RatingAggregate:
canonical_title: str
imdb_score: Score100 | None
imdb_votes: int | None
link: URL
media_type: str
original_title: str | None
user_scores: list[Score100]
year: int
@classmethod
def from_movie(cls, movie: models.Movie, *, ratings: Iterable[models.Rating] = []):
return cls(
canonical_title=movie.title,
imdb_score=movie.imdb_score,
imdb_votes=movie.imdb_votes,
link=imdb.movie_url(movie.imdb_id),
media_type=movie.media_type,
original_title=movie.original_title,
user_scores=[r.score for r in ratings],
year=movie.release_year,
)
def aggregate_ratings(
ratings: Iterable[Rating], user_ids: Container[str]
) -> Iterable[RatingAggregate]:
aggr: dict[str, RatingAggregate] = {}
for r in ratings:
mov = aggr.setdefault(
r.movie_imdb_id,
RatingAggregate(
canonical_title=r.canonical_title,
imdb_score=r.imdb_score,
imdb_votes=r.imdb_votes,
link=imdb.movie_url(r.movie_imdb_id),
media_type=r.media_type,
original_title=r.original_title,
user_scores=[],
year=r.release_year,
),
)
# XXX do we need this? why don't we just get the ratings we're supposed to aggregate?
if r.user_score is not None and r.user_id in user_ids:
mov.user_scores.append(r.user_score)
return aggr.values()