list-stats

#!/usr/bin/env python

"""
Calculate and report some statistics for each data source JSON file.
"""

import argparse
import csv
import gzip
import json
import re

import movielib

class IMDB(object):
    titles = {}
    def __init__(self):
        path = 'title.basics.tsv.gz'
        with gzip.open(path, 'rb') as f:
            reader = csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE)
            for row in reader:
                info = {
                    'title': row[3], # Original title
                    'year': row[5],
                }
                if '\\N' != row[7]:
                    info['duration'] = row[7] # minutes
                self.titles[row[0]] = info
    def __getitem__(self, key):
        if key in self.titles:
            return self.titles[key]
        else:
            return None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('files', metavar='N', type=str, nargs='*',
                        default=['free-movies-manual.json'],
                        help='input file(s)')
    parser.add_argument('--list-unique', action='store_true', default=False)
    parser.add_argument('--list-all', action='store_true', default=False)
    parser.add_argument('--create-complete', action='store_true', default=False)
    args = parser.parse_args()

    if args.create_complete:
        ids = IMDB()
    else:
        ids = None

    complete = {}
    s = {}
    i = {}

    for file in args.files:
        with open(file, 'rt') as input:
            l = json.load(input)
            imdb = 0
            noimdb = 0
            for e in l:
                if not dict is type(l[e]): # Ignore comment entries
                    continue
                if args.list_all or \
                   ('status' in l[e] and "free" == l[e]['status']):
                   m = re.match('https?://www.imdb.com/title/(.+)/', e)
                   if m:
                      imdb = imdb + 1
                      imdbid = m.group(1)
                      if e not in complete:
                          complete[e] = {
                              'file': [],
                              'freenessurls': [],
                              'status': 'free',
                          }
                          if ids:
                              imdbinfo = ids[imdbid]
                          else:
                              imdbinfo = None
                          if imdbinfo:
                              complete[e].update(imdbinfo)
                      complete[e]['file'].append(file)
                      # Keep list sorted to get predictable output
                      complete[e]['file'] = \
                        sorted(complete[e]['file'])
                      for key in ('archive', 'archive1', 'archive2',
                                  'archive3', 'archive4', 'archive5',
                                  'freenessurl', 'freenessurl1', 'freenessurl2',
                                  'freenessurl3', 'freenessurl4', 'freenessurl5',
                                  'freenessurl6', 'freenessurl7', 'freenessurl8',
                                  'freenessurl9', 'freenessurl10', 'freenessurl11',
                                  'freenessurl12', 'freenessurl13', 'freenessurl14',
                                  'freenessurl15', 'freenessurl16'):
                          if key in l[e] and 'n/a' != l[e][key]:
                              url = l[e][key]
                              if url not in complete[e]['freenessurls']:
                                  complete[e]['freenessurls'].append(url)
                                  # Keep list sorted to get predictable output
                                  complete[e]['freenessurls'] = \
                                    sorted(complete[e]['freenessurls'])
                   else:
                       noimdb = noimdb + 1
                   s[file] = {
                       'imdb' : imdb,
                       'noimdb' : noimdb,
                   }
                   i[e] = l[e]

    usum = 0
    xsum = 0
    for file in sorted(s.keys()):
        unique = 0
        for imdburl in complete.keys():
            if complete[imdburl]['file'] == [file]:
                unique = unique + 1
        print("%5d entries (%5d unique) with and %5d without IMDB title ID in %s"
              % (s[file]['imdb'], unique, s[file]['noimdb'], file))
        usum = usum + unique
        xsum = xsum + s[file]['noimdb']
        if args.list_unique:
            def sortbyyear(imdburl):
                if 'year' in i[imdburl]:
                    return int(i[imdburl]['year']), imdburl
                else:
                    return imdburl
            for imdburl in sorted(complete, key=sortbyyear):
                if complete[imdburl]['file'] == [file]:
                    if 'year' in i[imdburl]:
                        y = i[imdburl]['year']
                    else:
                        y = ""
                    if 'title' in i[imdburl]:
                        t = i[imdburl]['title']
                    else:
                        t = ""
                    msg = "        %s %4s %s" % (imdburl, y, t)
                    print(msg.encode('utf-8'))

    print("%5d unique IMDB title IDs in total, %d only in one list, %d without IMDB title ID" %
          (len(complete.keys()), usum, xsum))
    if args.create_complete:
        movielib.savelist(complete, 'free-complete.json')

if __name__ == '__main__':
    main()