102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
from dataclasses import dataclass
|
|
from typing import Container, Iterable
|
|
|
|
from . import imdb, models, types
|
|
|
|
type URL = str
|
|
|
|
|
|
@dataclass
|
|
class Rating:
|
|
canonical_title: str
|
|
imdb_score: types.Score100 | None
|
|
imdb_votes: int | None
|
|
media_type: str
|
|
movie_imdb_id: types.ImdbMovieId
|
|
original_title: str | None
|
|
release_year: int
|
|
user_id: types.UserIdStr | None
|
|
user_score: types.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: types.Score100 | None
|
|
imdb_votes: int | None
|
|
link: URL
|
|
media_type: str
|
|
original_title: str | None
|
|
user_scores: list[types.Score100]
|
|
year: int
|
|
awards: list[str]
|
|
|
|
@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,
|
|
awards=[],
|
|
)
|
|
|
|
|
|
def _serialize_award(award: models.Award) -> str:
|
|
if award.category == "oscars":
|
|
return f"{award.category}:{award.name}"
|
|
elif award.category.startswith("imdb-"):
|
|
return f"{award.category}:{award.position}"
|
|
raise RuntimeError(f"Unsupported category: {award.category}")
|
|
|
|
|
|
def aggregate_ratings(
|
|
ratings: Iterable[Rating],
|
|
user_ids: Container[types.UserIdStr],
|
|
*,
|
|
awards_dict: dict[types.ImdbMovieId, list[models.Award]] | None = None,
|
|
) -> Iterable[RatingAggregate]:
|
|
if awards_dict is None:
|
|
awards_dict = {}
|
|
|
|
aggr: dict[types.ImdbMovieId, RatingAggregate] = {}
|
|
|
|
for r in ratings:
|
|
awards = awards_dict.get(r.movie_imdb_id, [])
|
|
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,
|
|
awards=sorted(_serialize_award(a) for a in awards),
|
|
),
|
|
)
|
|
# 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()
|