estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: mypy for all type check (#10921)

Browse Source
This commit is contained in:
yihong
2024-12-24 18:38:51 +08:00
committed by GitHub
parent c91e8b1737
commit 56e15d09a9
584 changed files with 3975 additions and 2826 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
api/extensions/storage/aliyun_oss_storage.py
View File

@@ -1,7 +1,7 @@
import posixpath
from collections.abc import Generator
import oss2 as aliyun_s3
import oss2 as aliyun_s3 # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -33,7 +33,7 @@ class AliyunOssStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
obj = self.client.get_object(self.__wrapper_folder_filename(filename))
data = obj.read()
data: bytes = obj.read()
return data
def load_stream(self, filename: str) -> Generator:
@@ -41,14 +41,14 @@ class AliyunOssStorage(BaseStorage):
while chunk := obj.read(4096):
yield chunk
def download(self, filename, target_filepath):
def download(self, filename: str, target_filepath):
self.client.get_object_to_file(self.__wrapper_folder_filename(filename), target_filepath)
def exists(self, filename):
def exists(self, filename: str):
return self.client.object_exists(self.__wrapper_folder_filename(filename))
def delete(self, filename):
def delete(self, filename: str):
self.client.delete_object(self.__wrapper_folder_filename(filename))
def __wrapper_folder_filename(self, filename) -> str:
def __wrapper_folder_filename(self, filename: str) -> str:
return posixpath.join(self.folder, filename) if self.folder else filename

8
api/extensions/storage/aws_s3_storage.py
View File

@@ -1,9 +1,9 @@
import logging
from collections.abc import Generator
import boto3
from botocore.client import Config
from botocore.exceptions import ClientError
import boto3 # type: ignore
from botocore.client import Config # type: ignore
from botocore.exceptions import ClientError # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -53,7 +53,7 @@ class AwsS3Storage(BaseStorage):
def load_once(self, filename: str) -> bytes:
try:
data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
data: bytes = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")

8
api/extensions/storage/azure_blob_storage.py
View File

@@ -27,7 +27,7 @@ class AzureBlobStorage(BaseStorage):
client = self._sync_client()
blob = client.get_container_client(container=self.bucket_name)
blob = blob.get_blob_client(blob=filename)
data = blob.download_blob().readall()
data: bytes = blob.download_blob().readall()
return data
def load_stream(self, filename: str) -> Generator:
@@ -63,11 +63,11 @@ class AzureBlobStorage(BaseStorage):
sas_token = cache_result.decode("utf-8")
else:
sas_token = generate_account_sas(
account_name=self.account_name,
account_key=self.account_key,
account_name=self.account_name or "",
account_key=self.account_key or "",
resource_types=ResourceTypes(service=True, container=True, object=True),
permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True),
expiry=datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1),
)
redis_client.set(cache_key, sas_token, ex=3000)
return BlobServiceClient(account_url=self.account_url, credential=sas_token)
return BlobServiceClient(account_url=self.account_url or "", credential=sas_token)

9
api/extensions/storage/baidu_obs_storage.py
View File

@@ -2,9 +2,9 @@ import base64
import hashlib
from collections.abc import Generator
from baidubce.auth.bce_credentials import BceCredentials
from baidubce.bce_client_configuration import BceClientConfiguration
from baidubce.services.bos.bos_client import BosClient
from baidubce.auth.bce_credentials import BceCredentials # type: ignore
from baidubce.bce_client_configuration import BceClientConfiguration # type: ignore
from baidubce.services.bos.bos_client import BosClient # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -36,7 +36,8 @@ class BaiduObsStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
response = self.client.get_object(bucket_name=self.bucket_name, key=filename)
return response.data.read()
data: bytes = response.data.read()
return data
def load_stream(self, filename: str) -> Generator:
response = self.client.get_object(bucket_name=self.bucket_name, key=filename).data

4
api/extensions/storage/google_cloud_storage.py
View File

@@ -3,7 +3,7 @@ import io
import json
from collections.abc import Generator
from google.cloud import storage as google_cloud_storage
from google.cloud import storage as google_cloud_storage # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -35,7 +35,7 @@ class GoogleCloudStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename)
data = blob.download_as_bytes()
data: bytes = blob.download_as_bytes()
return data
def load_stream(self, filename: str) -> Generator:

