refactor: split flask app structure
This commit is contained in:
+215
@@ -0,0 +1,215 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
|
||||
from flask import (
|
||||
Blueprint,
|
||||
abort,
|
||||
current_app,
|
||||
flash,
|
||||
jsonify,
|
||||
redirect,
|
||||
render_template,
|
||||
request,
|
||||
send_file,
|
||||
session,
|
||||
url_for,
|
||||
)
|
||||
from flask_login import current_user, login_required, login_user, logout_user
|
||||
|
||||
from .config import BASE_DIR
|
||||
from .extensions import db
|
||||
from .models import UploadRecord, User
|
||||
from .prediction import PredictionError, run_prediction
|
||||
from .security import new_captcha
|
||||
|
||||
bp = Blueprint("main", __name__)
|
||||
|
||||
|
||||
@bp.route("/")
|
||||
def index():
|
||||
if current_user.is_authenticated:
|
||||
return redirect(url_for("main.home"))
|
||||
return redirect(url_for("main.login"))
|
||||
|
||||
|
||||
@bp.route("/login", methods=["GET", "POST"])
|
||||
def login():
|
||||
if current_user.is_authenticated:
|
||||
return redirect(url_for("main.home"))
|
||||
|
||||
if request.method == "GET":
|
||||
session["captcha"] = new_captcha()
|
||||
return render_template("login.html", mode="login", captcha=session["captcha"])
|
||||
|
||||
username = request.form.get("username", "").strip()
|
||||
password = request.form.get("password", "")
|
||||
captcha_input = request.form.get("captcha", "").strip().upper()
|
||||
|
||||
if captcha_input != session.get("captcha", ""):
|
||||
flash("验证码错误", "error")
|
||||
session["captcha"] = new_captcha()
|
||||
return render_template("login.html", mode="login", captcha=session["captcha"]), 400
|
||||
|
||||
user = User.query.filter_by(username=username).first()
|
||||
if not user or not user.check_password(password):
|
||||
flash("用户名或密码错误", "error")
|
||||
session["captcha"] = new_captcha()
|
||||
return render_template("login.html", mode="login", captcha=session["captcha"]), 400
|
||||
|
||||
login_user(user, remember=bool(request.form.get("remember")))
|
||||
return redirect(url_for("main.home"))
|
||||
|
||||
|
||||
@bp.route("/register", methods=["GET", "POST"])
|
||||
def register():
|
||||
if not current_app.config["ALLOW_REGISTRATION"]:
|
||||
abort(404)
|
||||
|
||||
if request.method == "GET":
|
||||
return render_template("login.html", mode="register", captcha="")
|
||||
|
||||
username = request.form.get("username", "").strip()
|
||||
password = request.form.get("password", "")
|
||||
|
||||
if not username:
|
||||
flash("用户名不能为空", "error")
|
||||
return render_template("login.html", mode="register", captcha=""), 400
|
||||
if len(password) < 6:
|
||||
flash("密码至少需要 6 位", "error")
|
||||
return render_template("login.html", mode="register", captcha=""), 400
|
||||
if User.query.filter_by(username=username).first():
|
||||
flash("用户名已存在", "error")
|
||||
return render_template("login.html", mode="register", captcha=""), 400
|
||||
|
||||
user = User(username=username, is_admin=False)
|
||||
user.set_password(password)
|
||||
db.session.add(user)
|
||||
db.session.commit()
|
||||
flash("注册成功，请登录", "info")
|
||||
session["captcha"] = new_captcha()
|
||||
return render_template("login.html", mode="login", captcha=session["captcha"])
|
||||
|
||||
|
||||
@bp.route("/logout", methods=["POST"])
|
||||
@login_required
|
||||
def logout():
|
||||
logout_user()
|
||||
return redirect(url_for("main.login"))
|
||||
|
||||
|
||||
@bp.route("/home")
|
||||
@login_required
|
||||
def home():
|
||||
return render_template("home.html")
|
||||
|
||||
|
||||
@bp.route("/history")
|
||||
@login_required
|
||||
def history_page():
|
||||
records = (
|
||||
UploadRecord.query.filter_by(user_id=current_user.id)
|
||||
.order_by(UploadRecord.upload_time.desc())
|
||||
.all()
|
||||
)
|
||||
return render_template("history.html", records=records)
|
||||
|
||||
|
||||
@bp.route("/admin")
|
||||
@login_required
|
||||
def admin_dashboard():
|
||||
if not current_user.is_admin:
|
||||
abort(403)
|
||||
records = UploadRecord.query.order_by(UploadRecord.upload_time.desc()).all()
|
||||
return render_template("admin.html", records=records)
|
||||
|
||||
|
||||
@bp.route("/download/<int:record_id>/<file_type>")
|
||||
@login_required
|
||||
def download_file(record_id: int, file_type: str):
|
||||
record = db.session.get(UploadRecord, record_id)
|
||||
if record is None:
|
||||
abort(404)
|
||||
if not (current_user.is_admin or current_user.id == record.user_id):
|
||||
abort(403)
|
||||
if file_type == "original":
|
||||
file_path = record.saved_path
|
||||
elif file_type == "prediction":
|
||||
file_path = record.prediction_path
|
||||
else:
|
||||
abort(404)
|
||||
if not os.path.exists(file_path):
|
||||
abort(404)
|
||||
return send_file(file_path, as_attachment=True)
|
||||
|
||||
|
||||
@bp.route("/download_template")
|
||||
def download_template():
|
||||
template_path = BASE_DIR / "example.xlsx"
|
||||
if not template_path.exists():
|
||||
abort(404)
|
||||
return send_file(template_path, as_attachment=True, download_name="example.xlsx")
|
||||
|
||||
|
||||
@bp.route("/result")
|
||||
@login_required
|
||||
def result_page():
|
||||
result = session.get("last_result")
|
||||
return render_template("result.html", result=result)
|
||||
|
||||
|
||||
@bp.route("/predict", methods=["POST"])
|
||||
@login_required
|
||||
def predict():
|
||||
model = current_app.config.get("RSF_MODEL")
|
||||
if model is None:
|
||||
return jsonify({"error": "模型未成功加载，请检查模型文件。"}), 500
|
||||
|
||||
uploaded = request.files.get("file")
|
||||
if uploaded is None or uploaded.filename == "":
|
||||
return jsonify({"error": "未选择文件"}), 400
|
||||
|
||||
try:
|
||||
artifacts = run_prediction(uploaded, int(current_user.id), model)
|
||||
except PredictionError as exc:
|
||||
return jsonify({"error": exc.message}), exc.status_code
|
||||
|
||||
record = UploadRecord(
|
||||
user_id=current_user.id,
|
||||
original_filename=artifacts.original_filename,
|
||||
saved_path=str(artifacts.saved_path),
|
||||
prediction_path=str(artifacts.excel_path),
|
||||
image_path=str(artifacts.image_path),
|
||||
)
|
||||
db.session.add(record)
|
||||
db.session.commit()
|
||||
|
||||
last_result = {
|
||||
"original_filename": artifacts.original_filename,
|
||||
"generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"image_url": url_for("static", filename=f"images/{artifacts.image_filename}"),
|
||||
"importance_url": (
|
||||
url_for("static", filename=f"images/{artifacts.importance_filename}")
|
||||
if artifacts.importance_filename
|
||||
else None
|
||||
),
|
||||
"excel_url": url_for("main.download_file", record_id=record.id, file_type="prediction"),
|
||||
"result_url": url_for("main.result_page"),
|
||||
"sample_count": int(artifacts.sample_count),
|
||||
"summary_rows": artifacts.summary_rows[:3],
|
||||
"analysis_text": artifacts.analysis_text,
|
||||
}
|
||||
session["last_result"] = last_result
|
||||
|
||||
return jsonify(
|
||||
{
|
||||
"message": "预测成功",
|
||||
"image_url": last_result["image_url"],
|
||||
"importance_url": last_result["importance_url"],
|
||||
"excel_url": last_result["excel_url"],
|
||||
"result_url": last_result["result_url"],
|
||||
"sample_count": last_result["sample_count"],
|
||||
"original_filename": artifacts.original_filename,
|
||||
}
|
||||
)
|
||||
Reference in New Issue
Block a user