feat: Normalize the code

Author: ZQlQZ

veadk/database/tos/toshandler.py veadk/database/tos/tos_handler.pyveadk/database/tos/toshandler.py renamed to veadk/database/tos/tos_handler.py

@@ -4,7 +4,7 @@
 import tos
 from datetime import datetime
 import asyncio
-from typing import Literal
+from typing import Literal, Union
 from urllib.parse import urlparse
 
 logger = get_logger(__name__)
@@ -13,7 +13,7 @@
 class TOSHandler:
     def __init__(self):
         """Initialize TOS configuration information"""
-        self.region = getenv("VOLCENGINE_REGION")
+        self.region = getenv("DATABASE_TOS_REGION")
         self.ak = getenv("VOLCENGINE_ACCESS_KEY")
         self.sk = getenv("VOLCENGINE_SECRET_KEY")
         self.bucket_name = getenv("DATABASE_TOS_BUCKET")
@@ -63,24 +63,26 @@ def get_suffix(self, data_path: str) -> str:
                 return f".{candidate.lower()}"
         return f".{parts[-1].lower()}"
 
-    def gen_url(self, user_id, app_name, session_id, data_path):
+    def gen_url(
+        self, user_id: str, app_name: str, session_id: str, data_path: str
+    ) -> str:
         """generate TOS URL"""
-        suffix = self.get_suffix(data_path)
-        timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
-        url = (
+        suffix: str = self.get_suffix(data_path)
+        timestamp: str = datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3]
+        url: str = (
             f"{self.bucket_name}/{app_name}/{user_id}-{session_id}-{timestamp}{suffix}"
         )
         return url
 
-    def parse_url(self, url):
+    def parse_url(self, url: str) -> tuple[str, str]:
         """Parse the URL to obtain bucket_name and object_key"""
         """bucket_name/object_key"""
         parts = url.split("/", 1)
         if len(parts) < 2:
             raise ValueError("URL format error, it should be: bucket_name/object_key")
         return parts
 
-    def create_bucket(self, client, bucket_name):
+    def create_bucket(self, client: tos.TosClientV2, bucket_name: str) -> bool:
         """If the bucket does not exist, create it"""
         try:
             client.head_bucket(self.bucket_name)
@@ -99,7 +101,9 @@ def create_bucket(self, client, bucket_name):
             logger.error(f"Bucket creation failed: {str(e)}")
             return False
 
-    def upload_to_tos(self, url: str, data, data_type: Literal["file", "bytes"]):
+    def upload_to_tos(
+        self, url: str, data: Union[str, bytes], data_type: Literal["file", "bytes"]
+    ):
         if data_type not in ("file", "bytes"):
             error_msg = f"Upload failed: data_type error. Only 'file' and 'bytes' are supported, got {data_type}"
             logger.error(error_msg)
@@ -109,7 +113,7 @@ def upload_to_tos(self, url: str, data, data_type: Literal["file", "bytes"]):
         elif data_type == "bytes":
             return asyncio.to_thread(self._do_upload_bytes, url, data)
 
-    def _do_upload_bytes(self, url, bytes):
+    def _do_upload_bytes(self, url: str, bytes: bytes) -> bool:
         bucket_name, object_key = self.parse_url(url)
         client = self._init_tos_client()
         try:
@@ -127,7 +131,7 @@ def _do_upload_bytes(self, url, bytes):
             if client:
                 client.close()
 
-    def _do_upload_file(self, url, file_path):
+    def _do_upload_file(self, url: str, file_path: str) -> bool:
         bucket_name, object_key = self.parse_url(url)
         client = self._init_tos_client()
         try:
@@ -147,7 +151,7 @@ def _do_upload_file(self, url, file_path):
             if client:
                 client.close()
 
-    def download_from_tos(self, url, save_path):
+    def download_from_tos(self, url: str, save_path: str) -> bool:
         """download image from TOS"""
         try:
             bucket_name, object_key = self.parse_url(url)

veadk/runner.py

@@ -32,7 +32,7 @@
 from veadk.types import MediaMessage
 from veadk.utils.logger import get_logger
 from veadk.utils.misc import read_png_to_bytes
-from veadk.database.tos.toshandler import TOSHandler
+from veadk.database.tos.tos_handler import TOSHandler
 
 logger = get_logger(__name__)