diff --git a/ucsinfer/__main__.py b/ucsinfer/__main__.py index 1d62d46..c2a5eb7 100644 --- a/ucsinfer/__main__.py +++ b/ucsinfer/__main__.py @@ -17,6 +17,17 @@ def recommend_text(text: str, ctx: InferenceContext): def print_recommendation(path: str | None, text: str, ctx: InferenceContext, interactive_rename: bool): + """ + Print recommendations interactively. + + :returns: If interactive_rename is false or path is None, returns None. + If interactive, returns a tuple of bool, str | None, str | None where: + - if retval[0] is False, the user has requested processing quit. + - if retval[1] is a str, this is the text the user has entered to + perform a new inference instead of the file metadata. + `print_recommendation` should be called again with this argument. + - if retval[2] is a str, this is the catid the user has selected. + """ recs = ctx.classify_text_ranked(text) print("----------") if path: @@ -28,20 +39,36 @@ def print_recommendation(path: str | None, text: str, ctx: InferenceContext, print(f"- {i}: {r} ({cat}-{subcat})") if interactive_rename and path is not None: - response = input("Enter number, t for alternate text, or " - "return to skip: ") + response = input("#, t [text], ?, q > ") if m := match(r'^([0-9]+)', response): selection = int(m.group(1)) if 0 <= selection < len(recs): - new_name = recs[selection] + '_' + os.path.basename(path) - new_path = os.path.join(os.path.dirname(path), new_name) - os.rename(path, new_path) + return True, None, recs[selection] + else: + print(f"Invalid index {selection}") + return True, text, None + elif m := match(r'^t (.*)', response): print("searching for new matches") text = m.group(1) - print_recommendation(path, text, ctx, True) + return True, text, None + elif response.startswith("?"): + print(""" +Choices: +- Enter recommendation number to rename file, +- "t [text]" to search for new recommendations based on [text] +- "p" re-use the last selected cat-id +- "?" for this message +- "q" to quit +- or any other key to skip this file and continue to next file +""") + return True, text, None + elif response.startswith('q'): + return (False, None, None) + else: + return None @click.group(epilog="For more information see " "") @@ -84,13 +111,25 @@ def recommend(text, paths, model, interactive): for path in paths: text = ffmpeg_description(path) - if text: - print_recommendation(path, text, ctx, interactive) + if not text: + text = os.path.basename(path) - else: - filename = os.path.basename(path) - print_recommendation(path, filename, ctx, interactive) - + while True: + retval = print_recommendation(path, text, ctx, interactive) + if not retval: + break + if retval[0] is False: + return + elif retval[1] is not None: + text = retval[1] + continue + elif retval[2] is not None: + new_name = retval[2] + '_' + os.path.basename(path) + new_path = os.path.join(os.path.dirname(path), new_name) + print(f"Renaming {path} to {new_path}") + os.rename(path, new_path) + break + @ucsinfer.command('gather') @click.option('--outfile', type=click.File(mode='w', encoding='utf8'),