From 3caf3f0e60dd200fd7dff2cfa1bfe92c53d97970 Mon Sep 17 00:00:00 2001 From: devdudumuniz <82589615+devdudumuniz@users.noreply.github.com> Date: Tue, 1 Sep 2026 07:01:59 +0000 Subject: [PATCH] Extracted nested functions from `check` command in `pysus/cli/management.py` to module level functions `_do_check`, `_do_apply`, `_flush`, `_print_check_report`, and `_print_check_json`. This flattens the long function and improves readability and maintainability. --- pysus/cli/management.py | 261 ++++++++++++++++++++++------------------ 1 file changed, 146 insertions(+), 115 deletions(-) diff --git a/pysus/cli/management.py b/pysus/cli/management.py index b4eb1f19..e35013b6 100644 --- a/pysus/cli/management.py +++ b/pysus/cli/management.py @@ -63,6 +63,138 @@ def _journal_path( return None +def _flush() -> None: + import sys + + sys.stdout.flush() + sys.stderr.flush() + + +def _print_check_report(checks) -> None: + total_missing = sum(len(c.missing) for c in checks.values()) + total_outdated = sum(len(c.outdated) for c in checks.values()) + total_current = sum(len(c.current) for c in checks.values()) + + typer.echo( + f"{'DATABASE':<30}{'MISSING':>9}{'OUTDATED':>10}" + f"{'CURRENT':>9} STATUS" + ) + typer.echo("-" * 78) + for ds in sorted(checks): + c = checks[ds] + verdict = "needs update" if c.needs_update else "up to date" + typer.echo( + f"{ds:<30}{len(c.missing):>9}{len(c.outdated):>10}" + f"{len(c.current):>9} {verdict}" + ) + typer.echo("-" * 78) + typer.echo( + f"{'TOTAL':<30}{total_missing:>9}{total_outdated:>10}" + f"{total_current:>9}" + ) + if total_missing or total_outdated: + typer.echo( + f"\n{total_missing + total_outdated} of " + f"{total_missing + total_outdated + total_current} file(s) " + "need updating — re-run with --apply to mirror them." + ) + + +def _print_check_json(checks) -> None: + for ds in sorted(checks): + typer.echo(json.dumps({"dataset": ds, **checks[ds].summary()})) + + +async def _do_check(engine, datasets, json_out: bool) -> None: + await engine.__aenter__(lock=False) + try: + checks = await engine.check(datasets=datasets) + finally: + await engine.__aexit__(None, None, None) + _flush() + if json_out: + _print_check_json(checks) + else: + _print_check_report(checks) + _flush() + + +async def _do_apply( + engine, + datasets, + force: bool, + reupload_before: str | None, + workers: int, + ftp_connections: int, + checkpoint_every: int, + resume_keys: set, + journal: Path | None, + json_out: bool, +) -> dict[str, int]: + counts: dict[str, int] = {} + + def on_outcome(outcome) -> None: + counts[outcome.status] = counts.get(outcome.status, 0) + 1 + if json_out: + typer.echo( + json.dumps( + { + "status": outcome.status, + "dataset": outcome.key.dataset, + "group": outcome.key.group, + "year": outcome.key.year, + "month": outcome.key.month, + "state": outcome.key.state, + "stem": outcome.key.stem, + "origin": outcome.origin, + "detail": outcome.detail, + }, + default=str, + ) + ) + return + if outcome.status == "needs_update": + typer.echo(f"[needs_update] {outcome.detail}") + elif outcome.status in ("uploaded", "failed", "needs_token"): + typer.echo(f"[{outcome.status}] {outcome.detail}") + total = sum(counts.values()) + if total % 500 == 0: + typer.echo(f"progress: {counts}", err=True) + + async with engine: + report = await engine.run( + datasets=datasets, + force=force, + reupload_before=( + _parse_date(reupload_before) if reupload_before else None + ), + dry_run=False, + workers=workers, + ftp_connections=ftp_connections, + checkpoint_every=checkpoint_every, + on_outcome=on_outcome, + resume=resume_keys or None, + journal=journal, + ) + summary = report.summary() + _flush() + if json_out: + typer.echo(json.dumps({"summary": summary}), err=True) + _flush() + return summary + typer.echo( + "\n" + f" total: {summary['total']}\n" + f" needs_update: {summary['needs_update']}\n" + f" uploaded: {summary['uploaded']}\n" + f" skipped: {summary['skipped']}\n" + f" failed: {summary['failed']}\n" + f" needs_token: {summary['needs_token']}" + ) + _flush() + return summary + + @app.command() def check( name: list[str] = typer.Argument( # noqa: B008 @@ -137,60 +269,8 @@ def check( datasets = [d.upper() for d in name] if name else None - def _flush() -> None: - import sys - - sys.stdout.flush() - sys.stderr.flush() - - def _print_check_report(checks) -> None: - total_missing = sum(len(c.missing) for c in checks.values()) - total_outdated = sum(len(c.outdated) for c in checks.values()) - total_current = sum(len(c.current) for c in checks.values()) - - typer.echo( - f"{'DATABASE':<30}{'MISSING':>9}{'OUTDATED':>10}" - f"{'CURRENT':>9} STATUS" - ) - typer.echo("-" * 78) - for ds in sorted(checks): - c = checks[ds] - verdict = "needs update" if c.needs_update else "up to date" - typer.echo( - f"{ds:<30}{len(c.missing):>9}{len(c.outdated):>10}" - f"{len(c.current):>9} {verdict}" - ) - typer.echo("-" * 78) - typer.echo( - f"{'TOTAL':<30}{total_missing:>9}{total_outdated:>10}" - f"{total_current:>9}" - ) - if total_missing or total_outdated: - typer.echo( - f"\n{total_missing + total_outdated} of " - f"{total_missing + total_outdated + total_current} file(s) " - "need updating — re-run with --apply to mirror them." - ) - - def _print_check_json(checks) -> None: - for ds in sorted(checks): - typer.echo(json.dumps({"dataset": ds, **checks[ds].summary()})) - - async def _run_check() -> None: - await engine.__aenter__(lock=False) - try: - checks = await engine.check(datasets=datasets) - finally: - await engine.__aexit__(None, None, None) - _flush() - if json_out: - _print_check_json(checks) - else: - _print_check_report(checks) - _flush() - if not apply: - _run_sync(_run_check()) + _run_sync(_do_check(engine, datasets, json_out)) return journal = _journal_path(resume, reupload_before) @@ -198,68 +278,19 @@ async def _run_check() -> None: if journal is not None and journal.exists(): resume_keys = load_journal_keys(journal) - counts: dict[str, int] = {} - - def on_outcome(outcome) -> None: - counts[outcome.status] = counts.get(outcome.status, 0) + 1 - if json_out: - typer.echo( - json.dumps( - { - "status": outcome.status, - "dataset": outcome.key.dataset, - "group": outcome.key.group, - "year": outcome.key.year, - "month": outcome.key.month, - "state": outcome.key.state, - "stem": outcome.key.stem, - "origin": outcome.origin, - "detail": outcome.detail, - }, - default=str, - ) - ) - return - if outcome.status == "needs_update": - typer.echo(f"[needs_update] {outcome.detail}") - elif outcome.status in ("uploaded", "failed", "needs_token"): - typer.echo(f"[{outcome.status}] {outcome.detail}") - total = sum(counts.values()) - if total % 500 == 0: - typer.echo(f"progress: {counts}", err=True) - - async def _run() -> dict[str, int]: - async with engine: - report = await engine.run( - datasets=datasets, - force=force, - reupload_before=_parse_date(reupload_before), - dry_run=False, - workers=workers, - ftp_connections=ftp_connections, - checkpoint_every=checkpoint_every, - on_outcome=on_outcome, - resume=resume_keys or None, - journal=journal, - ) - summary = report.summary() - _flush() - if json_out: - typer.echo(json.dumps({"summary": summary}), err=True) - _flush() - return summary - typer.echo( - "\n" - f" total: {summary['total']}\n" - f" needs_update: {summary['needs_update']}\n" - f" uploaded: {summary['uploaded']}\n" - f" skipped: {summary['skipped']}\n" - f" failed: {summary['failed']}\n" - f" needs_token: {summary['needs_token']}" + summary = _run_sync( + _do_apply( + engine=engine, + datasets=datasets, + force=force, + reupload_before=reupload_before, + workers=workers, + ftp_connections=ftp_connections, + checkpoint_every=checkpoint_every, + resume_keys=resume_keys, + journal=journal, + json_out=json_out, ) - _flush() - return summary - - summary = _run_sync(_run()) + ) if summary and summary["failed"]: raise typer.Exit(code=1)