diff --git a/unwind/__main__.py b/unwind/__main__.py index 8ca995f..0d5bde6 100644 --- a/unwind/__main__.py +++ b/unwind/__main__.py @@ -1,153 +1,25 @@ import argparse import asyncio import logging -import secrets import sys -from base64 import b64encode -from pathlib import Path -from . import cli, config, db, models, utils -from .db import close_connection_pool, open_connection_pool -from .imdb import refresh_user_ratings_from_imdb -from .imdb_import import download_datasets, import_from_file +from . import cli, config log = logging.getLogger(__name__) -async def run_add_user(user_id: str, name: str, overwrite_existing: bool): - if not user_id.startswith("ur"): - raise ValueError(f"Invalid IMDb user ID: {user_id!a}") - - await open_connection_pool() - - async with db.new_connection() as conn: - user = await db.get(conn, models.User, imdb_id=user_id) - - if user is not None: - if overwrite_existing: - log.warning("⚠️ Overwriting existing user: %a", user) - else: - log.error("❌ User already exists: %a", user) - return - - secret = secrets.token_bytes() - - user = models.User(name=name, imdb_id=user_id, secret=utils.phc_scrypt(secret)) - async with db.transaction() as conn: - await db.add_or_update_user(conn, user) - - user_data = { - "secret": b64encode(secret), - "user": models.asplain(user), - } - - log.info("✨ User created: %a", user_data) - - await close_connection_pool() - - -async def run_load_user_ratings_from_imdb(): - await open_connection_pool() - - i = 0 - async for _ in refresh_user_ratings_from_imdb(): - i += 1 - - log.info("✨ Imported %s new ratings.", i) - - await close_connection_pool() - - -async def run_import_imdb_dataset(basics_path: Path, ratings_path: Path): - await open_connection_pool() - - await import_from_file(basics_path=basics_path, ratings_path=ratings_path) - - await close_connection_pool() - - -async def run_download_imdb_dataset(basics_path: Path, ratings_path: Path): - await download_datasets(basics_path=basics_path, ratings_path=ratings_path) - - def getargs(): parser = argparse.ArgumentParser(prog="unwind", allow_abbrev=False) commands = parser.add_subparsers(title="commands", metavar="COMMAND", dest="mode") - parser_import_imdb_dataset = commands.add_parser( - "import-imdb-dataset", - help="Import IMDb datasets.", - description=""" - Import IMDb datasets. - New datasets available from https://www.imdb.com/interfaces/. - """, - ) - parser_import_imdb_dataset.add_argument( - dest="mode", - action="store_const", - const="import-imdb-dataset", - ) - parser_import_imdb_dataset.add_argument( - "--basics", metavar="basics_file.tsv.gz", type=Path, required=True - ) - parser_import_imdb_dataset.add_argument( - "--ratings", metavar="ratings_file.tsv.gz", type=Path, required=True - ) - - parser_download_imdb_dataset = commands.add_parser( - "download-imdb-dataset", - help="Download IMDb datasets.", - description=""" - Download IMDb datasets. - """, - ) - parser_download_imdb_dataset.add_argument( - dest="mode", - action="store_const", - const="download-imdb-dataset", - ) - parser_download_imdb_dataset.add_argument( - "--basics", metavar="basics_file.tsv.gz", type=Path, required=True - ) - parser_download_imdb_dataset.add_argument( - "--ratings", metavar="ratings_file.tsv.gz", type=Path, required=True - ) - - parser_load_user_ratings_from_imdb = commands.add_parser( - "load-user-ratings-from-imdb", - help="Load user ratings from imdb.com.", - description=""" - Refresh user ratings for all registered users live from IMDb's website. - """, - ) - parser_load_user_ratings_from_imdb.add_argument( - dest="mode", - action="store_const", - const="load-user-ratings-from-imdb", - ) - - parser_add_user = commands.add_parser( - "add-user", - help="Add a new user.", - description=""" - Add a new user. - """, - ) - parser_add_user.add_argument( - dest="mode", - action="store_const", - const="add-user", - ) - parser_add_user.add_argument("--name", required=True) - parser_add_user.add_argument("--imdb-id", required=True) - parser_add_user.add_argument( - "--overwrite-existing", - action="store_true", - help="Allow overwriting an existing user. WARNING: This will reset the user's password!", - ) - for module in cli.modules: - cmd = commands.add_parser(module.name, help=module.help, allow_abbrev=False) + help_, *descr = module.help.splitlines() + cmd = commands.add_parser( + module.name, + help=help_, + description="\n".join(descr) or help_, + allow_abbrev=False, + ) module.add_args(cmd) try: @@ -173,15 +45,6 @@ def main(): args = getargs() - if args.mode == "load-user-ratings-from-imdb": - asyncio.run(run_load_user_ratings_from_imdb()) - elif args.mode == "add-user": - asyncio.run(run_add_user(args.imdb_id, args.name, args.overwrite_existing)) - elif args.mode == "import-imdb-dataset": - asyncio.run(run_import_imdb_dataset(args.basics, args.ratings)) - elif args.mode == "download-imdb-dataset": - asyncio.run(run_download_imdb_dataset(args.basics, args.ratings)) - modes = {m.name: m.main for m in cli.modules} if handler := modes.get(args.mode): asyncio.run(handler(args)) diff --git a/unwind/cli/add_user.py b/unwind/cli/add_user.py new file mode 100644 index 0000000..cc3d305 --- /dev/null +++ b/unwind/cli/add_user.py @@ -0,0 +1,56 @@ +import argparse +import logging +import secrets + +from unwind import db, models, utils + +log = logging.getLogger(__name__) + +name = "add-user" +help = "Add a new user." + + +def add_args(cmd: argparse.ArgumentParser) -> None: + cmd.add_argument("--name", required=True) + cmd.add_argument("--imdb-id", required=True) + cmd.add_argument( + "--overwrite-existing", + action="store_true", + help="Allow overwriting an existing user. WARNING: This will reset the user's password!", + ) + + +async def main(args: argparse.Namespace) -> None: + user_id: str = args.imdb_id + name: str = args.name + overwrite_existing: bool = args.overwrite_existing + + if not user_id.startswith("ur"): + raise ValueError(f"Invalid IMDb user ID: {user_id!a}") + + await db.open_connection_pool() + + async with db.new_connection() as conn: + user = await db.get(conn, models.User, imdb_id=user_id) + + if user is not None: + if overwrite_existing: + log.warning("⚠️ Overwriting existing user: %a", user) + else: + log.error("❌ User already exists: %a", user) + return + + secret = secrets.token_bytes() + + user = models.User(name=name, imdb_id=user_id, secret=utils.phc_scrypt(secret)) + async with db.transaction() as conn: + await db.add_or_update_user(conn, user) + + user_data = { + "secret": utils.b64encode(secret), + "user": models.asplain(user), + } + + log.info("✨ User created: %a", user_data) + + await db.close_connection_pool() diff --git a/unwind/cli/download_imdb_dataset.py b/unwind/cli/download_imdb_dataset.py new file mode 100644 index 0000000..5b81045 --- /dev/null +++ b/unwind/cli/download_imdb_dataset.py @@ -0,0 +1,24 @@ +import argparse +import logging +from pathlib import Path + +from unwind.imdb_import import download_datasets + +log = logging.getLogger(__name__) + +name = "download-imdb-dataset" +help = "Download IMDb datasets." + + +def add_args(cmd: argparse.ArgumentParser) -> None: + cmd.add_argument("--basics", metavar="basics_file.tsv.gz", type=Path, required=True) + cmd.add_argument( + "--ratings", metavar="ratings_file.tsv.gz", type=Path, required=True + ) + + +async def main(args: argparse.Namespace) -> None: + basics_path: Path = args.basics + ratings_path: Path = args.ratings + + await download_datasets(basics_path=basics_path, ratings_path=ratings_path) diff --git a/unwind/cli/import_imdb_dataset.py b/unwind/cli/import_imdb_dataset.py new file mode 100644 index 0000000..3adb5da --- /dev/null +++ b/unwind/cli/import_imdb_dataset.py @@ -0,0 +1,31 @@ +import argparse +import logging +from pathlib import Path + +from unwind import db +from unwind.imdb_import import import_from_file + +log = logging.getLogger(__name__) + +name = "import-imdb-dataset" +help = """Import IMDb datasets. +New datasets available from https://www.imdb.com/interfaces/. +""" + + +def add_args(cmd: argparse.ArgumentParser) -> None: + cmd.add_argument("--basics", metavar="basics_file.tsv.gz", type=Path, required=True) + cmd.add_argument( + "--ratings", metavar="ratings_file.tsv.gz", type=Path, required=True + ) + + +async def main(args: argparse.Namespace) -> None: + basics_path: Path = args.basics + ratings_path: Path = args.ratings + + await db.open_connection_pool() + + await import_from_file(basics_path=basics_path, ratings_path=ratings_path) + + await db.close_connection_pool() diff --git a/unwind/cli/load_user_ratings_from_imdb.py b/unwind/cli/load_user_ratings_from_imdb.py new file mode 100644 index 0000000..b4a8e0f --- /dev/null +++ b/unwind/cli/load_user_ratings_from_imdb.py @@ -0,0 +1,28 @@ +import argparse +import logging + +from unwind import db +from unwind.imdb import refresh_user_ratings_from_imdb + +log = logging.getLogger(__name__) + +name = "load-user-ratings-from-imdb" +help = """Load user ratings from imdb.com. +Refresh user ratings for all registered users live from IMDb's website. +""" + + +def add_args(cmd: argparse.ArgumentParser) -> None: + pass + + +async def main(args: argparse.Namespace) -> None: + await db.open_connection_pool() + + i = 0 + async for _ in refresh_user_ratings_from_imdb(): + i += 1 + + log.info("✨ Imported %s new ratings.", i) + + await db.close_connection_pool()