import asyncio
import os
from playwright.async_api import async_playwright
class FileHandler:
def __init__(self, download_dir="./downloads"):
self.download_dir = download_dir
os.makedirs(download_dir, exist_ok=True)
async def wait_for_upload_completion(self, page, success_selector=None, timeout=30000):
"""等待文件上传完成"""
completion_indicators = [
lambda: page.wait_for_selector(success_selector, timeout=timeout) if success_selector else None,
lambda: page.wait_for_selector('.upload-complete', timeout=timeout),
lambda: page.wait_for_selector('text=上传成功', timeout=timeout),
lambda: page.wait_for_selector('text=Upload Complete', timeout=timeout)
]
for indicator in completion_indicators:
if indicator:
try:
await indicator()
return True
except:
continue
return False
async def wait_for_download_completion(self, page, download_trigger=None, timeout=60000):
"""等待文件下载完成"""
try:
# 监听下载开始
async with page.expect_download(timeout=timeout) as download_info:
if download_trigger:
await download_trigger()
download = await download_info.value
print(f"开始下载: {download.suggested_filename}")
# 等待下载完成
file_path = await download.path()
print(f"下载完成: {file_path}")
# 保存到指定目录
save_path = os.path.join(self.download_dir, download.suggested_filename)
await download.save_as(save_path)
# 验证文件
if self.validate_downloaded_file(save_path):
print(f"文件验证成功: {save_path}")
return save_path
else:
print(f"文件验证失败: {save_path}")
return None
except Exception as e:
print(f"下载过程出错: {e}")
return None
def validate_downloaded_file(self, file_path):
"""验证下载的文件"""
if not os.path.exists(file_path):
return False
file_size = os.path.getsize(file_path)
if file_size == 0:
return False
# 可以根据需要添加更多验证逻辑
# 例如文件类型、内容校验等
return True
# 使用示例
async def main():
handler = FileHandler()
async with async_playwright() as p:
browser = await p.chromium.launch(headless=False)
page = await browser.new_page()
# 文件上传示例
await page.goto('https://example.com/upload')
await page.set_input_files('input[type="file"]', 'test_file.txt')
await page.click('#upload-button')
# 等待上传完成
upload_success = await handler.wait_for_upload_completion(
page,
success_selector='.upload-success-message'
)
if upload_success:
print("文件上传成功")
else:
print("文件上传失败或超时")
# 文件下载示例
await page.goto('https://example.com/download')
download_path = await handler.wait_for_download_completion(
page,
download_trigger=lambda: page.click('#download-btn')
)
if download_path:
print(f"文件下载并保存到: {download_path}")
await browser.close()
if __name__ == "__main__":
asyncio.run(main())
|