From 6686b74a05c84517ed72d50aa74ba9912b19f4e6 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 24 Dec 2024 15:22:49 +0800 Subject: [PATCH] fix: enhance load method with overloads for better type hinting and simplify data loading in process_trace_tasks Signed-off-by: -LAN- --- api/extensions/ext_storage.py | 8 +++++++- api/tasks/ops_trace_task.py | 6 +----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/api/extensions/ext_storage.py b/api/extensions/ext_storage.py index 42422263c4dd03..588bdb2d2717e0 100644 --- a/api/extensions/ext_storage.py +++ b/api/extensions/ext_storage.py @@ -1,6 +1,6 @@ import logging from collections.abc import Callable, Generator -from typing import Union +from typing import Literal, Union, overload from flask import Flask @@ -79,6 +79,12 @@ def save(self, filename, data): logger.exception(f"Failed to save file {filename}") raise e + @overload + def load(self, filename: str, /, *, stream: Literal[False] = False) -> bytes: ... + + @overload + def load(self, filename: str, /, *, stream: Literal[True]) -> Generator: ... + def load(self, filename: str, /, *, stream: bool = False) -> Union[bytes, Generator]: try: if stream: diff --git a/api/tasks/ops_trace_task.py b/api/tasks/ops_trace_task.py index fa6b7afeab88ac..bb3b9e17ead6d2 100644 --- a/api/tasks/ops_trace_task.py +++ b/api/tasks/ops_trace_task.py @@ -1,6 +1,5 @@ import json import logging -from collections.abc import Generator from celery import shared_task # type: ignore from flask import current_app @@ -27,10 +26,7 @@ def process_trace_tasks(file_info): app_id = file_info.get("app_id") file_id = file_info.get("file_id") file_path = f"{OPS_FILE_PATH}{app_id}/{file_id}.json" - loaded_data = storage.load(file_path) - if isinstance(loaded_data, Generator): - loaded_data = b"".join(loaded_data) - file_data = json.loads(loaded_data) + file_data = json.loads(storage.load(file_path)) trace_info = file_data.get("trace_info") trace_info_type = file_data.get("trace_info_type") trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)