From 6356bd53f8e320eafbdd40567e88580da1708ebe Mon Sep 17 00:00:00 2001 From: EdgeNeko Date: Fri, 5 Jul 2024 12:27:34 +0800 Subject: [PATCH] feat(breaking): Refactor CLI interface to be more clear and allow specifying categories when performing local indexing. --- .idea/NekoImageGallery.iml | 2 +- main.py | 152 +++++++++++++++++++++---------------- requirements.txt | 6 +- scripts/local_indexing.py | 10 +-- scripts/local_utility.py | 11 --- 5 files changed, 98 insertions(+), 83 deletions(-) delete mode 100644 scripts/local_utility.py diff --git a/.idea/NekoImageGallery.iml b/.idea/NekoImageGallery.iml index 1535ef5..8d250ed 100644 --- a/.idea/NekoImageGallery.iml +++ b/.idea/NekoImageGallery.iml @@ -8,6 +8,6 @@ - \ No newline at end of file diff --git a/main.py b/main.py index 4d41847..bc92634 100644 --- a/main.py +++ b/main.py @@ -1,71 +1,95 @@ -import argparse import asyncio +from pathlib import Path +from typing import Annotated, Optional +import rich +import typer import uvicorn - -def parse_args(): - parser = argparse.ArgumentParser(prog="NekoImageGallery Server", - description='Ciallo~ Welcome to NekoImageGallery Server.', - epilog="Build with ♥ By EdgeNeko. Github: " - "https://github.com/hv0905/NekoImageGallery") - - actions = parser.add_argument_group('Actions').add_mutually_exclusive_group() - - actions.add_argument('--show-config', action='store_true', help="Print the current configuration and exit.") - actions.add_argument('--init-database', action='store_true', - help="Initialize qdrant database using connection settings in " - "configuration. When this flag is set, will not" - "start the server.") - actions.add_argument('--migrate-db', dest="migrate_from_version", type=int, - help="Migrate qdrant database using connection settings in config from version specified." - "When this flag is set, will not start the server.") - actions.add_argument('--local-index', dest="local_index_target_dir", type=str, - help="Index all the images in this directory and copy them to " - "static folder set in config.py. When this flag is set, " - "will not start the server.") - actions.add_argument('--local-create-thumbnail', action='store_true', - help='Create thumbnail for all local images in static folder set in config.py. When this flag ' - 'is set, will not start the server.') - - server_options = parser.add_argument_group("Server Options") - - server_options.add_argument('--port', type=int, default=8000, help="Port to listen on, default is 8000") - server_options.add_argument('--host', type=str, default="0.0.0.0", help="Host to bind on, default is 0.0.0.0") - server_options.add_argument('--root-path', type=str, default="", - help="Root path of the server if your server is deployed behind a reverse proxy. " - "See https://fastapi.tiangolo.com/advanced/behind-a-proxy/ for detail.") - - parser.add_argument('--version', action='version', version='%(prog)s v1.1.0') - return parser.parse_args() +__version__ = '1.2.0' + +parser = typer.Typer(name='NekoImageGallery Server', + help='Ciallo~ Welcome to NekoImageGallery Server.\n\nBy default, running without command will ' + 'start the server. You can perform other actions by using the commands below.', + epilog="Build with ♥ By EdgeNeko. Github: " + "https://github.com/hv0905/NekoImageGallery" + ) + + +def version_callback(value: bool): + if value: + print(f"NekoImageGallery v{__version__}") + raise typer.Exit() + + +@parser.callback(invoke_without_command=True) +def server(ctx: typer.Context, + host: Annotated[str, typer.Option(help='The host to bind on.')] = '0.0.0.0', + port: Annotated[int, typer.Option(help='The port to listen on.')] = 8000, + root_path: Annotated[str, typer.Option( + help='Root path of the server if your server is deployed behind a reverse proxy. See ' + 'https://fastapi.tiangolo.com/advanced/behind-a-proxy/ for detail.')] = '', + _: Annotated[ + Optional[bool], typer.Option("--version", callback=version_callback, is_eager=True, + help="Show version and exit.") + ] = None + ): + """ + Start the server with the specified host, port and root path. + """ + if ctx.invoked_subcommand is not None: + return + uvicorn.run("app.webapp:app", host=host, port=port, root_path=root_path) + + +@parser.command('show-config') +def show_config(): + """ + Print the current configuration and exit. + """ + from app.config import config + rich.print_json(config.model_dump_json()) + + +@parser.command('init-database') +def init_database(): + """ + Initialize qdrant database using connection settings in configuration. + Note. The server will automatically initialize the database if it's not initialized. So you don't need to run this + command unless you want to explicitly initialize the database. + """ + from scripts import qdrant_create_collection + asyncio.run(qdrant_create_collection.main()) + + +@parser.command("local-index") +def local_index( + target_dir: Annotated[ + Path, typer.Argument(dir_okay=True, file_okay=False, exists=True, resolve_path=True, readable=True, + help="The directory to index.")], + categories: Annotated[Optional[list[str]], typer.Option(help="Categories for the indexed images.")] = None, + starred: Annotated[bool, typer.Option(help="Whether the indexed images are starred.")] = False, +): + """ + Index all the images in the specified directory. + The images will be copied to the local storage directory set in configuration. + """ + from scripts import local_indexing + if categories is None: + categories = [] + asyncio.run(local_indexing.main(target_dir, categories, starred)) + + +@parser.command('local-create-thumbnail', deprecated=True) +def local_create_thumbnail(): + """ + Create thumbnail for all local images in static folder, this won't affect non-local images. + This is generally not required since the server will automatically create thumbnails for new images by default. + This option will be refactored in the future. + """ + from scripts import local_create_thumbnail + asyncio.run(local_create_thumbnail.main()) if __name__ == '__main__': - args = parse_args() - if args.show_config: - from app.config import config - - print(config.model_dump_json(indent=2)) - elif args.init_database: - from scripts import qdrant_create_collection - from app.config import config - - asyncio.run(qdrant_create_collection.main()) - - elif args.migrate_from_version is not None: - from scripts import db_migrations - - asyncio.run(db_migrations.migrate(args.migrate_from_version)) - elif args.local_index_target_dir is not None: - from app.config import environment - - environment.local_indexing = True - from scripts import local_indexing - - asyncio.run(local_indexing.main(args.local_index_target_dir)) - elif args.local_create_thumbnail: - from scripts import local_create_thumbnail - - asyncio.run(local_create_thumbnail.main()) - else: - uvicorn.run("app.webapp:app", host=args.host, port=args.port, root_path=args.root_path) + parser() diff --git a/requirements.txt b/requirements.txt index d92992f..b4f416d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ python-multipart>=0.0.9 uvicorn[standard] pydantic pydantic-settings - +typer # AI - Manually install cuda-capable pytorch torch>=2.1.0 @@ -30,4 +30,6 @@ aiopath wcmatch pyyaml loguru -httpx \ No newline at end of file +httpx +pytest +rich \ No newline at end of file diff --git a/scripts/local_indexing.py b/scripts/local_indexing.py index 0b1db26..b25aa01 100644 --- a/scripts/local_indexing.py +++ b/scripts/local_indexing.py @@ -13,13 +13,13 @@ services: ServiceProvider | None = None -async def index_task(file_path: Path): +async def index_task(file_path: Path, categories: list[str], starred: bool): try: img_id = await services.upload_service.assign_image_id(file_path) image_data = ImageData(id=img_id, local=True, - categories=[], - starred=False, + categories=categories, + starred=starred, format=file_path.suffix[1:], # remove the dot index_date=datetime.now()) await services.upload_service.sync_upload_image(image_data, file_path.read_bytes(), skip_ocr=False, @@ -31,7 +31,7 @@ async def index_task(file_path: Path): @logger.catch() -async def main(root_directory): +async def main(root_directory: Path, categories: list[str], starred: bool): global services services = ServiceProvider() await services.onload() @@ -39,6 +39,6 @@ async def main(root_directory): for idx, item in enumerate(glob_local_files(root, '**/*')): logger.info("[{}] Indexing {}", idx, str(item)) - await index_task(item) + await index_task(item, categories, starred) logger.success("Indexing completed!") diff --git a/scripts/local_utility.py b/scripts/local_utility.py deleted file mode 100644 index 669fdca..0000000 --- a/scripts/local_utility.py +++ /dev/null @@ -1,11 +0,0 @@ -from pathlib import Path -from app.util.generate_uuid import generate_uuid - - -def calculate_uuid(file_path: Path) -> str: - return str(generate_uuid(file_path)) - - -def fetch_path_uuid_list(file_path: Path | list[Path]) -> list[tuple[Path, str]]: - file_path = [file_path] if isinstance(file_path, Path) else file_path - return [(itm, calculate_uuid(itm)) for itm in file_path]