diff --git a/pdfca.py b/pdfca.py index 847c26e..92646ae 100644 --- a/pdfca.py +++ b/pdfca.py @@ -12,6 +12,15 @@ import PyPDF2 +def checkIndex(dataframe): + key = click.prompt('Please enter item to explore') + while True: + if key in dataframe.index: + return key + else: + key = click.prompt('Please enter a valid key') + + def count(x, term): search = re.findall(term.strip().lower(), x.strip().lower()) num = len(search) @@ -235,9 +244,10 @@ def search(term, binary, form, search_type, number): else: results = dropped.sum().sort_values(by=[term]) click.echo(f'Top {number} results:') - click.echo(results.tail(number)) + results = results.tail(number) + click.echo(results) if click.confirm('\nWould you like to drill down?'): - key = click.prompt('Please enter item to explore') + key = checkIndex(results) filtered = grouped.get_group(key).sort_values(by=[term]) filtered = filtered.set_index('page') click.echo()