4
api/extensions/storage/huawei_obs_storage.py
View File

@@ -1,6 +1,6 @@
from collections.abc import Generator
from obs import ObsClient
from obs import ObsClient # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -23,7 +23,7 @@ class HuaweiObsStorage(BaseStorage):
self.client.putObject(bucketName=self.bucket_name, objectKey=filename, content=data)
def load_once(self, filename: str) -> bytes:
data = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response.read()
data: bytes = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response.read()
return data
def load_stream(self, filename: str) -> Generator:

8
api/extensions/storage/opendal_storage.py
View File

@@ -3,7 +3,7 @@ import os
from collections.abc import Generator
from pathlib import Path
import opendal
import opendal # type: ignore[import]
from dotenv import dotenv_values
from extensions.storage.base_storage import BaseStorage
@@ -18,7 +18,7 @@ def _get_opendal_kwargs(*, scheme: str, env_file_path: str = ".env", prefix: str
if key.startswith(config_prefix):
kwargs[key[len(config_prefix) :].lower()] = value
file_env_vars = dotenv_values(env_file_path)
file_env_vars: dict = dotenv_values(env_file_path) or {}
for key, value in file_env_vars.items():
if key.startswith(config_prefix) and key[len(config_prefix) :].lower() not in kwargs and value:
kwargs[key[len(config_prefix) :].lower()] = value
@@ -48,7 +48,7 @@ class OpenDALStorage(BaseStorage):
if not self.exists(filename):
raise FileNotFoundError("File not found")
content = self.op.read(path=filename)
content: bytes = self.op.read(path=filename)
logger.debug(f"file {filename} loaded")
return content
@@ -75,7 +75,7 @@ class OpenDALStorage(BaseStorage):
# error handler here when opendal python-binding has a exists method, we should use it
# more https://github.com/apache/opendal/blob/main/bindings/python/src/operator.rs
try:
res = self.op.stat(path=filename).mode.is_file()
res: bool = self.op.stat(path=filename).mode.is_file()
logger.debug(f"file {filename} checked")
return res
except Exception:

6
api/extensions/storage/oracle_oci_storage.py
View File

@@ -1,7 +1,7 @@
from collections.abc import Generator
import boto3
from botocore.exceptions import ClientError
import boto3 # type: ignore
from botocore.exceptions import ClientError # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -27,7 +27,7 @@ class OracleOCIStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
try:
data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
data: bytes = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")

2
api/extensions/storage/supabase_storage.py
View File

@@ -32,7 +32,7 @@ class SupabaseStorage(BaseStorage):
self.client.storage.from_(self.bucket_name).upload(filename, data)
def load_once(self, filename: str) -> bytes:
content = self.client.storage.from_(self.bucket_name).download(filename)
content: bytes = self.client.storage.from_(self.bucket_name).download(filename)
return content
def load_stream(self, filename: str) -> Generator:

4
api/extensions/storage/tencent_cos_storage.py
View File

@@ -1,6 +1,6 @@
from collections.abc import Generator
from qcloud_cos import CosConfig, CosS3Client
from qcloud_cos import CosConfig, CosS3Client # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -25,7 +25,7 @@ class TencentCosStorage(BaseStorage):
self.client.put_object(Bucket=self.bucket_name, Body=data, Key=filename)
def load_once(self, filename: str) -> bytes:
data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].get_raw_stream().read()
data: bytes = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].get_raw_stream().read()
return data
def load_stream(self, filename: str) -> Generator:

4
api/extensions/storage/volcengine_tos_storage.py
View File

@@ -1,6 +1,6 @@
from collections.abc import Generator
import tos
import tos # type: ignore
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
@@ -24,6 +24,8 @@ class VolcengineTosStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
data = self.client.get_object(bucket=self.bucket_name, key=filename).read()
if not isinstance(data, bytes):
raise TypeError("Expected bytes, got {}".format(type(data).__name__))
return data
def load_stream(self, filename: str) -> Generator: