Skip to content

Instantly share code, notes, and snippets.

@Decstasy
Last active May 3, 2023 19:17
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Decstasy/a97d911f2a41b01b9ff28eba19210204 to your computer and use it in GitHub Desktop.
Save Decstasy/a97d911f2a41b01b9ff28eba19210204 to your computer and use it in GitHub Desktop.
Creates a heatmap for a series based on imdb ratings, so you know if a series gets bad in advance ;)
#!/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__ = "git@decstasy.de"
__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')
@Decstasy
Copy link
Author

Decstasy commented Apr 5, 2023

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment