import os import json from . import common from . import tsdb def breakdown_database(subreddit=None, username=None): ''' Given a database, return a json dict breaking down the submission / comment count for users (if a subreddit database) or subreddits (if a user database). ''' if (subreddit is None) == (username is None): raise Exception('Enter subreddit or username but not both') breakdown_results = {} def _ingest(names, subkey): for name in names: breakdown_results.setdefault(name, {}) breakdown_results[name].setdefault(subkey, 0) breakdown_results[name][subkey] += 1 if subreddit: database = tsdb.TSDB.for_subreddit(subreddit, do_create=False) else: database = tsdb.TSDB.for_user(username, do_create=False) cur = database.sql.cursor() for table in ['submissions', 'comments']: if subreddit: cur.execute('SELECT author FROM %s' % table) elif username: cur.execute('SELECT subreddit FROM %s' % table) names = (row[0] for row in common.fetchgenerator(cur)) _ingest(names, table) for name in breakdown_results: breakdown_results[name].setdefault('submissions', 0) breakdown_results[name].setdefault('comments', 0) return breakdown_results def breakdown_argparse(args): if args.subreddit: database = tsdb.TSDB.for_subreddit(args.subreddit, do_create=False) else: database = tsdb.TSDB.for_user(args.username, do_create=False) breakdown_results = breakdown_database( subreddit=args.subreddit, username=args.username, ) def sort_name(name): return name.lower() def sort_submissions(name): invert_score = -1 * breakdown_results[name]['submissions'] return (invert_score, name.lower()) def sort_comments(name): invert_score = -1 * breakdown_results[name]['comments'] return (invert_score, name.lower()) def sort_total_posts(name): invert_score = breakdown_results[name]['submissions'] + breakdown_results[name]['comments'] invert_score = -1 * invert_score return (invert_score, name.lower()) breakdown_sorters = { 'name': sort_name, 'submissions': sort_submissions, 'comments': sort_comments, 'total_posts': sort_total_posts, } breakdown_names = list(breakdown_results.keys()) if args.sort is not None: try: sorter = breakdown_sorters[args.sort.lower()] except KeyError: message = '{sorter} is not a sorter. Choose from {options}' message = message.format(sorter=args.sort, options=list(breakdown_sorters.keys())) raise KeyError(message) breakdown_names.sort(key=sorter) dump = ' "{name}": {{"submissions": {submissions}, "comments": {comments}}}' dump = [dump.format(name=name, **breakdown_results[name]) for name in breakdown_names] dump = ',\n'.join(dump) dump = '{\n' + dump + '\n}\n' else: dump = json.dumps(breakdown_results) if args.sort is None: breakdown_basename = '%s_breakdown.json' else: breakdown_basename = '%%s_breakdown_%s.json' % args.sort breakdown_basename = breakdown_basename % database.filepath.replace_extension('').basename breakdown_filepath = database.breakdown_dir.with_child(breakdown_basename) breakdown_filepath.parent.makedirs(exist_ok=True) breakdown_file = breakdown_filepath.open('w') with breakdown_file: breakdown_file.write(dump) print('Wrote', breakdown_filepath.relative_path) return breakdown_results