superx-budget-overview/superx_budget/pyramid/templates.py

""" Views for the templates part """

from pathlib import Path

from pyramid.view import view_config
from pyramid.response import FileResponse
from pyramid.httpexceptions import HTTPFound

from . import XLSX_CONTENT_TYPE, Root
from ..budget import parse_budget_file
from ..helpers import list_budget_files, is_budget_file_name
from ..overview import create_overview  # noqa: F401
from ..exceptions import BudgetParserError, SuperXParserError  # noqa: F401


@view_config(
    context=Root,
    name="templates",
    request_method="GET",
    renderer="superx_budget:pyramid/templates/templates.jinja2",
    permission="view",
)
def templates(context, request):
    if "f" in request.GET:
        file_name = request.GET["f"]
        file_path = request.budgets_dir / file_name
        if not file_path.is_file():
            return HTTPFound("/templates")
        response = FileResponse(
            file_path,
            request=request,
            cache_max_age=0,
            content_type=XLSX_CONTENT_TYPE,
        )
        response.headers[
            "Content-Disposition"
        ] = f"attachment;filename={file_name}"
        return response

    tmp = list_budget_files(request.budgets_dir)
    tmp = sorted(tmp, key=lambda p: p.stem[-4:])
    budget_files = [path.name for path in tmp]
    return {"budget_files": budget_files}


@view_config(
    context=Root,
    name="templates",
    request_method="POST",
    renderer="superx_budget:pyramid/templates/templates.jinja2",
    permission="view",
)
def templates_update(context, request):
    upload = request.POST.get("budget")

    if upload == b"" or not is_budget_file_name(upload.filename):
        request.session.flash("No Excel file uploaded.", "error")
        return HTTPFound("/templates")

    try:
        parse_budget_file(upload.file)
    except BudgetParserError:
        request.session.flash(
            "File does not appear to be in the required format.", "error"
        )
        return HTTPFound("/templates")

    # sanitizing upload filename:
    file_name = Path(upload.filename).name

    tmp = list_budget_files(request.budgets_dir)
    available_files = {p.name.lower(): p for p in tmp}
    old_path = available_files.get(file_name.lower())
    if old_path:
        old_path.unlink()

    new_path = request.budgets_dir / file_name
    with new_path.open("wb") as fh:
        upload.file.seek(0)
        fh.write(upload.file.read())

    request.session.flash("File upload successful.", "info")

    return HTTPFound("/templates")