Skip to content

Commit c0702aa

Browse files
Use typing.Literal to replace str places (langgenius#24099)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
1 parent 670d479 commit c0702aa

File tree

8 files changed

+34
-26
lines changed

8 files changed

+34
-26
lines changed

api/controllers/console/app/annotation.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Literal
2+
13
from flask import request
24
from flask_login import current_user
35
from flask_restful import Resource, marshal, marshal_with, reqparse
@@ -24,7 +26,7 @@ class AnnotationReplyActionApi(Resource):
2426
@login_required
2527
@account_initialization_required
2628
@cloud_edition_billing_resource_check("annotation")
27-
def post(self, app_id, action):
29+
def post(self, app_id, action: Literal["enable", "disable"]):
2830
if not current_user.is_editor:
2931
raise Forbidden()
3032

@@ -38,8 +40,6 @@ def post(self, app_id, action):
3840
result = AppAnnotationService.enable_app_annotation(args, app_id)
3941
elif action == "disable":
4042
result = AppAnnotationService.disable_app_annotation(app_id)
41-
else:
42-
raise ValueError("Unsupported annotation reply action")
4343
return result, 200
4444

4545

api/controllers/console/datasets/datasets_document.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import logging
22
from argparse import ArgumentTypeError
3-
from typing import cast
3+
from typing import Literal, cast
44

55
from flask import request
66
from flask_login import current_user
@@ -758,7 +758,7 @@ class DocumentProcessingApi(DocumentResource):
758758
@login_required
759759
@account_initialization_required
760760
@cloud_edition_billing_rate_limit_check("knowledge")
761-
def patch(self, dataset_id, document_id, action):
761+
def patch(self, dataset_id, document_id, action: Literal["pause", "resume"]):
762762
dataset_id = str(dataset_id)
763763
document_id = str(document_id)
764764
document = self.get_document(dataset_id, document_id)
@@ -784,8 +784,6 @@ def patch(self, dataset_id, document_id, action):
784784
document.paused_at = None
785785
document.is_paused = False
786786
db.session.commit()
787-
else:
788-
raise InvalidActionError()
789787

790788
return {"result": "success"}, 200
791789

@@ -840,7 +838,7 @@ class DocumentStatusApi(DocumentResource):
840838
@account_initialization_required
841839
@cloud_edition_billing_resource_check("vector_space")
842840
@cloud_edition_billing_rate_limit_check("knowledge")
843-
def patch(self, dataset_id, action):
841+
def patch(self, dataset_id, action: Literal["enable", "disable", "archive", "un_archive"]):
844842
dataset_id = str(dataset_id)
845843
dataset = DatasetService.get_dataset(dataset_id)
846844
if dataset is None:

api/controllers/console/datasets/metadata.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Literal
2+
13
from flask_login import current_user
24
from flask_restful import Resource, marshal_with, reqparse
35
from werkzeug.exceptions import NotFound
@@ -100,7 +102,7 @@ class DatasetMetadataBuiltInFieldActionApi(Resource):
100102
@login_required
101103
@account_initialization_required
102104
@enterprise_license_required
103-
def post(self, dataset_id, action):
105+
def post(self, dataset_id, action: Literal["enable", "disable"]):
104106
dataset_id_str = str(dataset_id)
105107
dataset = DatasetService.get_dataset(dataset_id_str)
106108
if dataset is None:

api/controllers/service_api/app/annotation.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Literal
2+
13
from flask import request
24
from flask_restful import Resource, marshal, marshal_with, reqparse
35
from werkzeug.exceptions import Forbidden
@@ -15,7 +17,7 @@
1517

1618
class AnnotationReplyActionApi(Resource):
1719
@validate_app_token
18-
def post(self, app_model: App, action):
20+
def post(self, app_model: App, action: Literal["enable", "disable"]):
1921
parser = reqparse.RequestParser()
2022
parser.add_argument("score_threshold", required=True, type=float, location="json")
2123
parser.add_argument("embedding_provider_name", required=True, type=str, location="json")
@@ -25,8 +27,6 @@ def post(self, app_model: App, action):
2527
result = AppAnnotationService.enable_app_annotation(args, app_model.id)
2628
elif action == "disable":
2729
result = AppAnnotationService.disable_app_annotation(app_model.id)
28-
else:
29-
raise ValueError("Unsupported annotation reply action")
3030
return result, 200
3131

3232

api/controllers/service_api/dataset/dataset.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Literal
2+
13
from flask import request
24
from flask_restful import marshal, marshal_with, reqparse
35
from werkzeug.exceptions import Forbidden, NotFound
@@ -358,14 +360,14 @@ def delete(self, _, dataset_id):
358360
class DocumentStatusApi(DatasetApiResource):
359361
"""Resource for batch document status operations."""
360362

361-
def patch(self, tenant_id, dataset_id, action):
363+
def patch(self, tenant_id, dataset_id, action: Literal["enable", "disable", "archive", "un_archive"]):
362364
"""
363365
Batch update document status.
364366
365367
Args:
366368
tenant_id: tenant id
367369
dataset_id: dataset id
368-
action: action to perform (enable, disable, archive, un_archive)
370+
action: action to perform (Literal["enable", "disable", "archive", "un_archive"])
369371
370372
Returns:
371373
dict: A dictionary with a key 'result' and a value 'success'

api/controllers/service_api/dataset/metadata.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Literal
2+
13
from flask_login import current_user # type: ignore
24
from flask_restful import marshal, reqparse
35
from werkzeug.exceptions import NotFound
@@ -77,7 +79,7 @@ def get(self, tenant_id):
7779

7880
class DatasetMetadataBuiltInFieldActionServiceApi(DatasetApiResource):
7981
@cloud_edition_billing_rate_limit_check("knowledge", "dataset")
80-
def post(self, tenant_id, dataset_id, action):
82+
def post(self, tenant_id, dataset_id, action: Literal["enable", "disable"]):
8183
dataset_id_str = str(dataset_id)
8284
dataset = DatasetService.get_dataset(dataset_id_str)
8385
if dataset is None:

api/services/dataset_service.py

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import time
77
import uuid
88
from collections import Counter
9-
from typing import Any, Optional
9+
from typing import Any, Literal, Optional
1010

1111
from flask_login import current_user
1212
from sqlalchemy import func, select
@@ -51,7 +51,7 @@
5151
RetrievalModel,
5252
SegmentUpdateArgs,
5353
)
54-
from services.errors.account import InvalidActionError, NoPermissionError
54+
from services.errors.account import NoPermissionError
5555
from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError
5656
from services.errors.dataset import DatasetNameDuplicateError
5757
from services.errors.document import DocumentIndexingError
@@ -1800,14 +1800,16 @@ def estimate_args_validate(cls, args: dict):
18001800
raise ValueError("Process rule segmentation max_tokens is invalid")
18011801

18021802
@staticmethod
1803-
def batch_update_document_status(dataset: Dataset, document_ids: list[str], action: str, user):
1803+
def batch_update_document_status(
1804+
dataset: Dataset, document_ids: list[str], action: Literal["enable", "disable", "archive", "un_archive"], user
1805+
):
18041806
"""
18051807
Batch update document status.
18061808
18071809
Args:
18081810
dataset (Dataset): The dataset object
18091811
document_ids (list[str]): List of document IDs to update
1810-
action (str): Action to perform (enable, disable, archive, un_archive)
1812+
action (Literal["enable", "disable", "archive", "un_archive"]): Action to perform
18111813
user: Current user performing the action
18121814
18131815
Raises:
@@ -1890,9 +1892,10 @@ def batch_update_document_status(dataset: Dataset, document_ids: list[str], acti
18901892
raise propagation_error
18911893

18921894
@staticmethod
1893-
def _prepare_document_status_update(document, action: str, user):
1894-
"""
1895-
Prepare document status update information.
1895+
def _prepare_document_status_update(
1896+
document: Document, action: Literal["enable", "disable", "archive", "un_archive"], user
1897+
):
1898+
"""Prepare document status update information.
18961899
18971900
Args:
18981901
document: Document object to update
@@ -2355,7 +2358,9 @@ def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset
23552358
db.session.commit()
23562359

23572360
@classmethod
2358-
def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset, document: Document):
2361+
def update_segments_status(
2362+
cls, segment_ids: list, action: Literal["enable", "disable"], dataset: Dataset, document: Document
2363+
):
23592364
# Check if segment_ids is not empty to avoid WHERE false condition
23602365
if not segment_ids or len(segment_ids) == 0:
23612366
return
@@ -2413,8 +2418,6 @@ def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset
24132418
db.session.commit()
24142419

24152420
disable_segments_from_index_task.delay(real_deal_segment_ids, dataset.id, document.id)
2416-
else:
2417-
raise InvalidActionError()
24182421

24192422
@classmethod
24202423
def create_child_chunk(

api/tasks/deal_dataset_vector_index_task.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import logging
22
import time
3+
from typing import Literal
34

45
import click
56
from celery import shared_task # type: ignore
@@ -13,7 +14,7 @@
1314

1415

1516
@shared_task(queue="dataset")
16-
def deal_dataset_vector_index_task(dataset_id: str, action: str):
17+
def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "add", "update"]):
1718
"""
1819
Async deal dataset from index
1920
:param dataset_id: dataset_id

0 commit comments

Comments
 (0)