Movie Diary.py

Movie Diary.py

# coding: utf-8

import requests
import json
import appex
import dialogs
import re
import datetime
import keychain
import console
import sys
import cPickle
from urllib import quote, unquote

class NoResultsError (Exception): pass
class NoMoviePickError (Exception): pass
class NoRatingError (Exception): pass
class TmdbConnectionError (Exception): pass
class AirtableConnectionError (Exception): pass
class MissingConfigError (Exception): pass
class InvalidColumnError (Exception): pass

class MovieDiary():

	def __init__(self):
		moviediary_config = keychain.get_password('Movie Diary', 'Config')
		
		if moviediary_config == None:
			moviedb_api = keychain.get_password('MovieDB', 'API')
			airtable_api = keychain.get_password('Airtable', 'API')
			airtable_db = keychain.get_password('Airtable', 'Movie Diary')
			airtable_table = keychain.get_password('Airtable', 'Movie Diary Table')
			
			moviediary_config = dialogs.form_dialog(title='Movie Diary Configuration', sections=[('MovieDB', [{'type': 'text', 'key': 'moviedb_api', 'value': moviedb_api if moviedb_api is not None else '84cef43ccf02b1ba6093c9694ed671c9', 'title': 'MovieDB API Token'}]), ('Airtable', [{'type': 'text', 'key': 'airtable_api', 'value': airtable_api if airtable_api is not None else '', 'title': 'Airtable API Key'}, {'type': 'text', 'key': 'airtable_db', 'value': airtable_db if airtable_db is not None else '', 'title': 'Airtable database ID'}, {'type': 'text', 'key': 'airtable_table', 'value': airtable_table if airtable_table is not None else 'Table 1', 'title': 'Airtable table name'}]), ('Custom', [{'type': 'switch', 'key': 'set_date_manually', 'value': False, 'title': 'Set date manually'}, {'type': 'switch', 'key': 'add_time_to_date', 'value': False, 'title': 'Add time to date'}]),('Extra Fields', [{'type': 'switch', 'key': 'directors_field', 'value': True, 'title': 'Directors'}, {'type': 'switch', 'key': 'genres_field', 'value': False, 'title': 'Genres'}, {'type': 'switch', 'key': 'runtime_field', 'value': False, 'title': 'Runtime'}, {'type': 'switch', 'key': 'cast_field', 'value': False, 'title': 'Cast'}, {'type': 'switch', 'key': 'imdb_field', 'value': False, 'title': 'IMDB URL'}])])
			
			if moviediary_config == None:
				raise MissingConfigError('You must setup and confirm the Movie Diary configuration before continuing.')
			else:
				if moviediary_config['moviedb_api'] == '':
					moviedb_api = console.input_alert('Insert your TMDB API key', '', '84cef43ccf02b1ba6093c9694ed671c9')
					if moviedb_api == None:
						raise MissingConfigError('You need a valid MovieDB API key')
					else:
						moviediary_config['moviedb_api'] = moviedb_api

				if moviediary_config['airtable_api'] == '':
					airtable_api = console.input_alert('Insert your Airtable API key')
					if airtable_api == None:
						raise MissingConfigError('You need a valid Airtable API key')
					else:
						moviediary_config['airtable_api'] = airtable_api

				if moviediary_config['airtable_db'] == '':
					airtable_db = console.input_alert('Insert your Airtable database ID')
					if airtable_db == None:
						raise MissingConfigError('You need the ID of your database')
					else:
						moviediary_config['airtable_db'] = airtable_db

				if moviediary_config['airtable_table'] == '':
					airtable_table = console.input_alert('Insert the name of yout Airtable table', '', 'Table 1')
					if airtable_table == None:
						raise MissingConfigError('You must insert the name of the table in your database.')
					else:
						moviediary_config['airtable_table'] = quote(airtable_table)
				
				config = moviediary_config
				keychain.set_password('Movie Diary', 'Config', cPickle.dumps(config))
		else:
			config = cPickle.loads(moviediary_config)

		self.moviedb_api = config['moviedb_api']
		self.airtable_api = config['airtable_api']
		self.airtable_db = config['airtable_db']
		self.airtable_table = config['airtable_table']
		self.set_date_manually = config['set_date_manually']
		self.add_time_to_date = config['add_time_to_date']
		self.directors_field = config['directors_field']
		self.genres_field = config['genres_field']
		self.runtime_field = config['runtime_field']
		self.cast_field = config['cast_field']
		self.imdb_field = config['imdb_field']

	@staticmethod
	def getyear(d, raw=False):
		if d is None or d == '':
			return ''
		elif raw:
			return str(d[:4])
		else:
			return ' (%s)' % d[:4]

	@staticmethod
	def getgenres(genres):
		return '/'.join([genre['name'] for genre in genres])

	def getcredits(self, url, params):
		req = requests.get('%s/credits' % (url), params=params)

		if req.status_code == 200:
			res = json.loads(req.text)
			directors = []
			cast = []
			
			if self.directors_field:
				directors = ', '.join([director['name'] for director in res['crew'] if director['job'] == 'Director'])
			if self.cast_field:
				cast = ', '.join([res['cast'][i]['name'] for i in range(min(5, len(res['cast'])))])
			
			return (cast, directors)
		else:
			raise TmdbConnectionError(req.text)

	def getdate(self):
		if self.set_date_manually and self.add_time_to_date:
			return dialogs.datetime_dialog().isoformat()
		elif self.set_date_manually:
			return dialogs.date_dialog().isoformat()
		elif self.add_time_to_date:
			return datetime.datetime.now().isoformat()
		else:
			return datetime.datetime.now().date().isoformat()

	def journal(self, data):
		headers = {
			'Authorization': 'Bearer %s' % self.airtable_api,
			'Content-type': 'application/json'
		}

		req = requests.post('https://api.airtable.com/v0/{0}/{1}'.format(self.airtable_db, self.airtable_table), headers=headers, data=json.dumps({'fields': data}))

		if req.status_code == 200:
			console.hud_alert('Added movie', 'success')
		elif req.status_code == 422:
			raise InvalidColumnError(json.loads(req.text))
		else:
			raise AirtableConnectionError(req.text)

	def getmovie(self, url, params):
		req = requests.get(url, params=params)

		if req.status_code == 200:
			res = json.loads(req.text)

			fields = {
				'Overview': res['overview'],
				'Title': res['title'],
				'Year': self.getyear(res['release_date'], True),
				'Date': self.getdate(),
				'Rating': dialogs.list_dialog("Rate '{0}'".format(res['title']), ['★★★★★', '★★★★½', '★★★★', '★★★½', '★★★', '★★½', '★★', '★½', '★', '½'])
			}
			
			if self.cast_field or self.directors_field:
				credits = self.getcredits(url, params)
				if self.cast_field:
					fields['Cast'] = credits[0]
				if self.directors_field:
					fields['Directors'] = credits[1]
			
			if self.runtime_field:
				fields['Runtime'] = res['runtime']
			
			if self.imdb_field:
				fields['IMDB'] = 'http://www.imdb.com/title/%s/' % res['imdb_id']
			
			if self.genres_field:
				fields['Genres'] = self.getgenres(res['genres'])

			if res['poster_path'] is not None:
				fields['Poster'] = [{'url': 'https://image.tmdb.org/t/p/original%s' % res['poster_path']}]
			if fields['Rating'] is not None:
				return self.journal(fields)
			else:
				raise NoRatingError()
		else:
			raise TmdbConnectionError(req.text)

	def log(self):
		console.show_activity()
		try:
			url_match = re.match(r'^https?://(?:www\.)?imdb\.com/title/(tt\d+)/?', appex.get_url())
			params = {
				'api_key': self.moviedb_api,
				'external_source': 'imdb_id'
			}
			return self.getmovie('https://api.themoviedb.org/3/movie/%s' % (url_match.group(1)), params)
		except TypeError:
			params = {
				'api_key': self.moviedb_api,
				'query': console.input_alert('Search for movie', '', sys.argv[1] if len(sys.argv) > 1 else '')
			}
			req = requests.post('https://api.themoviedb.org/3/search/movie', params=params)

			if req.status_code == 200:
				res = json.loads(req.text)

				if res['total_results'] > 1:
					results_map = {e['title'] + self.getyear(e['release_date']) : e for e in res['results']}
					movie_pick = dialogs.list_dialog('Pick a movie', [e['title'] + self.getyear(e['release_date']) for e in res['results']])
					if movie_pick is not None:
						return self.getmovie('https://api.themoviedb.org/3/movie/%s' % (results_map[movie_pick]['id']), {'api_key': self.moviedb_api})
					else:
						raise NoMoviePickError()
				elif res['total_results'] == 1:
					return self.getmovie('https://api.themoviedb.org/3/movie/%s' % (res['results'][0]['id']), {'api_key': self.moviedb_api})
				else:
					raise NoResultsError()
			else:
				raise TmdbConnectionError(req.text)

if __name__ == '__main__':
	md = MovieDiary()
	try:
		md.log()
	except MissingConfigError as e:
		console.alert('Missing configuration', e)
	except NoResultsError:
		console.alert('No Results', 'Couldn\'t find any movie matching your query.')
	except AirtableConnectionError as e:
		console.alert('Failed to connect to Airtable', e)
	except TmdbConnectionError as e:
		console.alert('Failed to connect to MovieDB', e)
	except NoMoviePickError:
		console.alert('No movie selected', 'You gotta pick a movie for the script to work.')
	except NoRatingError:
		console.alert('No rating', 'You gotta rate the movie for the script to work.')
	except InvalidColumnError as e:
		console.alert('Invalid Data', '{0}: {1}'.format(e[0]['error']['type'], e[0]['error']['message']))