228 lines
7.2 KiB
Python
Raw Permalink Normal View History

2025-11-12 00:28:07 +08:00
"""
数据模型定义
定义了项目中使用的主要数据结构
"""
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import List, Optional, Dict, Any
from pathlib import Path
class PlatformType(Enum):
"""支持的平台类型"""
XIAOHONGSHU = "xiaohongshu"
DOUYIN = "douyin"
class UploadStatus(Enum):
"""上传状态枚举"""
PENDING = "pending" # 等待上传
UPLOADING = "uploading" # 上传中
PROCESSING = "processing" # 处理中
COMPLETED = "completed" # 上传完成
FAILED = "failed" # 上传失败
CANCELLED = "cancelled" # 取消上传
class PublishStatus(Enum):
"""发布状态枚举"""
DRAFT = "draft" # 草稿
PENDING = "pending" # 等待发布
PUBLISHING = "publishing" # 发布中
PUBLISHED = "published" # 已发布
FAILED = "failed" # 发布失败
SCHEDULED = "scheduled" # 定时发布
@dataclass
class AccountInfo:
"""账号信息"""
platform: PlatformType = field()
username: str = field()
cookie_file: str = field()
is_active: bool = True
user_data_dir: Optional[str] = None
extra_config: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
"""后处理,确保路径存在"""
if not self.cookie_file.endswith('.json'):
self.cookie_file = f"{self.cookie_file}.json"
@dataclass
class Content:
"""内容基类"""
title: str = field()
description: str = field()
tags: List[str] = field(default_factory=list)
visibility: str = "public" # public, private, friends
location: Optional[str] = None
scheduled_time: Optional[datetime] = None
extra_data: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
"""后处理,清理数据"""
self.title = self.title.strip()
self.description = self.description.strip()
self.tags = [tag.strip() for tag in self.tags if tag.strip()]
@dataclass
class ImageNote(Content):
"""图文笔记内容"""
images: List[str] = field(default_factory=list)
cover_image: Optional[str] = None
image_captions: List[str] = field(default_factory=list) # 图片说明
def __post_init__(self):
super().__post_init__()
# 验证图片文件
self.images = [str(Path(img)) for img in self.images if Path(img).exists()]
if self.cover_image and not Path(self.cover_image).exists():
self.cover_image = None
# 如果没有指定封面,使用第一张图片
if not self.cover_image and self.images:
self.cover_image = self.images[0]
@dataclass
class VideoContent(Content):
"""视频内容"""
video_path: str
cover_image: Optional[str] = None
duration: Optional[int] = None # 秒
resolution: Optional[tuple] = None # (width, height)
file_size: Optional[int] = None # bytes
def __post_init__(self):
super().__post_init__()
video_path = Path(self.video_path)
if video_path.exists():
self.video_path = str(video_path)
else:
raise FileNotFoundError(f"视频文件不存在: {self.video_path}")
if self.cover_image and not Path(self.cover_image).exists():
self.cover_image = None
@dataclass
class PublishTask:
"""发布任务"""
id: str = field()
platform: PlatformType = field()
account: AccountInfo = field()
content: Content = field()
scheduled_time: Optional[datetime] = None
retry_count: int = 0
max_retries: int = 3
created_at: datetime = field(default_factory=datetime.now)
metadata: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
"""生成任务ID如果没有提供"""
if not self.id:
self.id = f"{self.platform.value}_{self.account.username}_{int(self.created_at.timestamp())}"
@property
def should_retry(self) -> bool:
"""是否应该重试"""
return self.retry_count < self.max_retries
def increment_retry(self):
"""增加重试次数"""
self.retry_count += 1
@dataclass
class PublishResult:
"""发布结果"""
task_id: str = field()
platform: PlatformType = field()
account: str = field()
success: bool = field()
message: str = field()
published_url: Optional[str] = None
upload_status: UploadStatus = UploadStatus.PENDING
publish_status: PublishStatus = PublishStatus.DRAFT
timestamp: datetime = field(default_factory=datetime.now)
error_details: Optional[Dict[str, Any]] = None
duration: Optional[float] = None # 执行耗时(秒)
def to_dict(self) -> Dict[str, Any]:
"""转换为字典格式"""
return {
"task_id": self.task_id,
"platform": self.platform.value,
"account": self.account,
"success": self.success,
"message": self.message,
"published_url": self.published_url,
"upload_status": self.upload_status.value,
"publish_status": self.publish_status.value,
"timestamp": self.timestamp.isoformat(),
"error_details": self.error_details,
"duration": self.duration
}
@dataclass
class PlatformConfig:
"""平台配置信息"""
name: str
login_url: str
supported_content_types: List[str]
max_file_size: int # bytes
max_duration: Optional[int] = None # 秒
supported_formats: List[str] = field(default_factory=list)
extra_config: Dict[str, Any] = field(default_factory=dict)
@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> 'PlatformConfig':
"""从字典创建配置"""
return cls(
name=config_dict.get("name", ""),
login_url=config_dict.get("login_url", ""),
supported_content_types=config_dict.get("supported_content_types", []),
max_file_size=config_dict.get("max_file_size", 0),
max_duration=config_dict.get("max_duration"),
supported_formats=config_dict.get("supported_formats", []),
extra_config=config_dict.get("extra_config", {})
)
# 常用的平台配置模板
XIAOHONGSHU_CONFIG = PlatformConfig(
name="小红书",
login_url="https://creator.xiaohongshu.com/",
supported_content_types=["image_note", "video_note"],
max_file_size=500 * 1024 * 1024, # 500MB
max_duration=300, # 5分钟
supported_formats=["jpg", "jpeg", "png", "mp4", "mov"],
extra_config={
"max_images": 9,
"image_note_url": "https://creator.xiaohongshu.com/publish/publish",
"video_note_url": "https://creator.xiaohongshu.com/publish/video",
}
)
DOUYIN_CONFIG = PlatformConfig(
name="抖音",
login_url="https://creator.douyin.com/creator-micro/home",
supported_content_types=["video"],
max_file_size=2 * 1024 * 1024 * 1024, # 2GB
max_duration=600, # 10分钟
supported_formats=["mp4", "mov", "avi"],
extra_config={
"upload_url": "https://creator.douyin.com/creator-micro/content/upload",
}
)
# 平台配置映射
PLATFORM_CONFIGS = {
PlatformType.XIAOHONGSHU: XIAOHONGSHU_CONFIG,
PlatformType.DOUYIN: DOUYIN_CONFIG
}