chore: move existing CLI commands into separate files
This commit is contained in:
parent
f7fc84c050
commit
22c44bfa60
5 changed files with 147 additions and 145 deletions
|
|
@ -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))
|
||||
|
|
|
|||
56
unwind/cli/add_user.py
Normal file
56
unwind/cli/add_user.py
Normal file
|
|
@ -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()
|
||||
24
unwind/cli/download_imdb_dataset.py
Normal file
24
unwind/cli/download_imdb_dataset.py
Normal file
|
|
@ -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)
|
||||
31
unwind/cli/import_imdb_dataset.py
Normal file
31
unwind/cli/import_imdb_dataset.py
Normal file
|
|
@ -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()
|
||||
28
unwind/cli/load_user_ratings_from_imdb.py
Normal file
28
unwind/cli/load_user_ratings_from_imdb.py
Normal file
|
|
@ -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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue