Minor UX improvements and cleanup (#38)

* Enhance error handling and user interruption responses across multiple scripts

* Refactor error handling tests to improve connection simulation and patching

* Remove comments about exception propagation in main functions across multiple scripts

* Update versioning and add version argument to scripts for better tracking

.github/workflows/publish.yml

@@ -15,6 +15,14 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: "3.x"
+    
+    - name: Update version to match release
+      run: |
+        VERSION=${GITHUB_REF_NAME#v}
+        echo "Updating version to $VERSION"
+        sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml
+        grep "^version =" pyproject.toml
+        
     - name: Install pypa/build
       run: >-
         python3 -m pip install build --user

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imap-migration-tools"
-version = "1.0.0"
+version = "1.1.0"
 authors = [
   { name="Javier Callico", email="jcallico@callicode.com" },
   { name="Nathan Moinvaziri", email="nathan@nathanm.com" }

src/backup_imap_emails.py

@@ -681,6 +681,7 @@ def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=
 
 def main():
     parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
 
     # Source
     default_src_host = os.getenv("SRC_IMAP_HOST")
@@ -885,11 +886,15 @@ def main():
             print("Use this file when restoring to reapply read/starred status to emails.")
 
     except KeyboardInterrupt:
-        print("\nBackup interrupted by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
+        raise
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/compare_imap_folders.py

@@ -59,6 +59,7 @@ Examples:
 
 import argparse
 import os
+import sys
 
 import imap_common
 import imap_oauth2
@@ -99,6 +100,7 @@ def main():
 
     # Phase 2: full parser with conditional requirements
     parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
     parser.add_argument(
         "--src-path",
         default=default_src_path,
@@ -304,26 +306,32 @@ def main():
         print("-" * len(header))
         print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
 
-    except Exception as e:
-        print(f"\nFatal Error: {e}")
-        if "Too many simultaneous connections" in str(e):
-            print("Tip: Wait a few minutes for previous connections to timeout or check other active scripts.")
+    except KeyboardInterrupt:
+        # Re-raise to be handled by the outer block, but ensure finally runs
+        raise
 
     finally:
         # Check source connection state and logout if possible
         if src:
             try:
                 src.logout()
-            except Exception:
+            except BaseException:
                 pass
 
         # Check dest connection state and logout if possible
         if dest:
             try:
                 dest.logout()
-            except Exception:
+            except BaseException:
                 pass
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/count_imap_emails.py

@@ -144,6 +144,7 @@ def main(argv: Optional[list[str]] = None) -> None:
 
     # Phase 2: full parser with conditional requirements
     parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
 
     parser.add_argument(
         "--path",
@@ -236,4 +237,11 @@ def main(argv: Optional[list[str]] = None) -> None:
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
+        sys.exit(0)
+    except Exception as e:
+        print(f"Fatal Error: {e}")
+        sys.exit(1)

src/imap_common.py

@@ -16,12 +16,42 @@ import urllib.parse
 from email import policy
 from email.header import decode_header
 from email.parser import BytesParser
+from importlib.metadata import PackageNotFoundError, version
 
 import imap_compress
 import imap_oauth2
 import imap_retry
 import restore_cache
 
+
+def get_version() -> str:
+    """
+    Return the package version string.
+
+    Tries to get the version from the installed package metadata.
+    If not installed, falls back to reading pyproject.toml from the project root.
+    """
+    try:
+        return version("imap-migration-tools")
+    except PackageNotFoundError:
+        # Fallback: Try to read from pyproject.toml for local development
+        try:
+            # Assuming src/imap_common.py structure
+            root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            pyproject_path = os.path.join(root_dir, "pyproject.toml")
+
+            if os.path.exists(pyproject_path):
+                with open(pyproject_path, encoding="utf-8") as f:
+                    for line in f:
+                        # Looking for: version = "1.2.3"
+                        match = re.search(r'^version\s*=\s*"([^"]+)"', line.strip())
+                        if match:
+                            return match.group(1)
+        except Exception:
+            pass
+        return "0.0.0-dev"
+
+
 # Standard IMAP flags
 FLAG_SEEN = "\\Seen"
 FLAG_ANSWERED = "\\Answered"

src/migrate_imap_emails.py

@@ -758,6 +758,7 @@ def migrate_folder(
 
 def main():
     parser = argparse.ArgumentParser(description="Migrate emails between IMAP accounts.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
 
     # Positional arg for folder (optional) to keep backward compatibility with previous quick-fix
     parser.add_argument("folder", nargs="?", help="Specific folder to migrate (e.g. '[Gmail]/Important')")
@@ -1102,12 +1103,16 @@ def main():
         src_main.logout()
         dest_main.logout()
 
+    except KeyboardInterrupt:
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
     except KeyboardInterrupt:
         safe_print("\n\nProcess terminated by user.")
         sys.exit(0)
     except Exception as e:
         safe_print(f"Fatal Error: {e}")
-
-
-if __name__ == "__main__":
-    main()
+        sys.exit(1)

src/restore_imap_emails.py

@@ -693,6 +693,7 @@ def restore_gmail_with_labels(
 
 def main():
     parser = argparse.ArgumentParser(description="Restore IMAP emails from local .eml files.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
 
     # Source (Local Path)
     env_path = os.getenv("BACKUP_LOCAL_PATH")
@@ -959,7 +960,14 @@ def main():
         print("\nRestore completed successfully.")
 
     except KeyboardInterrupt:
-        print("\nRestore interrupted by user.")
+        raise
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nProcess terminated by user.")
         sys.exit(0)
     except Exception as e:
         print(f"Fatal Error: {e}")
@@ -967,7 +975,3 @@ def main():
 
         traceback.print_exc()
         sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()

test/test_migrate_imap_emails.py

@@ -357,9 +357,11 @@ class TestMigrateErrorHandling:
 
         from unittest.mock import patch
 
+        original_get_imap_connection = imap_common.get_imap_connection
+
         def side_effect(host, user, pwd, oauth2_token=None):
             if threading.current_thread() is threading.main_thread():
-                return imap_common.get_imap_connection(host, user, pwd, oauth2_token)
+                return original_get_imap_connection(host, user, pwd, oauth2_token)
             return None  # Simulate connection failure in worker
 
         env = _mock_migrate_env(p1, p2)
@@ -442,7 +444,7 @@ class TestMigrateErrorHandling:
 
         env = _mock_migrate_env(p1, p2)
         with (
-            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid),
+            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid, autospec=True),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
         ):