Extract note file getting into function

papis-marvin

@@ -6,6 +6,7 @@
 # annotations as 'csv' format and then point the script to the
 # resulting file.
 import sys
+from typing import List
 import papis.api
 import papis.commands.list
 import papis.database
@@ -26,29 +27,9 @@ def main(fpath, db):
 
             # switch to next book
             if not is_same_book(row["Title"]):
-                documents = db.query(
-                    f"author:({row['Author']}) title:({title_stripped})"
-                )
-                if not documents:
-                    print(
-                        f"No papis entry found for Marvin entry - {row['Author']}: {row['Title']}.\nPlease manually create."
-                    )
-                    continue
-                note_file = papis.commands.list.run(notes=True, documents=documents)
+                note_file = get_notefile(row["Author"], title_stripped)
                 if not note_file:
-                    print(
-                        f"Found reference entry but no note file for - {row['Author']}: {row['Title']}."
-                    )
-                    if input(f"Create note file now? [y/N] ") == "y":
-                        output = subprocess.run(
-                            ["papis", "edit", "-n"], capture_output=True, shell=True
-                        )
-                        print(output)
-                        # print("NOT IMPLEMENTED: Please create note file manually.")
-                        continue
-                    else:
-                        continue
-                note_file = str(note_file[0])
+                    continue
 
             text = format_entry(row)
 
@@ -59,6 +40,34 @@ def main(fpath, db):
 
     write_to_files(notes)
 
+
+def get_documents(author, title) -> List:
+    return db.query(f"author:({author}) title:({title})")
+
+
+def get_notefile(author, title) -> str | None:
+    documents = get_documents(author, title)
+    if not documents:
+        print(
+            f"No papis entry found for Marvin entry - {author}: {title}.\nPlease manually create."
+        )
+        return
+    note_file = papis.commands.list.run(notes=True, documents=documents)
+    if not note_file:
+        print(f"Found reference entry but no note file for - {author}: {title}.")
+        if input(f"Create note file now? [y/N] ") == "y":
+            output = subprocess.run(
+                ["papis", "edit", "-n"], capture_output=True, shell=True
+            )
+            print(output)
+            print("NOT IMPLEMENTED: Please create note file manually.")
+            return
+        else:
+            return
+    return str(note_file[0])
+
+
+# TODO implement custom formatting (akin to pubs-extract)
 def format_entry(row) -> str:
     text = f"> {row['HighlightText']}"
     if row["EntryText"]:
@@ -71,6 +80,8 @@ def format_entry(row) -> str:
 
 
 _old_title = ""
+
+
 def is_same_book(title):
     global _old_title
 
@@ -95,7 +106,10 @@ def write_to_files(notes):
                             num_added += 1
                 print(f"Added {num_added} entries to it.")
 
+
 title_strip_pattern = re.compile(r"([^\s\w]|_)+")
+
+
 def strip_book_title(title) -> str:
     return title_strip_pattern.sub("", title)