first commit
This commit is contained in:
+641
@@ -0,0 +1,641 @@
|
||||
"""
|
||||
XiaoHongShu Scraper - 核心采集逻辑
|
||||
支持单线程和多线程模式
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Optional, AsyncIterator
|
||||
from playwright.async_api import async_playwright, Browser, BrowserContext, Page, TimeoutError as PlaywrightTimeoutError
|
||||
from app import config
|
||||
from app.log_manager import LogManager
|
||||
|
||||
|
||||
class RiskControlException(Exception):
|
||||
"""风控异常"""
|
||||
def __init__(self, reason: str):
|
||||
self.reason = reason
|
||||
super().__init__(f"触发风控: {reason}")
|
||||
|
||||
|
||||
class XiaoHongShuScraper:
|
||||
"""小红书采集器"""
|
||||
|
||||
def __init__(self, event_queue: asyncio.Queue, log_manager: LogManager):
|
||||
self.event_queue = event_queue
|
||||
self.log_manager = log_manager
|
||||
self.playwright = None
|
||||
self.browser: Optional[Browser] = None
|
||||
self.is_running = False
|
||||
self._stop_requested = False
|
||||
|
||||
async def init_browser(self):
|
||||
"""初始化浏览器"""
|
||||
self.playwright = await async_playwright().start()
|
||||
self.browser = await self.playwright.chromium.launch(
|
||||
headless=config.HEADLESS,
|
||||
args=config.BROWSER_ARGS
|
||||
)
|
||||
await self.emit_log("SUCCESS", "浏览器初始化完成")
|
||||
|
||||
async def close_browser(self):
|
||||
"""关闭浏览器"""
|
||||
if self.browser:
|
||||
await self.browser.close()
|
||||
if self.playwright:
|
||||
await self.playwright.stop()
|
||||
await self.emit_log("INFO", "浏览器已关闭")
|
||||
|
||||
async def emit_log(self, level: str, message: str, thread_id: Optional[str] = None):
|
||||
"""发送日志事件"""
|
||||
# 持久化到日志管理器
|
||||
log_entry = self.log_manager.append(level, message, thread_id)
|
||||
|
||||
# 推送到 Web 前端
|
||||
await self.event_queue.put({
|
||||
"type": "log",
|
||||
**log_entry
|
||||
})
|
||||
|
||||
async def emit_progress(self, **kwargs):
|
||||
"""发送进度更新"""
|
||||
await self.event_queue.put({
|
||||
"type": "progress",
|
||||
**kwargs
|
||||
})
|
||||
|
||||
async def emit_user_status(self, user_id: str, status: str, note_count: int = 0, thread_id: Optional[str] = None):
|
||||
"""发送用户状态"""
|
||||
await self.event_queue.put({
|
||||
"type": "user_status",
|
||||
"user_id": user_id,
|
||||
"status": status,
|
||||
"note_count": note_count,
|
||||
"thread_id": thread_id
|
||||
})
|
||||
|
||||
async def emit_thread_status(self, thread_id: str, status: str, current_user: Optional[str] = None):
|
||||
"""发送线程状态"""
|
||||
await self.event_queue.put({
|
||||
"type": "thread_status",
|
||||
"thread_id": thread_id,
|
||||
"status": status,
|
||||
"current_user": current_user
|
||||
})
|
||||
|
||||
async def handle_login(self, context: BrowserContext) -> bool:
|
||||
"""
|
||||
处理登录
|
||||
返回是否已登录
|
||||
"""
|
||||
page = await context.new_page()
|
||||
await page.goto(config.LOGIN_URL)
|
||||
|
||||
await self.emit_log("WARNING", "请在浏览器中完成扫码登录...")
|
||||
await self.emit_log("INFO", "等待登录确认...")
|
||||
|
||||
# 等待登录成功（检查 URL 变化或特定元素）
|
||||
try:
|
||||
# 等待跳转到首页或出现用户头像
|
||||
await page.wait_for_url(
|
||||
f"{config.XIAOHONGSHU_BASE_URL}/**",
|
||||
timeout=config.LOGIN_TIMEOUT * 1000
|
||||
)
|
||||
await self.emit_log("SUCCESS", "登录成功！")
|
||||
|
||||
# 保存登录状态
|
||||
await context.storage_state(path=config.AUTH_STATE_FILE)
|
||||
await self.emit_log("INFO", "登录状态已保存")
|
||||
|
||||
await page.close()
|
||||
return True
|
||||
|
||||
except PlaywrightTimeoutError:
|
||||
await self.emit_log("ERROR", "登录超时，请重试")
|
||||
await page.close()
|
||||
return False
|
||||
|
||||
async def copy_cookies(self, target_context: BrowserContext, source_context: BrowserContext):
|
||||
"""复制 Cookies 到目标 Context"""
|
||||
cookies = await source_context.cookies()
|
||||
await target_context.add_cookies(cookies)
|
||||
|
||||
async def check_risk_control(self, page: Page) -> tuple[bool, str]:
|
||||
"""
|
||||
检测页面是否触发风控
|
||||
返回 (是否触发, 触发原因)
|
||||
"""
|
||||
try:
|
||||
content = await page.content()
|
||||
for keyword in config.RISK_KEYWORDS:
|
||||
if keyword in content:
|
||||
return True, f"检测到关键词: {keyword}"
|
||||
return False, ""
|
||||
except Exception as e:
|
||||
return False, f"检测失败: {str(e)}"
|
||||
|
||||
async def human_like_scroll(self, page: Page, thread_id: Optional[str] = None):
|
||||
"""
|
||||
模拟人类间歇性平滑滚动
|
||||
"""
|
||||
await self.emit_log("DEBUG", f"[{thread_id}] 开始平滑滚动加载笔记...", thread_id=thread_id)
|
||||
|
||||
scroll_count = 0
|
||||
last_height = 0
|
||||
|
||||
while scroll_count < config.MAX_SCROLL_COUNT:
|
||||
# 获取当前页面高度
|
||||
current_height = await page.evaluate("document.body.scrollHeight")
|
||||
|
||||
# 如果高度没有变化，可能到底了
|
||||
if current_height == last_height:
|
||||
await self.emit_log("DEBUG", f"[{thread_id}] 已滚动到底部", thread_id=thread_id)
|
||||
break
|
||||
|
||||
last_height = current_height
|
||||
|
||||
# 随机滚动距离
|
||||
scroll_step = random.randint(config.SCROLL_STEP_MIN, config.SCROLL_STEP_MAX)
|
||||
scroll_position = await page.evaluate("window.scrollY")
|
||||
|
||||
# 平滑滚动
|
||||
await page.evaluate(f'''
|
||||
window.scrollTo({{
|
||||
top: {scroll_position + scroll_step},
|
||||
behavior: 'smooth'
|
||||
}});
|
||||
''')
|
||||
|
||||
# 随机暂停
|
||||
pause_time = random.uniform(config.SCROLL_PAUSE_MIN, config.SCROLL_PAUSE_MAX)
|
||||
await asyncio.sleep(pause_time)
|
||||
|
||||
scroll_count += 1
|
||||
|
||||
# 每滚动几次检查一下风控
|
||||
if scroll_count % 5 == 0:
|
||||
is_risk, reason = await self.check_risk_control(page)
|
||||
if is_risk:
|
||||
raise RiskControlException(reason)
|
||||
|
||||
await self.emit_log("DEBUG", f"[{thread_id}] 滚动完成，共滚动 {scroll_count} 次", thread_id=thread_id)
|
||||
|
||||
async def extract_note_cards(self, page: Page, thread_id: Optional[str] = None) -> List[Dict]:
|
||||
"""
|
||||
从用户主页提取笔记卡片
|
||||
返回 [{'url': '...', 'title': '...'}]
|
||||
"""
|
||||
await self.emit_log("DEBUG", f"[{thread_id}] 提取笔记卡片...", thread_id=thread_id)
|
||||
|
||||
note_urls = []
|
||||
|
||||
try:
|
||||
# 尝试多种选择器
|
||||
selectors = [
|
||||
"a[href*='/explore/']",
|
||||
"section[class*='note'] a[href*='/explore/']",
|
||||
"div[class*='note-item'] a[href*='/explore/']",
|
||||
]
|
||||
|
||||
links = []
|
||||
for selector in selectors:
|
||||
try:
|
||||
elements = await page.query_selector_all(selector)
|
||||
if elements:
|
||||
links = elements
|
||||
await self.emit_log("DEBUG", f"[{thread_id}] 使用选择器: {selector}", thread_id=thread_id)
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# 提取 URL
|
||||
for link in links:
|
||||
url = await link.get_attribute('href')
|
||||
if url and '/explore/' in url:
|
||||
# 补全 URL
|
||||
if url.startswith('/'):
|
||||
url = f"{config.XIAOHONGSHU_BASE_URL}{url}"
|
||||
|
||||
# 尝试获取标题
|
||||
try:
|
||||
title_elem = await link.query_selector("span, div")
|
||||
title = await title_elem.inner_text() if title_elem else "无标题"
|
||||
except Exception:
|
||||
title = "无标题"
|
||||
|
||||
note_urls.append({
|
||||
'url': url,
|
||||
'title': title[:50] # 限制标题长度
|
||||
})
|
||||
|
||||
await self.emit_log("INFO", f"[{thread_id}] 提取到 {len(note_urls)} 条笔记", thread_id=thread_id)
|
||||
|
||||
return note_urls
|
||||
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 提取笔记卡片失败: {str(e)}", thread_id=thread_id)
|
||||
return []
|
||||
|
||||
async def extract_note_detail(self, page: Page, thread_id: Optional[str] = None) -> Optional[Dict]:
|
||||
"""
|
||||
从笔记详情页提取数据
|
||||
"""
|
||||
try:
|
||||
# 等待页面加载
|
||||
await asyncio.sleep(random.uniform(*config.PAGE_LOAD_DELAY))
|
||||
|
||||
# 检查风控
|
||||
is_risk, reason = await self.check_risk_control(page)
|
||||
if is_risk:
|
||||
raise RiskControlException(reason)
|
||||
|
||||
data = {}
|
||||
|
||||
# 提取标题
|
||||
try:
|
||||
title_elem = await page.query_selector(config.SELECTOR_DETAIL_TITLE)
|
||||
data['title'] = await title_elem.inner_text() if title_elem else "无标题"
|
||||
except Exception:
|
||||
data['title'] = "无标题"
|
||||
|
||||
# 提取发布时间
|
||||
try:
|
||||
time_elem = await page.query_selector(config.SELECTOR_DETAIL_PUBLISH_TIME)
|
||||
data['publish_time'] = await time_elem.inner_text() if time_elem else ""
|
||||
except Exception:
|
||||
data['publish_time'] = ""
|
||||
|
||||
# 提取发布地
|
||||
try:
|
||||
loc_elem = await page.query_selector(config.SELECTOR_DETAIL_LOCATION)
|
||||
data['location'] = await loc_elem.inner_text() if loc_elem else ""
|
||||
except Exception:
|
||||
data['location'] = ""
|
||||
|
||||
# 提取互动数据（尝试从页面内容解析）
|
||||
content = await page.content()
|
||||
|
||||
# 评论数
|
||||
data['comment_count'] = self._extract_count(content, ['评论', 'comment'])
|
||||
|
||||
# 点赞数
|
||||
data['like_count'] = self._extract_count(content, ['点赞', '赞', 'like'])
|
||||
|
||||
# 收藏数
|
||||
data['collect_count'] = self._extract_count(content, ['收藏', 'star', 'collect'])
|
||||
|
||||
# 笔记链接
|
||||
data['note_url'] = page.url
|
||||
|
||||
return data
|
||||
|
||||
except RiskControlException:
|
||||
raise
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 提取详情失败: {str(e)}", thread_id=thread_id)
|
||||
return None
|
||||
|
||||
def _extract_count(self, content: str, keywords: List[str]) -> int:
|
||||
"""从页面内容提取计数"""
|
||||
import re
|
||||
for keyword in keywords:
|
||||
# 匹配数字
|
||||
pattern = f"{keyword}[\\s:：]*([\\d.]+)"
|
||||
match = re.search(pattern, content)
|
||||
if match:
|
||||
count_str = match.group(1)
|
||||
# 处理万、k等单位
|
||||
if '万' in count_str:
|
||||
return int(float(count_str.replace('万', '')) * 10000)
|
||||
elif 'k' in count_str.lower():
|
||||
return int(float(count_str.lower().replace('k', '')) * 1000)
|
||||
else:
|
||||
return int(count_str)
|
||||
return 0
|
||||
|
||||
async def save_user_data(self, user_id: str, new_notes: List[Dict]):
|
||||
"""保存用户数据（增量）"""
|
||||
file_path = os.path.join(config.DATA_DIR, f"{user_id}.json")
|
||||
|
||||
# 读取现有数据
|
||||
existing_notes = []
|
||||
if os.path.exists(file_path):
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
existing_notes = data.get('notes', [])
|
||||
except Exception as e:
|
||||
await self.emit_log("WARNING", f"读取现有数据失败: {e}")
|
||||
|
||||
# 提取现有 URL
|
||||
existing_urls = {note['note_url'] for note in existing_notes}
|
||||
|
||||
# 追加新笔记
|
||||
added_count = 0
|
||||
for note in new_notes:
|
||||
if note['note_url'] not in existing_urls:
|
||||
existing_notes.append(note)
|
||||
existing_urls.add(note['note_url'])
|
||||
added_count += 1
|
||||
|
||||
# 保存
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
json.dump({
|
||||
"user_id": user_id,
|
||||
"last_updated": datetime.now().isoformat(),
|
||||
"total_notes": len(existing_notes),
|
||||
"notes": existing_notes
|
||||
}, f, ensure_ascii=False, indent=2)
|
||||
|
||||
return added_count
|
||||
|
||||
async def process_user(self, context: BrowserContext, user_id: str, thread_id: str = "Main") -> int:
|
||||
"""
|
||||
处理单个用户
|
||||
返回新增笔记数
|
||||
"""
|
||||
await self.emit_user_status(user_id, "processing", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "busy", user_id)
|
||||
|
||||
try:
|
||||
await self.emit_log("INFO", f"[{thread_id}] 开始处理用户: {user_id}", thread_id=thread_id)
|
||||
|
||||
# 创建新页面
|
||||
page = await context.new_page()
|
||||
|
||||
# 访问用户主页
|
||||
url = config.USER_PROFILE_URL_TEMPLATE.format(user_id=user_id)
|
||||
await self.emit_log("INFO", f"[{thread_id}] 访问用户主页: {url}", thread_id=thread_id)
|
||||
|
||||
try:
|
||||
await page.goto(url, timeout=30000)
|
||||
except PlaywrightTimeoutError:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 访问用户主页超时: {user_id}", thread_id=thread_id)
|
||||
await page.close()
|
||||
await self.emit_user_status(user_id, "failed", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
return 0
|
||||
|
||||
# 检查风控
|
||||
is_risk, reason = await self.check_risk_control(page)
|
||||
if is_risk:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 访问用户主页触发风控: {reason}", thread_id=thread_id)
|
||||
await page.close()
|
||||
await self.emit_user_status(user_id, "failed", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
raise RiskControlException(reason)
|
||||
|
||||
await asyncio.sleep(random.uniform(*config.PAGE_LOAD_DELAY))
|
||||
|
||||
# 滚动加载所有笔记
|
||||
await self.human_like_scroll(page, thread_id)
|
||||
|
||||
# 提取笔记卡片
|
||||
note_cards = await self.extract_note_cards(page, thread_id)
|
||||
|
||||
if not note_cards:
|
||||
await self.emit_log("WARNING", f"[{thread_id}] 未找到笔记，可能页面结构变化", thread_id=thread_id)
|
||||
await page.close()
|
||||
await self.emit_user_status(user_id, "completed", 0, thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
return 0
|
||||
|
||||
# 采集每条笔记
|
||||
collected_notes = []
|
||||
for i, card in enumerate(note_cards):
|
||||
if self._stop_requested:
|
||||
await self.emit_log("WARNING", f"[{thread_id}] 收到停止信号", thread_id=thread_id)
|
||||
break
|
||||
|
||||
await self.emit_log("INFO", f"[{thread_id}] 采集笔记 {i+1}/{len(note_cards)}: {card['title']}", thread_id=thread_id)
|
||||
|
||||
try:
|
||||
# 访问详情页
|
||||
await page.goto(card['url'], timeout=30000)
|
||||
|
||||
# 模拟阅读停留
|
||||
await asyncio.sleep(random.uniform(*config.NOTE_READ_DELAY))
|
||||
|
||||
# 提取数据
|
||||
note_data = await self.extract_note_detail(page, thread_id)
|
||||
|
||||
if note_data:
|
||||
collected_notes.append(note_data)
|
||||
await self.emit_log("SUCCESS", f"[{thread_id}] 成功采集: {note_data['title'][:30]}", thread_id=thread_id)
|
||||
|
||||
# 返回列表页
|
||||
await page.go_back(url, wait_until="networkidle")
|
||||
|
||||
# 间隔
|
||||
await asyncio.sleep(random.uniform(*config.NOTE_INTERVAL_DELAY))
|
||||
|
||||
except RiskControlException:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 采集笔记触发风控", thread_id=thread_id)
|
||||
raise
|
||||
except Exception as e:
|
||||
await self.emit_log("WARNING", f"[{thread_id}] 采集笔记失败: {str(e)}", thread_id=thread_id)
|
||||
continue
|
||||
|
||||
await page.close()
|
||||
|
||||
# 保存数据
|
||||
if collected_notes:
|
||||
added_count = await self.save_user_data(user_id, collected_notes)
|
||||
await self.emit_log("SUCCESS", f"[{thread_id}] 用户 {user_id} 采集完成，新增 {added_count} 条笔记", thread_id=thread_id)
|
||||
await self.emit_user_status(user_id, "completed", len(collected_notes), thread_id=thread_id)
|
||||
else:
|
||||
await self.emit_log("INFO", f"[{thread_id}] 用户 {user_id} 无新笔记", thread_id=thread_id)
|
||||
await self.emit_user_status(user_id, "completed", 0, thread_id=thread_id)
|
||||
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
return len(collected_notes)
|
||||
|
||||
except RiskControlException:
|
||||
await self.emit_user_status(user_id, "failed", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "error", None)
|
||||
raise
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 处理用户 {user_id} 失败: {str(e)}", thread_id=thread_id)
|
||||
await self.emit_user_status(user_id, "failed", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "error", None)
|
||||
return 0
|
||||
|
||||
async def run_single_thread(self, user_ids: List[str]):
|
||||
"""单线程模式"""
|
||||
await self.emit_log("INFO", "启动单线程采集模式")
|
||||
|
||||
# 创建 Context 并登录
|
||||
context = await self.browser.new_context(
|
||||
viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
|
||||
user_agent=config.USER_AGENT
|
||||
)
|
||||
|
||||
# 检查是否有保存的登录状态
|
||||
if os.path.exists(config.AUTH_STATE_FILE):
|
||||
await self.emit_log("INFO", "加载已保存的登录状态")
|
||||
await context.add_cookies([json.load(open(config.AUTH_STATE_FILE))])
|
||||
else:
|
||||
# 需要登录
|
||||
login_success = await self.handle_login(context)
|
||||
if not login_success:
|
||||
await context.close()
|
||||
return
|
||||
|
||||
self.is_running = True
|
||||
|
||||
for i, user_id in enumerate(user_ids):
|
||||
if self._stop_requested:
|
||||
await self.emit_log("WARNING", "收到停止请求")
|
||||
break
|
||||
|
||||
await self.emit_log("INFO", f"处理用户 {i+1}/{len(user_ids)}: {user_id}")
|
||||
|
||||
try:
|
||||
await self.process_user(context, user_id, thread_id="Main")
|
||||
|
||||
# 用户切换延迟
|
||||
if i < len(user_ids) - 1:
|
||||
delay = random.uniform(*config.USER_SWITCH_DELAY)
|
||||
await self.emit_log("INFO", f"等待 {delay:.1f} 秒后处理下一个用户...")
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
except RiskControlException as e:
|
||||
await self.emit_log("ERROR", f"触发风控: {e.reason}，终止采集")
|
||||
break
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"处理用户 {user_id} 失败: {str(e)}")
|
||||
continue
|
||||
|
||||
await context.close()
|
||||
self.is_running = False
|
||||
await self.emit_log("SUCCESS", "单线程采集完成")
|
||||
|
||||
async def worker(self, context: BrowserContext, task_queue: asyncio.Queue, thread_id: str):
|
||||
"""
|
||||
Worker 处理任务（多线程模式）
|
||||
"""
|
||||
# 线程启动时错开
|
||||
await asyncio.sleep(int(thread_id.split('-')[1]) * config.THREAD_OFFSET)
|
||||
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
|
||||
while True:
|
||||
if self._stop_requested:
|
||||
break
|
||||
|
||||
try:
|
||||
# 获取任务
|
||||
user_id = await asyncio.wait_for(task_queue.get(), timeout=1.0)
|
||||
await self.emit_thread_status(thread_id, "busy", user_id)
|
||||
|
||||
try:
|
||||
# 线程间最小间隔
|
||||
await asyncio.sleep(random.uniform(
|
||||
config.THREAD_MIN_DELAY,
|
||||
config.THREAD_MIN_DELAY + 3
|
||||
))
|
||||
|
||||
await self.process_user(context, user_id, thread_id)
|
||||
|
||||
except RiskControlException as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 触发风控: {e.reason}，停止该线程", thread_id=thread_id)
|
||||
await self.emit_thread_status(thread_id, "error", None)
|
||||
break
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] 处理失败: {str(e)}", thread_id=thread_id)
|
||||
finally:
|
||||
task_queue.task_done()
|
||||
await self.emit_thread_status(thread_id, "idle", None)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except Exception as e:
|
||||
await self.emit_log("ERROR", f"[{thread_id}] Worker 错误: {str(e)}", thread_id=thread_id)
|
||||
break
|
||||
|
||||
async def run_multi_thread(self, user_ids: List[str], max_threads: int = 3):
|
||||
"""多线程模式"""
|
||||
await self.emit_log("INFO", f"启动多线程采集模式，线程数: {max_threads}")
|
||||
|
||||
# 创建多个 Context
|
||||
contexts = []
|
||||
for i in range(max_threads):
|
||||
ctx = await self.browser.new_context(
|
||||
viewport={'width': config.VIEWPORT_WIDTH, 'height': config.VIEWPORT_HEIGHT},
|
||||
user_agent=config.USER_AGENT
|
||||
)
|
||||
|
||||
# 第一个 Context 登录，其他复制 Cookies
|
||||
if i == 0:
|
||||
if os.path.exists(config.AUTH_STATE_FILE):
|
||||
await self.emit_log("INFO", "加载已保存的登录状态")
|
||||
await ctx.add_cookies([json.load(open(config.AUTH_STATE_FILE))])
|
||||
else:
|
||||
login_success = await self.handle_login(ctx)
|
||||
if not login_success:
|
||||
await ctx.close()
|
||||
return
|
||||
else:
|
||||
await self.copy_cookies(ctx, contexts[0])
|
||||
await self.emit_log("INFO", f"Thread-{i+1} 已复制登录状态")
|
||||
|
||||
contexts.append(ctx)
|
||||
|
||||
# 创建任务队列
|
||||
task_queue = asyncio.Queue()
|
||||
for user_id in user_ids:
|
||||
await task_queue.put(user_id)
|
||||
|
||||
self.is_running = True
|
||||
|
||||
# 创建 workers
|
||||
workers = [
|
||||
asyncio.create_task(
|
||||
self.worker(contexts[i], task_queue, f"Thread-{i+1}")
|
||||
)
|
||||
for i in range(max_threads)
|
||||
]
|
||||
|
||||
# 等待完成
|
||||
await task_queue.join()
|
||||
|
||||
# 取消 workers
|
||||
for worker in workers:
|
||||
worker.cancel()
|
||||
|
||||
# 关闭 contexts
|
||||
for ctx in contexts:
|
||||
await ctx.close()
|
||||
|
||||
self.is_running = False
|
||||
await self.emit_log("SUCCESS", "多线程采集完成")
|
||||
|
||||
def stop(self):
|
||||
"""请求停止采集"""
|
||||
self._stop_requested = True
|
||||
self.is_running = False
|
||||
|
||||
async def run(self, user_ids: List[str], mode: str = "single", max_threads: int = 3):
|
||||
"""
|
||||
运行采集
|
||||
:param user_ids: 用户 ID 列表
|
||||
:param mode: single / multi
|
||||
:param max_threads: 多线程模式下的线程数
|
||||
"""
|
||||
self._stop_requested = False
|
||||
|
||||
await self.emit_log("INFO", f"开始采集，共 {len(user_ids)} 个用户")
|
||||
|
||||
await self.init_browser()
|
||||
|
||||
try:
|
||||
if mode == "single":
|
||||
await self.run_single_thread(user_ids)
|
||||
else:
|
||||
await self.run_multi_thread(user_ids, max_threads)
|
||||
finally:
|
||||
await self.close_browser()
|
||||
|
||||
await self.emit_log("SUCCESS", "所有采集任务完成")
|
||||
Reference in New Issue
Block a user