mirror of
https://github.com/Tencent/WeKnora.git
synced 2025-11-25 03:15:00 +08:00
2d66abedf0
移除日志设置与冗余代码，优化导入、类型提示及OCR后端管理 统一调整各文件模块导入路径为绝对导入 调整导入路径，移除部分导入，优化日志及注释 升级文档解析器为 Docx2Parser，优化超时与图片处理逻辑
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
import logging
|
|
import os
|
|
import tempfile
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TempFileContext:
|
|
def __init__(self, file_content: bytes, suffix: str):
|
|
"""
|
|
Initialize the context
|
|
:param file_content: Byte data to write to file
|
|
:param suffix: File suffix
|
|
"""
|
|
self.file_content = file_content
|
|
self.suffix = suffix
|
|
self.file = None
|
|
|
|
def __enter__(self):
|
|
"""
|
|
Create file when entering context
|
|
"""
|
|
self.temp_file = tempfile.NamedTemporaryFile(suffix=self.suffix, delete=False)
|
|
self.temp_file.write(self.file_content)
|
|
self.temp_file.flush()
|
|
logger.info(
|
|
f"Saved {self.suffix} content to temporary file: {self.temp_file.name}"
|
|
)
|
|
return self.temp_file.name
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""
|
|
Delete file when exiting context
|
|
"""
|
|
if self.temp_file:
|
|
self.temp_file.close()
|
|
if os.path.exists(self.temp_file.name):
|
|
os.remove(self.temp_file.name)
|
|
logger.info(f"File {self.temp_file.name} has been deleted.")
|
|
# Return False to propagate exception (if any exception occurred)
|
|
return False
|
|
|
|
|
|
class TempDirContext:
|
|
def __init__(self):
|
|
"""
|
|
Initialize the context
|
|
"""
|
|
self.temp_dir = None
|
|
|
|
def __enter__(self):
|
|
"""
|
|
Create directory when entering context
|
|
"""
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
logger.info(f"Created temporary directory: {self.temp_dir.name}")
|
|
return self.temp_dir.name
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""
|
|
Delete directory when exiting context
|
|
"""
|
|
if self.temp_dir and os.path.exists(self.temp_dir.name):
|
|
self.temp_dir.cleanup()
|
|
logger.info(f"Directory {self.temp_dir.name} has been deleted.")
|
|
# Return False to propagate exception (if any exception occurred)
|
|
return False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
example_bytes = b"Hello, this is a test file."
|
|
file_name = "test_file.txt"
|
|
|
|
# Using with statement
|
|
with TempFileContext(example_bytes, file_name) as temp_file:
|
|
# File operations can be performed within the context
|
|
print(f"Does file {file_name} exist: {os.path.exists(file_name)}")
|