diff --git a/src/documents/management/commands/base.py b/src/documents/management/commands/base.py index a2b0880f5..1d76460c0 100644 --- a/src/documents/management/commands/base.py +++ b/src/documents/management/commands/base.py @@ -304,7 +304,7 @@ class PaperlessCommand(RichCommand): Progress output is directed to stderr to match the convention that progress bars are transient UI feedback, not command output. This - mirrors tqdm's default behavior and prevents progress bar rendering + mirrors the convention that progress bars are transient UI feedback and prevents progress bar rendering from interfering with stdout-based assertions in tests or piped command output. diff --git a/src/documents/management/commands/document_exporter.py b/src/documents/management/commands/document_exporter.py index 52cadb4fe..f53feb0e3 100644 --- a/src/documents/management/commands/document_exporter.py +++ b/src/documents/management/commands/document_exporter.py @@ -8,7 +8,6 @@ from itertools import islice from pathlib import Path from typing import TYPE_CHECKING -import tqdm from allauth.mfa.models import Authenticator from allauth.socialaccount.models import SocialAccount from allauth.socialaccount.models import SocialApp @@ -19,7 +18,6 @@ from django.contrib.auth.models import Permission from django.contrib.auth.models import User from django.contrib.contenttypes.models import ContentType from django.core import serializers -from django.core.management.base import BaseCommand from django.core.management.base import CommandError from django.core.serializers.json import DjangoJSONEncoder from django.db import transaction @@ -38,6 +36,7 @@ if settings.AUDIT_LOG_ENABLED: from documents.file_handling import delete_empty_directories from documents.file_handling import generate_filename +from documents.management.commands.base import PaperlessCommand from documents.management.commands.mixins import CryptMixin from documents.models import Correspondent from documents.models import CustomField @@ -81,14 +80,17 @@ def serialize_queryset_batched( yield serializers.serialize("python", chunk) -class Command(CryptMixin, BaseCommand): +class Command(CryptMixin, PaperlessCommand): help = ( "Decrypt and rename all files in our collection into a given target " "directory. And include a manifest file containing document data for " "easy import." ) + supports_progress_bar = True + def add_arguments(self, parser) -> None: + super().add_arguments(parser) parser.add_argument("target") parser.add_argument( @@ -195,13 +197,6 @@ class Command(CryptMixin, BaseCommand): help="If set, only the database will be imported, not files", ) - parser.add_argument( - "--no-progress-bar", - default=False, - action="store_true", - help="If set, the progress bar will not be shown", - ) - parser.add_argument( "--passphrase", help="If provided, is used to encrypt sensitive data in the export", @@ -230,7 +225,6 @@ class Command(CryptMixin, BaseCommand): self.no_thumbnail: bool = options["no_thumbnail"] self.zip_export: bool = options["zip"] self.data_only: bool = options["data_only"] - self.no_progress_bar: bool = options["no_progress_bar"] self.passphrase: str | None = options.get("passphrase") self.batch_size: int = options["batch_size"] @@ -347,10 +341,12 @@ class Command(CryptMixin, BaseCommand): document_manifest = manifest_dict["documents"] # 3. Export files from each document - for index, document_dict in tqdm.tqdm( - enumerate(document_manifest), - total=len(document_manifest), - disable=self.no_progress_bar, + for index, document_dict in enumerate( + self.track( + document_manifest, + description="Exporting documents...", + total=len(document_manifest), + ), ): document = document_map[document_dict["pk"]] diff --git a/src/documents/management/commands/document_importer.py b/src/documents/management/commands/document_importer.py index 5cd743590..68ff49430 100644 --- a/src/documents/management/commands/document_importer.py +++ b/src/documents/management/commands/document_importer.py @@ -8,14 +8,12 @@ from pathlib import Path from zipfile import ZipFile from zipfile import is_zipfile -import tqdm from django.conf import settings from django.contrib.auth.models import Permission from django.contrib.auth.models import User from django.contrib.contenttypes.models import ContentType from django.core.exceptions import FieldDoesNotExist from django.core.management import call_command -from django.core.management.base import BaseCommand from django.core.management.base import CommandError from django.core.serializers.base import DeserializationError from django.db import IntegrityError @@ -25,6 +23,7 @@ from django.db.models.signals import post_save from filelock import FileLock from documents.file_handling import create_source_path_directory +from documents.management.commands.base import PaperlessCommand from documents.management.commands.mixins import CryptMixin from documents.models import Correspondent from documents.models import CustomField @@ -57,21 +56,17 @@ def disable_signal(sig, receiver, sender, *, weak: bool | None = None) -> Genera sig.connect(receiver=receiver, sender=sender, **kwargs) -class Command(CryptMixin, BaseCommand): +class Command(CryptMixin, PaperlessCommand): help = ( "Using a manifest.json file, load the data from there, and import the " "documents it refers to." ) - def add_arguments(self, parser) -> None: - parser.add_argument("source") + supports_progress_bar = True - parser.add_argument( - "--no-progress-bar", - default=False, - action="store_true", - help="If set, the progress bar will not be shown", - ) + def add_arguments(self, parser) -> None: + super().add_arguments(parser) + parser.add_argument("source") parser.add_argument( "--data-only", @@ -231,7 +226,6 @@ class Command(CryptMixin, BaseCommand): self.source = Path(options["source"]).resolve() self.data_only: bool = options["data_only"] - self.no_progress_bar: bool = options["no_progress_bar"] self.passphrase: str | None = options.get("passphrase") self.version: str | None = None self.salt: str | None = None @@ -365,7 +359,7 @@ class Command(CryptMixin, BaseCommand): filter(lambda r: r["model"] == "documents.document", self.manifest), ) - for record in tqdm.tqdm(manifest_documents, disable=self.no_progress_bar): + for record in self.track(manifest_documents, description="Copying files..."): document = Document.objects.get(pk=record["pk"]) doc_file = record[EXPORTER_FILE_NAME] diff --git a/src/documents/management/commands/mixins.py b/src/documents/management/commands/mixins.py index aa5befd6d..d1e70d74e 100644 --- a/src/documents/management/commands/mixins.py +++ b/src/documents/management/commands/mixins.py @@ -1,6 +1,5 @@ import base64 import os -from argparse import ArgumentParser from typing import TypedDict from cryptography.fernet import Fernet @@ -21,25 +20,6 @@ class CryptFields(TypedDict): fields: list[str] -class ProgressBarMixin: - """ - Many commands use a progress bar, which can be disabled - via this class - """ - - def add_argument_progress_bar_mixin(self, parser: ArgumentParser) -> None: - parser.add_argument( - "--no-progress-bar", - default=False, - action="store_true", - help="If set, the progress bar will not be shown", - ) - - def handle_progress_bar_mixin(self, *args, **options) -> None: - self.no_progress_bar = options["no_progress_bar"] - self.use_progress_bar = not self.no_progress_bar - - class CryptMixin: """ Fully based on: