diff --git a/novel_stats/novel_stats.py b/novel_stats/novel_stats.py index d63f779..3e08f85 100755 --- a/novel_stats/novel_stats.py +++ b/novel_stats/novel_stats.py @@ -28,53 +28,65 @@ def main(): arguments = sys.argv[1:] filename = arguments[0] chapter_number = None - chapter_word_count = 0 - chapter_status = None act_number = None - act_word_count = 0 total_word_count = 0 + word_count_by_chapter = collections.defaultdict(int) word_count_by_status = collections.defaultdict(int) + word_count_by_act = collections.defaultdict(int) + status_by_chapter = {} for line in open(filename).readlines(): if line.startswith(CHAPTER_MARKER): - if chapter_number: - print( - 'chapter {}: {} words{}'.format( - chapter_number, - chapter_word_count, - ' ({})'.format(chapter_status) if chapter_status else '' - ) - ) + word_count_by_act[act_number] += word_count_by_chapter[chapter_number] + total_word_count += word_count_by_chapter[chapter_number] + if chapter_number in status_by_chapter: + word_count_by_status[status_by_chapter[chapter_number]] += 1 + chapter_number = int(line[len(CHAPTER_MARKER):]) - act_word_count += chapter_word_count - total_word_count += chapter_word_count - chapter_word_count = 1 # Start at one, because the chapter number itself counts as a word. - if chapter_status: + word_count_by_chapter[chapter_number] = 1 # Start at one, because the chapter number itself counts as a word. + if chapter_number in status_by_chapter: word_count_by_status[chapter_status] += 1 - - chapter_status = None # Clear out the status from the previous chapter. elif line.startswith(STATUS_MARKER): - chapter_status = line[len(STATUS_MARKER):].strip('()\n') + status_by_chapter[chapter_number] = line[len(STATUS_MARKER):].strip('()\n') elif line.startswith(ACT_MARKER): - if act_number: - print('act {}: {} words'.format(act_number, act_word_count)) act_number = int(line[len(ACT_MARKER):].strip('()\n')) - act_word_count = 1 + word_count_by_act[act_number] = 1 else: line_word_count = count_words(line) - chapter_word_count += line_word_count + word_count_by_chapter[chapter_number] += line_word_count - if chapter_status: - word_count_by_status[chapter_status] += line_word_count + if chapter_number in status_by_chapter: + word_count_by_status[status_by_chapter[chapter_number]] += line_word_count - if chapter_status: - word_count_by_status[chapter_status] += 1 + # Do some final accounting after the last chapter. + word_count_by_act[act_number] += word_count_by_chapter[chapter_number] + total_word_count += word_count_by_chapter[chapter_number] + if chapter_number in status_by_chapter: + word_count_by_status[status_by_chapter[chapter_number]] += 1 + + # Print out word counts. + for chapter_number, chapter_word_count in word_count_by_chapter.items(): + if chapter_number is None: + continue + + chapter_status = status_by_chapter.get(chapter_number) + + print( + 'chapter {}: {} words{}'.format( + chapter_number, + chapter_word_count, + ' ({})'.format(chapter_status) if chapter_status else '', + ) + ) - print('chapter {}: {} words'.format(chapter_number, chapter_word_count)) - print('act {}: {} words'.format(act_number, act_word_count)) print() - total_word_count += chapter_word_count + + for act_number, act_word_count in word_count_by_act.items(): + if act_number is None: + continue + + print('act {}: {} words'.format(act_number, act_word_count)) for status, status_word_count in word_count_by_status.items(): print('{}: {} words'.format(status, status_word_count))