From 8460879f7333725e94964ac29cfb71ee234595ba Mon Sep 17 00:00:00 2001 From: Dan Helfman Date: Sat, 4 Nov 2023 21:39:49 -0700 Subject: [PATCH] Add support for chapter ranges. --- format_novel/format_novel.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/format_novel/format_novel.py b/format_novel/format_novel.py index ba21f68..69ebf41 100755 --- a/format_novel/format_novel.py +++ b/format_novel/format_novel.py @@ -6,17 +6,32 @@ import sys import docx -def filter_chapters(document, chapter_numbers): - if not chapter_numbers: +def expand_chapter_range(chapter_request): + if not '-' in chapter_request: + return (int(chapter_request),) + + chapter_range = chapter_request.split('-') + + return tuple(range(int(chapter_range[0]), int(chapter_range[1]) + 1)) + + +def filter_chapters(document, chapter_requests): + if not chapter_requests: return - chapter_number = 0 + chapter_numbers = { + chapter_number + for request in chapter_requests + for chapter_number in expand_chapter_range(request) + } + + current_chapter_number = 0 for paragraph in document.paragraphs: if paragraph.style.name == 'Heading 2': - chapter_number += 1 + current_chapter_number += 1 - if chapter_number not in chapter_numbers: + if current_chapter_number not in chapter_numbers: element = paragraph._element element.getparent().remove(element) @@ -109,8 +124,7 @@ def parse_arguments(unparsed_arguments): '--chapters', metavar='CHAPTER', nargs='+', - type=int, - help='Numbers of chapters to include in output', + help='Chapter numbers or ranges to include in output', ) return parser.parse_args(unparsed_arguments)