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 argparse
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import secrets
|
|
||||||
import sys
|
import sys
|
||||||
from base64 import b64encode
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
from . import cli, config, db, models, utils
|
from . import cli, config
|
||||||
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
|
|
||||||
|
|
||||||
log = logging.getLogger(__name__)
|
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():
|
def getargs():
|
||||||
parser = argparse.ArgumentParser(prog="unwind", allow_abbrev=False)
|
parser = argparse.ArgumentParser(prog="unwind", allow_abbrev=False)
|
||||||
commands = parser.add_subparsers(title="commands", metavar="COMMAND", dest="mode")
|
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:
|
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)
|
module.add_args(cmd)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
|
@ -173,15 +45,6 @@ def main():
|
||||||
|
|
||||||
args = getargs()
|
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}
|
modes = {m.name: m.main for m in cli.modules}
|
||||||
if handler := modes.get(args.mode):
|
if handler := modes.get(args.mode):
|
||||||
asyncio.run(handler(args))
|
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