IMDBtoCSV: Converts IMDB data sets in plain text to a simple CSV

$IMDBtoCSV.md

Como usar

1. Descargar y descomprimir el ultimo set de datos proporcionado por IMDB aca:

• actors.list.gz: ftp://ftp.funet.fi/pub/mirrors/ftp.imdb.com/pub/frozendata/actors.list.gz
• actresses.list.gz: ftp://ftp.funet.fi/pub/mirrors/ftp.imdb.com/pub/frozendata/actresses.list.gz

2. Llamar con: python imdbtocsv.py actresses.list actors.list o cada uno por separado

En imdbtocsv.py se puede leer un pequeño paso a paso de como se parsea el archivo

Opciones completas

python imdbtocsv.py [archivos] -flags con los flags siendo:

• -h, --help Show this help message and exit
• -t N, --test N Genera un archivo de test con N lineas elegidas al azar
• -kb , --Kevin-Bacon El archivo generado tendra la linea de Kevin Bacon como primera

Ejemplos de conversion

Ejemplo (acotado) de actress/actors.list

Bacon, Kevin	12th Annual Screen Actors Guild Awards (2006) (TV)  [Himself]
			15th Annual Critics' Choice Movie Awards (2010) (TV)  [Himself]
			17th Annual Screen Actors Guild Awards (2011) (TV)  [Himself]
			18th Annual Screen Actors Guild Awards (2012) (TV)  [Himself]
			1995 MTV Video Music Awards (1995) (TV)  [Himself]
			1996 MTV Video Music Awards (1996) (TV)  [Himself - Presenter]
			[...]
			A Few Good Men (1992)  [Capt. Jack Ross]  <4>
			A Few Good Men: From Stage to Screen (2001) (V)  (archive footage) (uncredited)  [Capt. Jack Ross]
			A Little Vicious (1991)  [Narrator]
			A Look Behind the Scenes: Super (2011)  [Himself]
			AFI's 100 Years... 100 Stars: America's Greatest Screen Legends (1999) (TV)  [Himself]
			American Empire (2011) {{SUSPENDED}}
			Animal House (1978)  [Chip Diller]  <7>
			[...]
			"Alter Egos" (2011) {Kevin Bacon (#1.4)}  [Himself]
			"America's Game: The Superbowl Champions" (2006) {1996 Green Bay Packers (#1.5)}  [Narrator]
			"Big Morning Buzz Live" (2011) {Kevin Bacon/Chrissy Teigen/The Bacon Brothers (#7.76)}  [Himself - Guest]  <2>
			"Biography" (1987) {Graham Nash: A Life in Harmony}  (voice)  [Himself - Narrator]
			"Biography" (1987) {Ron Howard: Hollywood's Favorite Son}  (archive footage)  [Jack Swigert]

Ejemplo de actresses/actors.csv

Bacon Kevin,A Few Good Men (1992),A Little Vicious (1991),Animal House (1978),Apollo 13 (1995/I),Balto (1995),Beauty Shop (2005),Beyond All Boundaries (2009),Black Mass (2015),Cavedweller (2004),Cop Car (2015),Crazy Stupid Love (2011),Criminal Law (1988),Death Sentence (2007),Digging to China (1997),Diner (1982),Elephant White (2011),End of the Line (1987),Enormous Changes at the Last Minute (1983),Flatliners (1990),Footloose (1984),Forty Deuce (1982),Friday the 13th (1980),Frost/Nixon (2008),He Said She Said (1991),Hero at Large (1980),Hollow Man (2000),Jayne Mansfields Car (2012),JFK (1991),Lemon Sky (1988),Loverboy (2005),Murder in the First (1995),My Dog Skip (2000),My One and Only (2009),Mystic River (2003),New York Skyride (1994),Only When I Laugh (1981),Patriots Day (2016),Picture Perfect (1997),Pyrates (1991),Queens Logic (1991),Quicksilver (1986),RIPD (2013),Rails & Ties (2007),Saving Angelo (2007),Shes Having a Baby (1988),Sleepers (1996),Starting Over (1979),Stir of Echoes (1999),Super (2010/I),Telling Lies in America (1997),The Air I Breathe (2007),The Air Up There (1994),The Big Green (2014),The Big Picture (1989),The Darkness (2016/I),The Making of Apollo 13 (1995),The River Wild (1994),The Woodsman (2004),These Vagabond Shoes (2009),Tough Day (2014),Trapped (2002/I),Tremors (1990),Where the Truth Lies (2005),White Water Summer (1987),Wild Things (1998),X First Class (2011)

actors.py

'''
Recibe una base de datos de imdb actors.list y/o actresses.list y lo pasa a un csv muy simplificado
Recibe lista de lineas originales y devuelve lista de lineas modificada

1) Separa al actor de su primer credito para mejor manejo
2) Saca la coma entre nombre y apellido del actor para no interpretar como otra columna de csv y saca los simbolos raros de los nombres
3) Descarta todos los creditos que no sirven, explicitados en PALABRAS_A_DESCARTAR
4) Saca todo lo siguiente al año de la pelicula, como el rol y el numero del credito
5) Saca todos los malos caracteres de las peliculas para evitar nombres como X-Men o Birdman or 'The unexp....
6) Genera una lista de listas por cada actor, y descarta todo actor sin creditos (por haber sido, por ejemplo, solo de series)
7) Pasa de lista de listas a lista de lineas y las devuelve
'''

import logging
import re

PALABRAS_A_DESCARTAR = [
	'Himself','Herself', # Premios y talkshows
	'(TV)','(V)', # TV and made for cable movies
	'"', # Episodios de television
	'SUSPENDED', # Peliculas nunca hechas
	'archive footage','uncredited','(scenes deleted)', # No acreditas o donde no hubo real actuacion
	'VG' # Videojuegos
	]
	
PALABRAS_MALAS_PELICULA = [
	'(voice)',
	'(????)'
	]
	
SIMBOLOS_RAROS_ACTOR = '[#\$\'\"\&]'
SIMBOLOS_RAROS_PELICULA = '[,\.\'\-\:]'

def actorstocsv(lineas_originales, flags={}):
	lineas = sacar_introduccion(lineas_originales)
	lineas = separar_primer_credito(lineas)
	lineas = sacar_simbolos_raros_actor(lineas)
	lineas = descartar_creditos_inservible(lineas)
	lineas = sacar_todo_despues_de_anio(lineas)
	lineas = sacar_simbolos_raros_pelicula(lineas)
	lineas = pasar_a_csv(lineas, flags.get('KevinBacon'))
	return lineas

def sacar_introduccion(lineas):
	for i,linea in enumerate(lineas):
		if "Name" in linea and "Titles" in linea:
			logging.debug("Sacada la introduccion")
			return lineas[i+1:]
			
def separar_primer_credito(lineas):
	for num,linea in enumerate(lineas):
		if es_actor(linea):
			for i,l in enumerate(linea.split('\t')):
				if l:
					if i == 0: lineas[num] = l
					elif num+1<len(lineas): lineas[num+1] = '\t\t\t' + l
	logging.debug("Separado el primer actor del credito")
	return lineas

def sacar_simbolos_raros_actor(lineas):
	for i,linea in enumerate(lineas):
		if es_actor(linea):
			linea = linea.replace('(I)','')
			linea = ''.join(linea.split(',',1))
			lineas[i] = re.sub(SIMBOLOS_RAROS_ACTOR,'',linea)
	logging.debug("Sacados los simbolos raros del actor: {}".format(SIMBOLOS_RAROS_ACTOR))
	return lineas
	
def descartar_creditos_inservible(lineas):
	for i,linea in enumerate(lineas):
		if any(palabra in linea for palabra in PALABRAS_A_DESCARTAR):
			lineas[i] = ''
	logging.debug("Descartados los creditos con las siguentes palabras: {}".format(PALABRAS_A_DESCARTAR))
	return [x for x in lineas if x]	

def sacar_todo_despues_de_anio(lineas):
	for i,linea in enumerate(lineas):
		if ')' in lineas[i]: lineas[i] = linea[:linea.rfind(')') +1]
	logging.debug("Sacado todo lo siguiente al año de la pelicula")
	return [x for x in lineas if x]	
	
def sacar_simbolos_raros_pelicula(lineas):
	for i,linea in enumerate(lineas):
		if not es_actor(lineas[i]):
			for mala_palabra in PALABRAS_MALAS_PELICULA:
				linea = linea.replace(mala_palabra,'')
			linea = linea.split('(as ', 1)[0]
			lineas[i] = re.sub(SIMBOLOS_RAROS_PELICULA, '', linea)
	logging.debug("Sacados los simbolos raros de la pelicula: {}".format(SIMBOLOS_RAROS_PELICULA))
	return [x for x in lineas if x != '\n']		

def pasar_a_csv(lineas, kb=False):
	new_lineas = []
	actor = []
	actor.append(lineas[0].strip())
	for i in range(1,len(lineas)):
		if es_actor(lineas[i]):
			new_lineas.append(actor)
			actor = []
			actor.append(lineas[i].strip())
		else:
			actor.append(lineas[i].replace('\t\t\t','').strip())
	if actor != []: new_lineas.append(actor)
	logging.debug("Pasado todo a una lista de lineas para ser pasado a csv")
	if kb:
		for i,nlinea in enumerate(new_lineas):
			if 'Bacon Kevin' == nlinea[0]:
				logging.info("Puesta la linea de Kevin Bacon al comienzo: \n {}".format(nlinea))
				new_lineas.insert(0, new_lineas.pop(i))
				break

	return [x for x in new_lineas if len(x)>1]
	
def es_actor(linea):
	return "\t\t\t" not in linea and linea != '\n'
	

imdbtocsv.py

'''
Recibe una base de datos de imdb como actors.list y/o actresses.list y lo pasa a un csv muy simplificado
Llamar con: python imdbtocsv.py actors.list actresses.list (Tambien se puede por separado cada uno)

0) Usando como lista de lineas al archivo recibido por parametro 
1) Llama al archivo correspondiente para la conversion que lo modificara
2) Exporta a csv la lista de lineas recibida
'''

from sys import argv
import csv
import actors
import argparse
import random
import logging
from time import sleep


def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('files',help='Archivos a procesar', nargs='*', action = 'store')
	parser.add_argument('-t', '--test', help='Genera un archivo de test con N lineas elegidas al azar', action='store')	
	parser.add_argument('-kb', '--Kevin-Bacon', help='El archivo generado tendra la linea de Kevin Bacon como primera', action='append_const', const=("KevinBacon",True), dest='flags')
	args = parser.parse_args()
	flags = dict(args.flags) if args.flags else {}
	logging.basicConfig(level=logging.DEBUG)
	
	for nombre_archivo in args.files:
		if '-' in nombre_archivo: return
		archivo = open(nombre_archivo, encoding="latin_1")
		lineas_originales = list(archivo.readlines())
		archivo.close()
		
		if 'actor' in nombre_archivo or 'actresses' in nombre_archivo:
			lista_de_lineas = actors.actorstocsv(lineas_originales, flags)
		
		if args.test: 
			lista_de_lineas = generate_test(lista_de_lineas, int(args.test), flags.get('KevinBacon'))
			nombre_archivo = 'test.list'
		
		exportar(nombre_archivo, lista_de_lineas)
		
def generate_test(lista_de_lineas, K, kb):
	if kb: linea_bacon = lista_de_lineas[0]
	
	# Reservoir sampling, sacado de https://stackoverflow.com/questions/2612648/reservoir-sampling
	result = []
	N = 0
	for item in lista_de_lineas:
		N += 1
		if len(result) < K:
			result.append( item )
		else:
			s = int(random.random() * N)
			if s < K:
				result[s] = item
	
	if linea_bacon: result[0] = linea_bacon
	logging.info("El archivo generado sera de prueba y contendra {} lineas".format(K))
	return result
		
def exportar(nombre_archivo,lista_de_lineas):
	logging.debug("Comenzando la exportacion de lista de lineas a archivo csv")
	nombre_csv = nombre_archivo.rsplit(".",1)[0]+'.csv'
	with open(nombre_csv,"w") as f:
		wr = csv.writer(f, lineterminator='\n')
		wr.writerows(lista_de_lineas)
	logging.info("Generado {}!".format(nombre_csv))

main()