feat: add --max-videos argument and wire recommendation flow in main
This commit is contained in:
parent
340293deba
commit
4c33f40289
15
Douyin.py
15
Douyin.py
@ -547,6 +547,12 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
default=DEFAULT_BROWSER_PORT,
|
default=DEFAULT_BROWSER_PORT,
|
||||||
help="附着到已启动 Chrome 的调试端口,默认 9223",
|
help="附着到已启动 Chrome 的调试端口,默认 9223",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--max-videos",
|
||||||
|
type=int,
|
||||||
|
default=50,
|
||||||
|
help="推荐流最大抓取数量,默认 50",
|
||||||
|
)
|
||||||
return parser
|
return parser
|
||||||
|
|
||||||
|
|
||||||
@ -560,6 +566,8 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
parser.error("--timeout 必须大于 0")
|
parser.error("--timeout 必须大于 0")
|
||||||
if args.browser_port is not None and args.browser_port <= 0:
|
if args.browser_port is not None and args.browser_port <= 0:
|
||||||
parser.error("--browser-port 必须大于 0")
|
parser.error("--browser-port 必须大于 0")
|
||||||
|
if args.max_videos <= 0:
|
||||||
|
parser.error("--max-videos 必须大于 0")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
target = resolve_cli_target(args.target, browser_port=args.browser_port)
|
target = resolve_cli_target(args.target, browser_port=args.browser_port)
|
||||||
@ -572,6 +580,13 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
browser_port=args.browser_port,
|
browser_port=args.browser_port,
|
||||||
auto_scroll=args.pages > 1,
|
auto_scroll=args.pages > 1,
|
||||||
)
|
)
|
||||||
|
elif target.kind == "recommendation":
|
||||||
|
total = collect_recommendations(
|
||||||
|
max_videos=args.max_videos,
|
||||||
|
timeout=args.timeout,
|
||||||
|
output_dir=Path(args.output_dir),
|
||||||
|
browser_port=args.browser_port,
|
||||||
|
)
|
||||||
elif target.kind == "single-video":
|
elif target.kind == "single-video":
|
||||||
total = collect_single_video(
|
total = collect_single_video(
|
||||||
target=target,
|
target=target,
|
||||||
|
|||||||
@ -404,6 +404,11 @@ class DouyinModuleTests(unittest.TestCase):
|
|||||||
self.assertEqual(page.visited_urls, ["https://www.douyin.com/video/7619989983668240802"])
|
self.assertEqual(page.visited_urls, ["https://www.douyin.com/video/7619989983668240802"])
|
||||||
mocked_download.assert_called_once()
|
mocked_download.assert_called_once()
|
||||||
|
|
||||||
|
def test_build_parser_has_max_videos_argument(self) -> None:
|
||||||
|
module = importlib.import_module("Douyin")
|
||||||
|
args = module.build_parser().parse_args(["--max-videos", "30"])
|
||||||
|
self.assertEqual(args.max_videos, 30)
|
||||||
|
|
||||||
def test_build_parser_defaults_to_zero_argument_current_page_flow(self) -> None:
|
def test_build_parser_defaults_to_zero_argument_current_page_flow(self) -> None:
|
||||||
module = importlib.import_module("Douyin")
|
module = importlib.import_module("Douyin")
|
||||||
args = module.build_parser().parse_args([])
|
args = module.build_parser().parse_args([])
|
||||||
@ -419,6 +424,26 @@ class DouyinModuleTests(unittest.TestCase):
|
|||||||
self.assertEqual(target.aweme_id, "7619989983668240802")
|
self.assertEqual(target.aweme_id, "7619989983668240802")
|
||||||
mocked_imports.assert_not_called()
|
mocked_imports.assert_not_called()
|
||||||
|
|
||||||
|
def test_main_dispatches_recommendation_flow_for_recommendation_url(self) -> None:
|
||||||
|
module = importlib.import_module("Douyin")
|
||||||
|
stdout = io.StringIO()
|
||||||
|
recommendation_target = module.ResolvedTarget(
|
||||||
|
kind="recommendation",
|
||||||
|
value="https://www.douyin.com/",
|
||||||
|
source="current-page",
|
||||||
|
)
|
||||||
|
with redirect_stdout(stdout):
|
||||||
|
with mock.patch.object(module, "resolve_cli_target", return_value=recommendation_target):
|
||||||
|
with mock.patch.object(module, "collect_recommendations", return_value=5) as mocked_collect:
|
||||||
|
exit_code = module.main([])
|
||||||
|
self.assertEqual(exit_code, 0)
|
||||||
|
mocked_collect.assert_called_once_with(
|
||||||
|
max_videos=50,
|
||||||
|
timeout=10,
|
||||||
|
output_dir=module.Path("video"),
|
||||||
|
browser_port=9223,
|
||||||
|
)
|
||||||
|
|
||||||
def test_main_without_target_dispatches_current_page_creator_flow(self) -> None:
|
def test_main_without_target_dispatches_current_page_creator_flow(self) -> None:
|
||||||
module = importlib.import_module("Douyin")
|
module = importlib.import_module("Douyin")
|
||||||
stdout = io.StringIO()
|
stdout = io.StringIO()
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user