make QMailAuthuser a context manager

1 files changed, 73 insertions, 70 deletions

diff --git a/src/jwebmail/webmail.py b/src/jwebmail/webmail.py
index d183b30..d260e8a 100644
--- a/src/jwebmail/webmail.py
+++ b/src/jwebmail/webmail.py
@@ -86,31 +86,32 @@ def about():
 
 
 def displayheaders(folder=""):
-    folders = get_read_mails_logged_in().folders()
-
-    if folder and folder not in folders:
-        return render_template("error", error="no_folder", links=folders), 404
-
-    page_bound = request.args.get("page_bound")
-    page_after = bool(request.args.get("page_after", type=int, default=True))
-    per_page = request.args.get("per_page", type=int, default=25)
-    sort = request.args.get("sort", "!date")
-    search = request.args.get("search")
-
-    s = sort[1:] if sort[0] == "!" else sort
-    if s not in ["date", "size", "sender"]:
-        abort(400)
-
-    count = get_read_mails_logged_in().count(folder)
-
-    headers, first, last = get_read_mails_logged_in().list_search(
-        folder=folder,
-        bound=page_bound,
-        after=page_after,
-        limit=per_page,
-        sort=sort,
-        search=search,
-    )
+    with get_read_mails_logged_in() as read_mails:
+        folders = read_mails.folders()
+
+        if folder and folder not in folders:
+            return render_template("error", error="no_folder", links=folders), 404
+
+        page_bound = request.args.get("page_bound")
+        page_after = bool(request.args.get("page_after", type=int, default=True))
+        per_page = request.args.get("per_page", type=int, default=25)
+        sort = request.args.get("sort", "!date")
+        search = request.args.get("search")
+
+        s = sort[1:] if sort[0] == "!" else sort
+        if s not in ["date", "size", "sender"]:
+            abort(400)
+
+        count = read_mails.count(folder)
+
+        headers, first, last = read_mails.list_search(
+            folder=folder,
+            bound=page_bound,
+            after=page_after,
+            limit=per_page,
+            sort=sort,
+            search=search,
+        )
 
     if headers:
         match s:
@@ -160,41 +161,41 @@ def displayheaders(folder=""):
 def readmail(msgid, folder=""):
     format = request.args.get("format", "html").lower()
 
-    if format == "html":
-        try:
-            mail = get_read_mails_logged_in().show(folder, msgid)
-        except QMAuthError:
-            return render_template("not_found.html"), 404
+    with get_read_mails_logged_in() as read_mails:
+        if format == "html":
+            try:
+                mail = read_mails.show(folder, msgid)
+            except QMAuthError:
+                return render_template("not_found.html"), 404
 
-        return render_template("readmail.html", msg=mail, folder=folder)
+            return render_template("readmail.html", msg=mail, folder=folder)
 
-    elif format == "raw":
-        path = request.args.get("path", "")
+        elif format == "raw":
+            path = request.args.get("path", "")
 
-        content = get_read_mails_logged_in().raw(folder, msgid, path)
+            content = read_mails.raw(folder, msgid, path)
 
-        headers = []
+            headers = []
 
-        cd = content["head"].get("content_disposition")
-        if cd and cd.lower() == "attachment":
-            headers.append(
-                (
-                    "Content-Disposition",
-                    f"attachment; filename={content['head']['filename']}",
+            cd = content["head"].get("content_disposition")
+            if cd and cd.lower() == "attachment":
+                headers.append(
+                    (
+                        "Content-Disposition",
+                        f"attachment; filename={content['head']['filename']}",
+                    )
                 )
-            )
-        ct = to_mime_type(content["head"])
-        if ct.startswith("text/"):
-            ct += "; charset=UTF-8"
-        headers.append(("Content-Type", ct))
+            ct = to_mime_type(content["head"])
+            if ct.startswith("text/"):
+                ct += "; charset=UTF-8"
+            headers.append(("Content-Type", ct))
 
-        return content["body"], headers
+            return content["body"], headers
 
-    elif format == "json":
-        mail = get_read_mails_logged_in().show(folder, msgid)
-        return mail
-    else:
-        abort(404)
+        elif format == "json":
+            return read_mails.show(folder, msgid)
+        else:
+            abort(404)
 
 
 def writemail():
@@ -208,16 +209,17 @@ def _take_common_req_args(mapping):
 
 
 def move(folder=""):
-    folders = get_read_mails_logged_in().folders()
+    with get_read_mails_logged_in() as read_mails:
+        folders = read_mails.folders()
 
-    mm = request.form.getlist("mail")
-    to_folder = request.form["select-folder"]
+        mm = request.form.getlist("mail")
+        to_folder = request.form["select-folder"]
 
-    if folder not in folders or to_folder not in folders:
-        raise ValueError("folder not valid")
+        if folder not in folders or to_folder not in folders:
+            raise ValueError("folder not valid")
 
-    for m in mm:
-        get_read_mails_logged_in().move(m, folder, to_folder)
+        for m in mm:
+            read_mails.move(m, folder, to_folder)
 
     flash(gettext("succ_move"))
     args = _take_common_req_args(request.form)
@@ -225,21 +227,22 @@ def move(folder=""):
 
 
 def remove(folder=""):
-    folders = get_read_mails_logged_in().add_folder("Trash")
+    with get_read_mails_logged_in() as read_mails:
+        folders = read_mails.add_folder("Trash")
 
-    mm = request.form.getlist("mail")
+        mm = request.form.getlist("mail")
 
-    folders = get_read_mails_logged_in().folders()
+        folders = read_mails.folders()
 
-    if folder not in folders:
-        raise ValueError("folder not valid")
+        if folder not in folders:
+            raise ValueError("folder not valid")
 
-    if folder == "Trash":
-        for m in mm:
-            get_read_mails_logged_in().remove(folder, m)
-    else:
-        for m in mm:
-            get_read_mails_logged_in().move(m, folder, "Trash")
+        if folder == "Trash":
+            for m in mm:
+                read_mails.remove(folder, m)
+        else:
+            for m in mm:
+                read_mails.move(m, folder, "Trash")
 
     flash(gettext("succ_remove"))
     args = _take_common_req_args(request.form)