fix archivebox shell and manage CLI commands

archivebox/cli/__init__.py

@@ -87,15 +87,18 @@ class ArchiveBoxGroup(click.Group):
 
 @click.group(cls=ArchiveBoxGroup, invoke_without_command=True)
 @click.option('--help', '-h', is_flag=True, help='Show help')
-@click.version_option(version=VERSION, package_name='archivebox', message='%(version)s')
+@click.version_option(VERSION, '-v', '--version', package_name='archivebox', message='%(version)s')
 @click.pass_context
 def cli(ctx, help=False):
     """ArchiveBox: The self-hosted internet archive"""
     
+    # if --help is passed or no subcommand is given, show custom help message
     if help or ctx.invoked_subcommand is None:
         ctx.invoke(ctx.command.get_command(ctx, 'help'))
     
-    if ctx.invoked_subcommand in ArchiveBoxGroup.archive_commands:
+    # if the subcommand is in the archive_commands dict and is not 'manage',
+    # then we need to set up the django environment and check that we're in a valid data folder
+    if ctx.invoked_subcommand in ArchiveBoxGroup.archive_commands and ctx.invoked_subcommand != 'manage':
         # print('SETUP DJANGO AND CHECK DATA FOLDER')
         from archivebox.config.django import setup_django
         from archivebox.misc.checks import check_data_folder

archivebox/cli/archivebox_manage.py

@@ -1,44 +1,33 @@
 #!/usr/bin/env python3
 
 __package__ = 'archivebox.cli'
-__command__ = 'archivebox manage'
 
-import sys
+import rich_click as click
-from pathlib import Path
+from archivebox.misc.util import docstring, enforce_types
-from typing import Optional, List, IO
-
-from archivebox.misc.util import docstring
-from archivebox.config import DATA_DIR
 
 
-
+@enforce_types
-# @enforce_types
+def manage(args: list[str] | None=None) -> None:
-def manage(args: Optional[List[str]]=None, out_dir: Path=DATA_DIR) -> None:
     """Run an ArchiveBox Django management command"""
 
-    check_data_folder()
+    from archivebox.config.common import SHELL_CONFIG
-    from django.core.management import execute_from_command_line
+    from archivebox.misc.logging import stderr
 
-    if (args and "createsuperuser" in args) and (IN_DOCKER and not SHELL_CONFIG.IS_TTY):
+    if (args and "createsuperuser" in args) and (SHELL_CONFIG.IN_DOCKER and not SHELL_CONFIG.IS_TTY):
         stderr('[!] Warning: you need to pass -it to use interactive commands in docker', color='lightyellow')
         stderr('    docker run -it archivebox manage {}'.format(' '.join(args or ['...'])), color='lightyellow')
         stderr('')
 
-    # import ipdb; ipdb.set_trace()
+    from django.core.management import execute_from_command_line
-
     execute_from_command_line(['manage.py', *(args or ['help'])])
 
 
-
+@click.command(add_help_option=False, context_settings=dict(ignore_unknown_options=True))
-
+@click.argument('args', nargs=-1)
-
 @docstring(manage.__doc__)
-def main(args: Optional[List[str]]=None, stdin: Optional[IO]=None, pwd: Optional[str]=None) -> None:
+def main(args: list[str] | None=None) -> None:
-    manage(
+    manage(args=args)
-        args=args,
-        out_dir=Path(pwd) if pwd else DATA_DIR,
-    )
 
 
 if __name__ == '__main__':
-    main(args=sys.argv[1:], stdin=sys.stdin)
+    main()

archivebox/cli/archivebox_shell.py

@@ -1,46 +1,27 @@
 #!/usr/bin/env python3
 
 __package__ = 'archivebox.cli'
-__command__ = 'archivebox shell'
 
-import sys
+from typing import Iterable
-import argparse
+
-from pathlib import Path
+import rich_click as click
-from typing import Optional, List, IO
 
 from archivebox.misc.util import docstring
-from archivebox.config import DATA_DIR
-from archivebox.misc.logging_util import SmartFormatter, reject_stdin
 
 
-
+def shell(args: Iterable[str]=()) -> None:
-#@enforce_types
-def shell(out_dir: Path=DATA_DIR) -> None:
     """Enter an interactive ArchiveBox Django shell"""
 
-    check_data_folder()
-
     from django.core.management import call_command
-    call_command("shell_plus")
+    call_command("shell_plus", *args)
-
-
 
 
+@click.command(add_help_option=False, context_settings=dict(ignore_unknown_options=True))
+@click.argument('args', nargs=-1)
 @docstring(shell.__doc__)
-def main(args: Optional[List[str]]=None, stdin: Optional[IO]=None, pwd: Optional[str]=None) -> None:
+def main(args: Iterable[str]=()) -> None:
-    parser = argparse.ArgumentParser(
+    shell(args=args)
-        prog=__command__,
-        description=shell.__doc__,
-        add_help=True,
-        formatter_class=SmartFormatter,
-    )
-    parser.parse_args(args or ())
-    reject_stdin(__command__, stdin)
-    
-    shell(
-        out_dir=Path(pwd) if pwd else DATA_DIR,
-    )
 
 
 if __name__ == '__main__':
-    main(args=sys.argv[1:], stdin=sys.stdin)
+    main()

archivebox/core/settings.py

@@ -223,8 +223,6 @@ MIGRATION_MODULES = {'signal_webhooks': None}
 # as much as I'd love this to be a UUID or ULID field, it's not supported yet as of Django 5.0
 DEFAULT_AUTO_FIELD = 'django.db.models.BigAutoField'
 
-
-if not IS_GETTING_VERSION_OR_HELP:             # dont create queue.sqlite3 file if we're just running to get --version or --help
 HUEY = {
     "huey_class": "huey.SqliteHuey",
     "filename": CONSTANTS.QUEUE_DATABASE_FILENAME,
@@ -257,6 +255,7 @@ if not IS_GETTING_VERSION_OR_HELP:             # dont create queue.sqlite3 file
     },
 }
 
+
 class HueyDBRouter:
     """
     A router to store all the Huey result k:v / Huey Monitor models in the queue.sqlite3 database.