instrumentation/next-share/BaseLib/Category/TestCategory.py

   1 # Written by Yuan Yuan
   2 # see LICENSE.txt for license information
   3
   4 import sys, os
   5 execpath = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), '..', '..')
   6 sys.path.append(execpath)
   7 #print sys.path
   8 from Utility.utility import getMetainfo
   9 from BaseLib.Category.Category import Category
  10
  11 DEBUG = False
  12
  13 def testFilter(catfilename, torrentpath):
  14     readCategorisationFile(catfilename)
  15     #print 'Install_dir is %s' % execpath
  16     c = Category.getInstance(execpath, None)
  17     total = porn = fn = fp = 0
  18     for tfilename,isporn in tdict.items():
  19         torrent = getMetainfo(os.path.join(torrentpath,tfilename))
  20         name = torrent['info']['name']
  21         cat = c.calculateCategory(torrent, name)
  22         fporn = (cat == ['xxx'])
  23         total+= 1
  24         porn += int(isporn)
  25         if isporn == fporn:
  26             if DEBUG:
  27                 print (isporn, fporn), 'good', name
  28
  29         elif isporn and not fporn:
  30             fn+=1
  31             print 'FALSE NEGATIVE'
  32             showTorrent(os.path.join(torrentpath,tfilename))
  33         elif not isporn and fporn:
  34             fp +=1
  35             print 'FALSE POSITIVE'
  36             showTorrent(os.path.join(torrentpath,tfilename))
  37
  38     print """
  39     Total torrents:   %(total)d
  40     XXX torrents:     %(porn)d
  41     Correct filtered: %(good)d
  42     False negatives:  %(fn)d
  43     False positives:  %(fp)d
  44     """ % {'total':total, 'porn':porn, 'fn':fn,'fp':fp,'good':total-fn-fp}
  45
  46 def readCategorisationFile(filename):
  47     global tdict
  48     tdict = {}
  49     try:
  50         f = file(filename, 'r')
  51         lines = f.read().splitlines()
  52         for line in lines:
  53             if line:
  54                 parts = line.split('\t')
  55                 tdict[parts[0]] = bool(int(parts[1]))
  56         f.close()
  57     except IOError:
  58         print 'No file %s found, starting with empty file' % filename
  59
  60 def getTorrentData(path, max_num=-1):
  61     torrents= []
  62     i = 0
  63     for fname in os.listdir(path):
  64         if fname.endswith('.torrent'):
  65             torrents.append(os.path.join(path,fname))
  66         if i%1000 == 0 and i:
  67             print 'Loaded: %d torrents' % i
  68         if i == int(max_num):
  69             break
  70         i+=1
  71     print 'Loaded %d torrents' % len(torrents)
  72     return torrents
  73
  74 def showTorrent(path):
  75     torrent = getMetainfo(os.path.join(path))
  76     name = torrent['info']['name']
  77     print '------------------------------'
  78     print '\tfiles  :'
  79     files_list = []
  80     __size_change = 1024
  81     try:
  82         # the multi-files mode
  83         for ifiles in torrent['info']["files"]:
  84             files_list.append((ifiles['path'][-1], ifiles['length'] / float(__size_change)))
  85     except KeyError:
  86         # single mode
  87         files_list.append((torrent['info']["name"],torrent['info']['length'] / float(__size_change)))
  88     for fname, fsize in files_list:
  89         print'\t\t%s\t%d kb' % (fname, fsize)
  90     print 'Torrent name: %s' % name
  91     print '\ttracker:%s' % torrent['announce']
  92     print '------------------------------'
  93
  94 def createTorrentDataSet(filename, torrentpath):
  95     initSaveFile(filename)
  96     f_out = file(filename, 'a')
  97     torrents = getTorrentData(torrentpath)
  98     for torrent in torrents:
  99         if os.path.split(torrent)[-1] in tset: # already done
 100             continue
 101         showTorrent(torrent)
 102         ans = None
 103         while ans not in ['q', 'y','n']:
 104             print 'Is this torrent porn? (y/n/q)'
 105             ans = sys.stdin.readline()[:-1].lower()
 106         if ans == 'q':
 107             break
 108         else:
 109             saveTorrent(f_out, torrent, (ans=='y'))
 110     f_out.close()
 111
 112 def saveTorrent(f_out, torrent, boolean):
 113     if torrent in tset:
 114         return
 115     tfilename = os.path.split(torrent)[-1]
 116     assert tfilename
 117     f_out.write('%s\t%d\n' % (tfilename, int(boolean)))
 118     f_out.flush()
 119     tset.add(torrent)
 120
 121 def initSaveFile(filename):
 122     global tset
 123     tset = set()
 124     try:
 125         f = file(filename, 'r')
 126         lines = f.read().splitlines()
 127         for line in lines:
 128             tset.add(line.split('\t')[0])
 129         f.close()
 130     except IOError:
 131         print 'No file %s found, starting with empty file' % filename
 132
 133
 134
 135 def main(args):
 136     if len(args) != 4 or args[1] not in ['categorise', 'test']:
 137         print 'Usage 1: %s categorise [torrent-dir] [torrent-data-file]' % args[0]
 138         print 'Usage 2: %s test [torrent-dir] [torrent-data-file]' % args[0]
 139         sys.exit(1)
 140     if args[1] == 'categorise':
 141         createTorrentDataSet(args[3], args[2])
 142     elif args[1] == 'test':
 143         testFilter(args[3], args[2])
 144     print 'ready'
 145
 146
 147 if __name__ == '__main__':
 148     main(sys.argv)