Backup

superset_tool/utils/fileio.py

@@ -13,24 +13,35 @@ from ..utils.logger import SupersetLogger
 
 @contextmanager
 def create_temp_file(
-    content: bytes,
+    content: Optional[bytes] = None,
     suffix: str = ".zip",
+    mode: str = 'wb',
     logger: Optional[SupersetLogger] = None
 ):
-    """Контекстный менеджер для создания временных файлов с логированием"""
+    """Расширенный контекстный менеджер для временных файлов/директорий"""
     logger = logger or SupersetLogger(name="fileio", console=False)
     try:
-        logger.debug(f"Создание временного файла с суффиксом {suffix}")
-        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
-            tmp.write(content)
-            tmp.flush()
-        yield Path(tmp.name)
+        logger.debug(f"Создание временного ресурса с суффиксом {suffix}")
+        
+        # Для директорий
+        if suffix.startswith('.dir'):
+            with tempfile.TemporaryDirectory(suffix=suffix) as tmp_dir:
+                logger.debug(f"Создана временная директория: {tmp_dir}")
+                yield Path(tmp_dir)
+        
+        # Для файлов
+        else:
+            with tempfile.NamedTemporaryFile(suffix=suffix, mode=mode, delete=False) as tmp:
+                if content:
+                    tmp.write(content)
+                    tmp.flush()
+                logger.debug(f"Создан временный файл: {tmp.name}")
+                yield Path(tmp.name)
     except Exception as e:
-        logger.error(
-            f"Ошибка создания временного файла: {str(e)}", exc_info=True)
+        logger.error(f"Ошибка создания временного ресурса: {str(e)}", exc_info=True)
         raise
     finally:
-        if Path(tmp.name).exists():
+        if 'tmp' in locals() and Path(tmp.name).exists() and not suffix.startswith('.dir'):
             Path(tmp.name).unlink()
             logger.debug(f"Временный файл {tmp.name} удален")