AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-13 20:16:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: abstract common function with local storage (#9811)

Browse Source
This commit is contained in:
ice yao 2024-10-24 21:53:37 +08:00 committed by GitHub
parent 22bb0414a1
commit e71f494839
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 23 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
api/extensions/storage/local_fs_storage.py
View File

@ -19,68 +19,48 @@ class LocalFsStorage(BaseStorage):
folder = os.path.join(current_app.root_path, folder) folder = os.path.join(current_app.root_path, folder)
self.folder = folder self.folder = folder
def save(self, filename, data): def _build_filepath(self, filename: str) -> str:
"""Build the full file path based on the folder and filename."""
if not self.folder or self.folder.endswith("/"): if not self.folder or self.folder.endswith("/"):
filename = self.folder + filename return self.folder + filename
else: else:
filename = self.folder + "/" + filename return self.folder + "/" + filename
folder = os.path.dirname(filename) def save(self, filename, data):
filepath = self._build_filepath(filename)
folder = os.path.dirname(filepath)
os.makedirs(folder, exist_ok=True) os.makedirs(folder, exist_ok=True)
Path(os.path.join(os.getcwd(), filepath)).write_bytes(data)
Path(os.path.join(os.getcwd(), filename)).write_bytes(data)
def load_once(self, filename: str) -> bytes: def load_once(self, filename: str) -> bytes:
if not self.folder or self.folder.endswith("/"): filepath = self._build_filepath(filename)
filename = self.folder + filename if not os.path.exists(filepath):
else:
filename = self.folder + "/" + filename
if not os.path.exists(filename):
raise FileNotFoundError("File not found") raise FileNotFoundError("File not found")
return Path(filepath).read_bytes()
data = Path(filename).read_bytes()
return data
def load_stream(self, filename: str) -> Generator: def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator: filepath = self._build_filepath(filename)
if not self.folder or self.folder.endswith("/"):
filename = self.folder + filename
else:
filename = self.folder + "/" + filename
if not os.path.exists(filename): def generate() -> Generator:
if not os.path.exists(filepath):
raise FileNotFoundError("File not found") raise FileNotFoundError("File not found")
with open(filepath, "rb") as f:
with open(filename, "rb") as f:
while chunk := f.read(4096): # Read in chunks of 4KB while chunk := f.read(4096): # Read in chunks of 4KB
yield chunk yield chunk
return generate() return generate()
def download(self, filename, target_filepath): def download(self, filename, target_filepath):
if not self.folder or self.folder.endswith("/"): filepath = self._build_filepath(filename)
filename = self.folder + filename if not os.path.exists(filepath):
else:
filename = self.folder + "/" + filename
if not os.path.exists(filename):
raise FileNotFoundError("File not found") raise FileNotFoundError("File not found")
shutil.copyfile(filepath, target_filepath)
shutil.copyfile(filename, target_filepath)
def exists(self, filename): def exists(self, filename):
if not self.folder or self.folder.endswith("/"): filepath = self._build_filepath(filename)
filename = self.folder + filename return os.path.exists(filepath)
else:
filename = self.folder + "/" + filename
return os.path.exists(filename)
def delete(self, filename): def delete(self, filename):
if not self.folder or self.folder.endswith("/"): filepath = self._build_filepath(filename)
filename = self.folder + filename if os.path.exists(filepath):
else: os.remove(filepath)
filename = self.folder + "/" + filename
if os.path.exists(filename):
os.remove(filename)