jwxt()

1. 将 jwxt() 移动到 njupt_api 下，实现根据设置选择教务系统登录方式。
2. 将 api_router.py 和 mcp_router.py 中的对 ZhengFang() 的调用全部替换为对 jwxt() 的调用。

.idea/dictionaries/project.xml

@@ -0,0 +1,7 @@
+<component name="ProjectDictionaryState">
+  <dictionary name="project">
+    <words>
+      <w>jwxt</w>
+    </words>
+  </dictionary>
+</component>

njupt_api/zhengfang/__init__.py

@@ -1,4 +1,5 @@
 from .createcourse import create_course_schedule
+from .lib import jwxt
 from .sso import SSO
 from .types import Course, course_dict_serializer, course_list_serializer
 from .zhengfang import ZhengFang
@@ -10,4 +11,5 @@ __all__ = [
     course_dict_serializer,
     course_list_serializer,
     ZhengFang,
+    jwxt,
 ]

njupt_api/zhengfang/exc.py

@@ -0,0 +1,2 @@
+class LoginError(Exception):
+    pass

njupt_api/zhengfang/lib.py

@@ -0,0 +1,40 @@
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from ..baselib import config
+from .exc import LoginError
+from .sso import SSO
+from .zhengfang import ZhengFang
+
+
+@asynccontextmanager
+async def jwxt(username: str, password: str) -> AsyncGenerator[ZhengFang, None]:
+    """
+    根据设置，选择 SSO 登录教务系统或直接登录教务系统。
+
+    Args:
+        username: 用户名，str
+        password: 密码，str
+
+    Yields:
+        zf: ZhengFang
+
+    Raises:
+        LoginError: 登录失败，包含下层返回的提示信息。
+    """
+    try:
+        if config.get("schedule", "jwxt_login_method", "sso"):
+            sso = SSO()
+            await sso.start()
+            await sso.login(username, password)
+            zf = await ZhengFang.init_from_sso(sso)
+        else:
+            zf = ZhengFang()
+            await zf.start()
+            await zf.login(username, password)
+    except LoginError as e:
+        raise e
+    else:
+        yield zf
+        await zf.close()
+        return

njupt_api/zhengfang/sso.py

@@ -1,5 +1,7 @@
 from njupt_api.baselib import PlayContextManager, logger
 
+from .exc import LoginError
+
 
 class SSO(PlayContextManager):
     def __init__(self) -> None:
@@ -14,6 +16,9 @@ class SSO(PlayContextManager):
 
         Returns:
             bool，表明判登录是否成功。
+
+        Raises:
+            LoginError: 登录失败，暂时不包含任何提示信息……
         """
         await self.page.goto("http://i.njupt.edu.cn/")
 
@@ -24,7 +29,7 @@ class SSO(PlayContextManager):
         await self.page.wait_for_load_state("networkidle")
         if "user-login" in self.page.url:
             logger.error(f"{username} | 登录失败，请检查学号和密码是否正确。")
-            return False
+            raise LoginError("unknown")
 
         logger.info(f"{username} | 登录南邮统一身份认证成功。")
         self.isLogin = True

njupt_api/zhengfang/zhengfang.py

@@ -2,9 +2,11 @@ from ddddocr import DdddOcr
 from playwright.async_api import Browser, BrowserContext, Page, Playwright
 
 from njupt_api.baselib import PlayContextManager, logger
-from njupt_api.zhengfang import Course
-from njupt_api.zhengfang.createcourse import create_course_schedule
-from njupt_api.zhengfang.sso import SSO
+
+from .createcourse import create_course_schedule
+from .exc import LoginError
+from .sso import SSO
+from .types import Course
 
 
 class ZhengFang(PlayContextManager):
@@ -29,6 +31,9 @@ class ZhengFang(PlayContextManager):
 
         Returns:
             bool，表明登录是否成功。
+
+        Raises:
+            LoginError: 登录失败，包含提示信息
         """
         await self.page.goto("http://jwxt.njupt.edu.cn")
 
@@ -58,7 +63,7 @@ class ZhengFang(PlayContextManager):
             return await self.login(username, password)
         await dialog.accept()
         logger.error(f"{username} | 登录失败，教务系统提示信息为: {dialog.message}")
-        return False
+        raise LoginError(dialog.message)
 
     async def get_class_schedule(self) -> list[Course]:
         await self.page.locator("a.top_link:has-text('公用信息')").click()

router/api_router.py

@@ -7,10 +7,11 @@ from sqlmodel import Session, select
 
 from njupt_api.baselib import logger
 from njupt_api.zhengfang import (
-    ZhengFang,
     course_dict_serializer,
     course_list_serializer,
+    jwxt,
 )
+from njupt_api.zhengfang.exc import LoginError
 from router.enhance.lib import ReturnDto, ScheduleQueryDto, apply_enhance, get_session
 from router.enhance.model import Course
 
@@ -32,20 +33,16 @@ async def post_schedule_class(
         logger.success(f"{student.week=} 从数据库中返回一次性存储的班级课表。")
         return await apply_enhance(course_list, student.week, student.img)
     if student.username and student.password:
-        async with ZhengFang() as zf:
-            if await zf.login(student.username, student.password):
+        try:
+            async with jwxt(student.username, student.password) as zf:
                 course_list = course_list_serializer(await zf.get_class_schedule())
                 logger.success(
                     f"{student.username} | {student.week=} 获取指定学生的班级课表成功。",
                 )
                 return await apply_enhance(course_list, student.week, student.img)
-            logger.error(
-                f"{student.username} | 获取课程表失败，请检查账号密码是否正确后再试。",
-            )
-            return ReturnDto(
-                success=False,
-                message="获取课程表失败，请检查账号密码是否正确后再试。",
-            )
+        except LoginError as e:
+            return ReturnDto(success=False, message=str(e))
+
     else:
         logger.error(
             f"参数错误，请同时携带或同时不携带学号和密码参数: {student.username=} | {student.password=}",
@@ -65,18 +62,13 @@ async def post_schedule_student(student: ScheduleQueryDto) -> ReturnDto:
             message="查询学生课表需要同时提供学号和密码参数。",
         )
 
-    async with ZhengFang() as zf:
-        if await zf.login(student.username, student.password):
+    try:
+        async with jwxt(student.username, student.password) as zf:
             course_list = course_list_serializer(await zf.get_student_schedule())
             logger.success(f"{student.username} | 获取学生个人课表成功。")
             return await apply_enhance(course_list, student.week, student.img)
-        logger.error(
-            f"{student.username} | 获取课程表失败，请检查账号密码是否正确后再试。",
-        )
-        return ReturnDto(
-            success=False,
-            message="获取课程表失败，请检查账号密码是否正确后再试。",
-        )
+    except LoginError as e:
+        return ReturnDto(success=False, message=str(e))
 
 
 @api_router.get("/schedule/img/{name}")

router/mcp_router.py

@@ -9,10 +9,11 @@ from sqlmodel import Session, select
 
 from njupt_api.baselib import LoggingMiddleware, logger
 from njupt_api.zhengfang import (
-    ZhengFang,
     course_dict_serializer,
     course_list_serializer,
+    jwxt,
 )
+from njupt_api.zhengfang.exc import LoginError
 from router.enhance.lib import ReturnDto, apply_enhance
 from router.enhance.model import Course, engine
 
@@ -76,16 +77,13 @@ async def tool_schedule_class_special(
     week: WEEK_TYPE = 0,
     img: IMG_TYPE = False,
 ) -> ReturnDto:
-    async with ZhengFang() as zf:
-        if await zf.login(username, password):
+    try:
+        async with jwxt(username, password) as zf:
             final_course_list = course_list_serializer(await zf.get_class_schedule())
             logger.success(f"{username} | 获取指定学生的班级课表成功。")
             return await apply_enhance(final_course_list, week, img)
-        logger.error(f"{username} | 获取课程表失败，请检查账号密码是否正确后再试。")
-        return ReturnDto(
-            success=False,
-            message="获取课程表失败，请检查账号密码是否正确后再试。",
-        )
+    except LoginError as e:
+        return ReturnDto(success=False, message=str(e))
 
 
 @mcp.tool(
@@ -106,16 +104,13 @@ async def tool_schedule_student_special(
     week: WEEK_TYPE = 0,
     img: IMG_TYPE = False,
 ) -> ReturnDto:
-    async with ZhengFang() as zf:
-        if await zf.login(username, password):
+    try:
+        async with jwxt(username, password) as zf:
             final_course_list = course_list_serializer(await zf.get_student_schedule())
             logger.success(f"{username} | 获取指定学生的个人课表成功。")
             return await apply_enhance(final_course_list, week, img)
-        logger.error(f"{username} | 获取课程表失败，请检查账号密码是否正确后再试。")
-        return ReturnDto(
-            success=False,
-            message="获取课程表失败，请检查账号密码是否正确后再试。",
-        )
+    except LoginError as e:
+        return ReturnDto(success=False, message=str(e))
 
 
 @mcp.tool(

server.py

@@ -19,7 +19,6 @@ from njupt_api.baselib import (
     log_record_serialize,
     logger,
 )
-from njupt_api.zhengfang.zhengfang import ZhengFang
 from router import __version__
 from router.admin_router import admin_router
 from router.api_router import api_router
@@ -54,16 +53,6 @@ async def life_span(_: FastAPI) -> AsyncGenerator[None, None]:
         logger.info("配置文件监听任务已结束。")
 
 
-@asynccontextmanager
-async def jwxt(username: str, password: str) -> AsyncGenerator[ZhengFang, None]:
-    zf = ZhengFang()
-    await zf.start()
-    await zf.login(username, password)
-    yield zf
-    await zf.close()
-    return
-
-
 app = FastAPI(lifespan=combine_lifespans(life_span, mcp_app.lifespan))