124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
import argparse
|
|
import asyncio
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from . import 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__)
|
|
|
|
|
|
async def run_load_user_ratings_from_imdb():
|
|
await open_connection_pool()
|
|
|
|
i = 0
|
|
async for rating 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()
|
|
commands = parser.add_subparsers(required=True)
|
|
|
|
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",
|
|
)
|
|
|
|
try:
|
|
args = parser.parse_args()
|
|
except TypeError:
|
|
parser.print_usage()
|
|
raise
|
|
|
|
return args
|
|
|
|
|
|
def main():
|
|
logging.basicConfig(
|
|
format="%(asctime)s.%(msecs)03d [%(name)s:%(process)d] %(levelname)s: %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
level=config.loglevel,
|
|
)
|
|
log.debug(f"Log level: {config.loglevel}")
|
|
|
|
try:
|
|
args = getargs()
|
|
except:
|
|
return
|
|
|
|
if args.mode == "load-user-ratings-from-imdb":
|
|
asyncio.run(run_load_user_ratings_from_imdb())
|
|
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))
|
|
|
|
|
|
main()
|