advice.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from cStringIO import StringIO
from urllib import urlopen

from PIL import Image, ImageDraw, ImageFont
from requests import get

# Poster width and hight
WIDTH = 214
HIGHT = 317


def get_poster_url(imdb_id):
    """Return poster URL if exist or False otherwise"""

    response = get('http://imdbapi.org/', params={'ids': imdb_id})
    json = response.json()

    if json[0].__contains__('poster'):
        return json[0]['poster']
    else:
        return False


def main():

    # Getting list of unique genres from input file
    genres = [line.split(',')[2] for line in open('films.txt', 'r')]
    genres = [item.strip() for item in genres]
    unique_genres = [item for item in set(genres)]

    # Getting maximum value of posters in column (genre)
    counts = [''.join(genres).count(item) for item in unique_genres]
    max_posters = max(counts)

    # Creating canvas
    canv_width = 20 + (20 + WIDTH) * len(unique_genres)
    canv_hight = 100 + HIGHT * max_posters + 20
    canvas = Image.new('RGB', (canv_width, canv_hight), 'black')
    draw = ImageDraw.Draw(canvas)

    # Adding genres
    font = ImageFont.truetype(
            "/usr/share/fonts/ubuntu-font-family/Ubuntu-C.ttf",
            42)
    shift = 0

    for genre in unique_genres:
        draw.text((20 + shift, 20), genre, fill='white', font=font)
        shift += (20 + WIDTH)

    # Adding posters
    y = [100] * len(unique_genres)  # List of y values for each column (genre)
    count = [0] * len(unique_genres)  # Number of posters in each column
    xx = 0
    yy = 0
    with open('films.txt', 'r') as f:
        for line in f:
            imdb_id = line.split(',')[0]
            title = line.split(',')[1].strip()
            genre = line.split(',')[2].strip()

            poster_url = get_poster_url(imdb_id)

            if poster_url:
                # Open image(poster) from URL
                fs = StringIO(urlopen(poster_url).read())
                poster = Image.open(fs)

                if poster.size != (WIDTH, HIGHT):
                    poster = poster.resize (WIDTH, HIGHT)

                # Sorting by genre
                for col in range(0, len(unique_genres)):
                    if genre == unique_genres[col]:
                        xx = 20 + (20 + WIDTH)*col
                        y[col] = y[col] + HIGHT * count[col]
                        yy = y[col]
                        count[col] = count[col] + 1

                canvas.paste(poster, (xx, yy))

            else:
                # Drawing caption if no poster available
                #draw.text((xx, yy), title, fill='white')
                pass

    canvas.save('advice.png')


if __name__ == '__main__':
    main()