Creates a heatmap for a series based on imdb ratings, so you know if a series gets bad in advance ;)

imdbheatmap.py

#!/usr/bin/env python3

__author__ = "Dennis Ullrich"
__copyright__ = "Copyright 2023, Dennis Ullrich"
__credits__ = ["Dennis Ullrich"]
__license__ = "GPL"
__version__ = "3.0.0"
__maintainer__ = "Dennis Ullrich"
__email__ = "[email protected]"
__status__ = "Fun"

from pprint import pprint
# pip3.9 install git+https://github.com/cinemagoer/cinemagoer matplotlib seaborn
import imdb
import matplotlib
import seaborn as sns

class SeriesHandler:
    def __init__(self, ia, SeriesName):
        self.ia = ia
        self._series = self._search(SeriesName)
        self.ratings = self._get_ratings()

    def __repr__(self):
        return str(self._series)

    @staticmethod
    def _fill_with_zero(original_list):
    # Seaborn braucht homogene längen der listen
        max_len = max(len(sublist) for sublist in original_list)
        return [sublist + [0] * (max_len - len(sublist)) for sublist in original_list]
    
    def _replace_none_with_zero(self, lst):
    # Manchmal kommen None objects in den zurückgegebenen listen vor, die müssen für seaborn ersetzt werden
        for i in range(len(lst)):
            if isinstance(lst[i], list):
                self._replace_none_with_zero(lst[i])
            elif lst[i] is None:
                lst[i] = 0
        return lst

    def _cleanup_zero_lists(self, lst):
    # Im Fall der Serie Fargo kommt eine 5. Staffel raus, die momentan nur ein None gibt [None].
    # Das wird durch die anderen funktionen aufgefüllt und ersetzt sodass ein [0,0,0,0,0] rauskommt.
    # Da uns so ein scheiß nichts bringt, durchsuchen wir die liste der listen auf listen mit ausschließlich nullen.
        i = 0
        while i < len(lst):
            if isinstance(lst[i], list):
                if all(val == 0 for val in lst[i]):
                    del lst[i]
                else:
                    self._cleanup_zero_lists(lst[i])
                    i += 1
            else:
                i += 1
        return lst
    
    def _cleanup_ratings(self, ratings):
    # Wrapper um alle nötigen cleanups auf die Rating-Daten anzuwenden
        tmp = self._fill_with_zero(ratings)
        tmp = self._replace_none_with_zero(tmp)
        return self._cleanup_zero_lists(tmp)

    def _search(self, SeriesName):
    # Sucht eien serie und gibt diese zurück
        search_results = self.ia.search_movie(SeriesName)
        series_id = search_results[0].getID()
        return self.ia.get_movie(series_id)
    
    def _get_ratings(self):
    # Holt sich die Ratings jeder Episode und plegt diese in eine zweidimensionale Liste ein
    # Die erste dimension sind die staffeln, die zweite die episoden. Beispiel zwei staffeln, je drei folgen:
    # [[8.4, 8.2, 8.9], [7.2, 7.7, 7.4]]
        self.ia.update(self._series, "episodes")
        seasons = sorted(self._series['episodes'].keys())
        ratings = [[self._series['episodes'][season][episode]['rating'] if 'rating' in self._series['episodes'][season][episode].keys() else None for episode in self._series['episodes'][season]] for season in seasons]
        return self._cleanup_ratings(ratings)
        
    def get_x_label(self):
    # Label für x-achse also episoden (seaborn stinkt)
        return list(range(1,max(len(sublist) for sublist in self.ratings)+1))
    
    def get_y_label(self):
    # Label für y-achse also staffeln
        return list(range(1,len(self.ratings)+1))
    
    def get_xy_figsize(self, xmin=3, ymin=3):
        x, y = float(len(self.get_x_label()))*0.8, float(len(self.get_y_label())*0.8)
        if y <= ymin:
            y = ymin
        if x <= xmin:
            x = xmin
        return (x, y)
    
    @staticmethod
    def get_locator(label):
    # Die Beschriftungen sollten mittig sein, daher müssen wir etwas rechnen um den seaborn locator zu füttern
        return [i-0.5 for i in label]

if __name__ == '__main__':
    # IMDB instanz erzeugen
    ia = imdb.Cinemagoer()
    series_name = input("Name der Serie: ")
    s = SeriesHandler(ia, series_name)

    # Drecksplot bauen
    fig, ax = matplotlib.pyplot.subplots(figsize=s.get_xy_figsize())
    cmap = sns.color_palette("RdYlGn", as_cmap=True)
    cmap.set_bad(color='grey')
    sns.heatmap(s.ratings, cmap=cmap, vmin=1, vmax=10, annot=True, square=True, cbar_kws={'orientation': 'vertical', 'shrink': 0.8}, linewidths=2)

    ax.set_title(f'IMDB Bewertungen von "{s}"')

    ax.set_ylabel('Staffel')
    ylabel = s.get_y_label()
    ax.yaxis.set_major_locator(matplotlib.ticker.FixedLocator(s.get_locator(ylabel)))
    ax.set_yticks=ylabel
    ax.set(yticklabels=ylabel)

    ax.set_xlabel('Folge')
    xlabel = s.get_x_label()
    ax.xaxis.set_major_locator(matplotlib.ticker.FixedLocator(s.get_locator(xlabel)))
    ax.set_xticks=xlabel
    ax.set(xticklabels=xlabel)

    matplotlib.pyplot.savefig(f'{s}.png')

The comments are in german because its just a fun project... Just use a translator lol