diff --git a/app/core/views.py b/app/core/views.py index 98c90d5..84d4c33 100644 --- a/app/core/views.py +++ b/app/core/views.py @@ -1,5 +1,3 @@ -from typing import Optional - from flask import g from flask_apikit.views import APIView from app.models.user import User @@ -7,5 +5,5 @@ class MoeAPIView(APIView): @property - def current_user(self) -> Optional[User]: + def current_user(self) -> User: return g.get("current_user") diff --git a/app/models/user.py b/app/models/user.py index dd92cd4..5f50db7 100644 --- a/app/models/user.py +++ b/app/models/user.py @@ -36,7 +36,7 @@ from app.models.message import Message from app.models.project import Project, ProjectRole, ProjectUserRelation from app.models.site_setting import SiteSetting -from app.models.team import Team, TeamPermission, TeamUserRelation +from app.models.team import Team, TeamPermission, TeamRole, TeamUserRelation from app.regexs import EMAIL_REGEX, USER_NAME_REGEX from app.constants.locale import Locale from app.utils.hash import md5 @@ -339,7 +339,7 @@ def projects( projects = mongo_slice(projects, skip, limit) return projects - def get_project_relation(self, project): + def get_project_relation(self, project) -> ProjectUserRelation | None: """获取与某个项目的关系""" relation = ProjectUserRelation.objects(user=self, group=project).first() if relation: @@ -354,14 +354,18 @@ def join_project(self, project, role=None): return ProjectUserRelation(user=self, group=project, role=role).save() # =====加入流程===== - def invitations(self, group=None, status=None, skip=None, limit=None): + def invitations( + self, group=None, status=None, skip=None, limit=None + ) -> list[Invitation]: """获取对于个人的所有的邀请""" invitations = Invitation.get( user=self, group=group, status=status, skip=skip, limit=limit ) return invitations - def applications(self, group=None, status=None, skip=None, limit=None): + def applications( + self, group=None, status=None, skip=None, limit=None + ) -> list[Application]: """获取自己发出的所有的申请""" applications = Application.get( user=self, @@ -372,7 +376,13 @@ def applications(self, group=None, status=None, skip=None, limit=None): ) return applications - def invite(self, user, group, role, message=""): + def invite( + self, + user: "User", + group: Team | Project, + role: ProjectRole | TeamRole, + message="", + ): """邀请某个用户加入某个团体""" # 判断团体是否已满员 if group.is_full(): @@ -478,7 +488,7 @@ def apply(self, group, message=""): return {"message": gettext("申请成功,请等待管理员审核")} # =====自动鉴别部分===== - def get_relation(self, group): + def get_relation(self, group: Team | Project): """返回与目标的关系,返回关系对象或None,以此判断是否加入""" if isinstance(group, Team): return self.get_team_relation(group) diff --git a/app/scripts/__init__.py b/app/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/scripts/fsck.py b/app/scripts/fsck.py new file mode 100644 index 0000000..58f94aa --- /dev/null +++ b/app/scripts/fsck.py @@ -0,0 +1,68 @@ +from dataclasses import dataclass +from typing import Literal +import io +import click +from app.models.file import File +from app.models.team import Team +from app.models.project import ProjectSet, Project +import csv +import logging + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, kw_only=True) +class SizeEntry: + team_id: str + team_name: str + project_set_id: str + project_set_name: str + project_id: str + project_name: str + num_files: int + total_size_kb: int + + +@click.command("fsck") +@click.option("--export-csv", type=click.File("w"), help="export statistics to csv") +@click.option( + "--sum-by", + type=click.Choice(["team", "project_set", "project"]), + default="team", + help="export statistics to csv", +) +def fsck( + *, export_csv: io.FileIO | None, sum_by: Literal["team", "project_set", "project"] +): + logger.info("Running fsck") + size_entries: list[SizeEntry] = [ + _sum_by_project(t, ps, p) + for t in Team.objects() + for ps in ProjectSet.objects(team=t) + for p in Project.objects(project_set=ps) + ] + if export_csv: + writer = csv.DictWriter(export_csv, fieldnames=SizeEntry.__dataclass_fields__) + writer.writeheader() + for entry in size_entries: + writer.writerow(entry.__dict__) + export_csv.close() + else: + for entry in size_entries: + logger.info(entry) + + +def _sum_by_project(t: Team, ps: ProjectSet, p: Project) -> SizeEntry: + files = File.objects(project=p) + num_files = len(files) + size_kb = sum(f.file_size for f in files) + return SizeEntry( + team_id=t.id, + team_name=t.name, + project_set_id=ps.id, + project_set_name=ps.name, + project_id=p.id, + project_name=p.name, + total_size_kb=size_kb, + num_files=num_files, + ) diff --git a/app/utils/mongo.py b/app/utils/mongo.py index d76be3c..97d29b9 100644 --- a/app/utils/mongo.py +++ b/app/utils/mongo.py @@ -3,7 +3,11 @@ T = TypeVar("T") -def mongo_order(objects: List[T], order_by, default_order_by) -> List[T]: +def mongo_order( + objects: List[T], + order_by: None | list[str] | str, + default_order_by: str | list[str], +) -> List[T]: """处理排序""" # 设置排序默认值 if order_by is None or order_by == []: diff --git a/manage.py b/manage.py index 3ff8a6a..7f3ce85 100644 --- a/manage.py +++ b/manage.py @@ -4,6 +4,7 @@ import logging from app import flask_app from app.factory import init_db +from app.scripts.fsck import fsck logging.basicConfig( level=logging.INFO, @@ -81,6 +82,8 @@ def mit_preprocess_dir(dir: str): main.add_command(list_translations) main.add_command(mit_preprocess_file) main.add_command(mit_preprocess_dir) +main.add_command(fsck) + if __name__ == "__main__": main()