Refactored recommendation code into a new file

ucsinfer/__main__.py

@@ -1,7 +1,6 @@
 import os
 import sys
 import csv
-from re import match
 
 from sentence_transformers import SentenceTransformer
 import tqdm
@@ -9,67 +8,9 @@ import click
 from tabulate import tabulate, SEPARATING_LINE
 
 from .inference import InferenceContext, load_ucs
+from .recommend import print_recommendation
 from .util import ffmpeg_description, parse_ucs
 
-
-def recommend_text(text: str, ctx: InferenceContext):
-    return ctx.classify_text_ranked(text)
-
-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:
-        print(f"Path: {path}")
-
-    print(f"Text: {text or '<None>'}")
-    for i, r in enumerate(recs):
-        cat, subcat, _ = ctx.lookup_category(r)
-        print(f"- {i}: {r} ({cat}-{subcat})")
-
-    if interactive_rename and path is not None:
-        response = input("#, t [text], ?, q > ")
-        
-        if m := match(r'^([0-9]+)', response):
-            selection = int(m.group(1))
-            if 0 <= selection < len(recs):
-                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)
-            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 "
              "<https://git.squad51.us/jamie/ucsinfer>")
 # @click.option('--verbose', flag_value='verbose', help='Verbose output')

ucsinfer/recommend.py

@@ -0,0 +1,62 @@
+# recommend.py
+
+from re import match
+
+from .inference import 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:
+        print(f"Path: {path}")
+
+    print(f"Text: {text or '<None>'}")
+    for i, r in enumerate(recs):
+        cat, subcat, _ = ctx.lookup_category(r)
+        print(f"- {i}: {r} ({cat}-{subcat})")
+
+    if interactive_rename and path is not None:
+        response = input("#, t [text], ?, q > ")
+        
+        if m := match(r'^([0-9]+)', response):
+            selection = int(m.group(1))
+            if 0 <= selection < len(recs):
+                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)
+            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
+
+