Remove archive on import during move mode (#6551)

Fixes #1663.

This modifies the `cleanup` method on the `ArchiveImportTask` to remove
the original archive file on successful import when `move: yes`.

I decided to do this such that partial imports do not destructively
remove the archive. Perhaps this would be better as a configuration
option, but I think that may go beyond the scope for this initial
feature.

I am unsure whether my method for "complete import" is strictly correct,
I reasoned that we would check that all files had been moved out, but
surely this data exists more concretely somewhere. Let me know if this
should be polished, or there's more nuance I'm not seeing!

Additionally, I added unit tests for this behavior under the *old* test
framework since existing helpers are present for the archive backend
setup. This way we minimize the diff and maximize probability that my
test is behaving correctly. I will gladly port this entire archive test
section to pytest in a future PR if this is a dealbreaker.

Author: snejus

beets/importer/tasks.py

@@ -803,19 +803,27 @@ def _emit_imported(self, lib):
 class ArchiveImportTask(SentinelImportTask):
     """An import task that represents the processing of an archive.
 
-    `toppath` must be a `zip`, `tar`, or `rar` archive. Archive tasks
-    serve two purposes:
+    `toppath` must be a `zip`, `tar`, `rar`, or `7z` archive. Archive tasks
+    serve three purposes:
     - First, it will unarchive the files to a temporary directory and
       return it. The client should read tasks from the resulting
       directory and send them through the pipeline.
     - Second, it will clean up the temporary directory when it proceeds
       through the pipeline. The client should send the archive task
       after sending the rest of the music tasks to make this work.
+    - Third, when the import mode is ``move`` and every file in the
+      archive was successfully imported, it will remove the source
+      archive itself. Archives are preserved on partial imports and in
+      non-move modes.
     """
 
     def __init__(self, toppath):
         super().__init__(toppath, ())
         self.extracted = False
+        # ``extract()`` reassigns ``self.toppath`` to the temp extraction
+        # directory; here we track the original archive location so
+        # ``cleanup()`` can remove it when the import mode demands.
+        self.archive_path = toppath
 
     @classmethod
     def is_archive(cls, path):
@@ -862,13 +870,37 @@ def handlers(cls) -> list[ArchiveHandler]:
         return _handlers
 
     def cleanup(self, copy=False, delete=False, move=False):
-        """Removes the temporary directory the archive was extracted to."""
-        if self.extracted and self.toppath:
+        """Remove the temporary extraction directory and optionally the archive.
+
+        In ``move`` mode, if the extraction directory is empty after the
+        pipeline has run (i.e. every file in the archive was successfully
+        imported) also remove the source archive. Archives are preserved on
+        partial imports and in non-move modes.
+        """
+        if not self.extracted:
+            return
+
+        all_files_imported = move and not any(
+            files for _, _, files in os.walk(util.syspath(self.toppath))
+        )
+
+        log.debug(
+            "Removing extracted directory: {}",
+            util.displayable_path(self.toppath),
+        )
+        shutil.rmtree(util.syspath(self.toppath))
+
+        if all_files_imported:
+            log.debug(
+                "Removing imported archive: {}",
+                util.displayable_path(self.archive_path),
+            )
+            util.remove(self.archive_path)
+        elif move:
             log.debug(
-                "Removing extracted directory: {}",
-                util.displayable_path(self.toppath),
+                "Not removing partially imported archive: {}",
+                util.displayable_path(self.archive_path),
             )
-            shutil.rmtree(util.syspath(self.toppath))
 
     def extract(self):
         """Extracts the archive to a temporary directory and sets

docs/changelog.rst

@@ -19,6 +19,11 @@ New features
   would be updated"* instead of *"N playlists updated"*. The ``--format`` option
   allows customizing the track line format. The ``--pretend-paths`` option was
   removed (use ``--format='$path'`` instead). :bug:`6183`
+- :ref:`import-cmd`: When importing an archive (zip, tar, rar, or 7z) with
+  ``move: yes``, the source archive is now removed after a successful import.
+  Archives are preserved if any file in the archive was not imported (e.g.
+  skipped as a duplicate, or the import was aborted), and in non-move import
+  modes.
 
 ..
     Bug fixes

docs/reference/cli.rst

@@ -71,9 +71,11 @@ Optional command flags:
 
 - By default, the command copies files to your library directory and updates the
   ID3 tags on your music. In order to move the files, instead of copying, use
-  the ``-m`` (move) option. If you'd like to leave your music files untouched,
-  try the ``-C`` (don't copy) and ``-W`` (don't write tags) options. You can
-  also disable this behavior by default in the configuration file (below).
+  the ``-m`` (move) option. When importing an archive with ``-m``, if all files
+  are imported, the archive is removed from disk. If you'd like to leave your
+  music files untouched, try the ``-C`` (don't copy) and ``-W`` (don't write
+  tags) options. You can also disable this behavior by default in the
+  configuration file (below).
 - Also, you can disable the autotagging behavior entirely using ``-A`` (don't
   autotag)---then your music will be imported with its existing metadata.
 - During a long tagging import, it can be useful to keep track of albums that

docs/reference/config.rst

@@ -581,6 +581,9 @@ beets will copy and then delete when a simple rename is impossible.) Moving
 files can be risky—it's a good idea to keep a backup in case beets doesn't do
 what you expect with your files.
 
+In the case of a ``move`` when importing an archive, the archive will be removed
+if all contents were successfully imported.
+
 This option *overrides* ``copy``, so enabling it will always move (and not copy)
 files. The ``-c`` switch to the ``beet import`` command, however, still takes
 precedence.

test/test_importer.py

@@ -180,6 +180,9 @@ def setUp(self):
         super().setUp()
         self.want_resume = False
         self.config["incremental"] = False
+        self.config["copy"] = False
+        self.config["move"] = False
+        self.config["delete"] = False
         self._old_home = None
 
     def test_rm(self):
@@ -191,6 +194,48 @@ def test_rm(self):
         archive_task.finalize(self)
         assert not tmp_path.exists()
 
+    def test_archive_removed_on_move_complete(self):
+        zip_path = create_archive(self)
+        archive_task = importer.ArchiveImportTask(zip_path)
+        archive_task.extract()
+        for root, _, files in os.walk(syspath(archive_task.toppath)):
+            for f in files:
+                os.remove(os.path.join(root, f))
+        assert Path(os.fsdecode(zip_path)).exists()
+        archive_task.cleanup(move=True)
+        assert not Path(os.fsdecode(zip_path)).exists()
+
+    def test_archive_preserved_on_move_partial(self):
+        zip_path = create_archive(self)
+        archive_task = importer.ArchiveImportTask(zip_path)
+        archive_task.extract()
+        archive_task.cleanup(move=True)
+        assert Path(os.fsdecode(zip_path)).exists()
+
+    def test_archive_preserved_on_copy(self):
+        zip_path = create_archive(self)
+        archive_task = importer.ArchiveImportTask(zip_path)
+        archive_task.extract()
+        archive_task.cleanup(copy=True)
+        assert Path(os.fsdecode(zip_path)).exists()
+
+    def test_tempdir_removed_in_all_modes(self):
+        for cleanup_kwargs in (
+            {},
+            {"move": True},
+            {"copy": True},
+            {"copy": True, "delete": True},
+        ):
+            zip_path = create_archive(self)
+            archive_task = importer.ArchiveImportTask(zip_path)
+            archive_task.extract()
+            tmp_path = Path(os.fsdecode(archive_task.toppath))
+            assert tmp_path.exists(), f"extract failed for {cleanup_kwargs}"
+            archive_task.cleanup(**cleanup_kwargs)
+            assert not tmp_path.exists(), (
+                f"tempdir {tmp_path} not removed for {cleanup_kwargs}"
+            )
+
 
 class ImportZipTest(AsIsImporterMixin, ImportTestCase):
     def test_import_zip(self):