From d97223c1aabdef9d85f6273ab6e083f58293ece3 Mon Sep 17 00:00:00 2001 From: Xie Yanbo Date: Thu, 10 Aug 2023 14:49:53 +0800 Subject: [PATCH] feat: delete task, remote and local files (#85) --- task/baidupcs.py | 3 ++ task/models.py | 12 ++++++ task/tests/test_api.py | 89 ++++++++++++++++++++++++++++----------- task/views.py | 67 ++++++++++++++++++++++++----- ui/templates/ui/base.html | 2 +- 5 files changed, 136 insertions(+), 37 deletions(-) diff --git a/task/baidupcs.py b/task/baidupcs.py index 0136685..4a2b9b7 100644 --- a/task/baidupcs.py +++ b/task/baidupcs.py @@ -132,6 +132,9 @@ def leech(self, remote_dir, local_dir, sample_size=0): self.download_dir(remote_dir, local_dir, sample_size=sample_size) + def delete(self, remote_dir): + self.api.remove(remote_dir) + def remotepath_exists(api, name: str, rd: str, _cache={}) -> bool: names = _cache.get(rd) diff --git a/task/models.py b/task/models.py index c8a4383..685d6d9 100644 --- a/task/models.py +++ b/task/models.py @@ -1,7 +1,9 @@ +import shutil from json import dumps from json import loads from os import makedirs from os import walk +from os.path import exists from os.path import getsize from os.path import join from pathlib import Path @@ -361,3 +363,13 @@ def recoverable(self): # assume task is not recoverable by default to avoid flood requests return False + + def delete_files(self): + if exists(self.sample_path): + shutil.rmtree(self.sample_path) + if exists(self.data_path): + shutil.rmtree(self.data_path) + + def erase(self): + self.delete_files() + self.delete() diff --git a/task/tests/test_api.py b/task/tests/test_api.py index 07bcfa8..d36272f 100644 --- a/task/tests/test_api.py +++ b/task/tests/test_api.py @@ -15,31 +15,30 @@ def setUp(self): self.task = Task.objects.create( shared_link="https://pan.baidu.com/s/123abc?pwd=def", ) - self.task.set_files( - [ - { - "path": "张楚", - "is_dir": True, - "is_file": False, - "size": 0, - "md5": None, - }, - { - "path": "张楚/孤独的人是可耻的.mp3", - "is_dir": False, - "is_file": True, - "size": 9518361, - "md5": "6d5bea8001e9db88f8cd8145aaf8cce4", - }, - { - "path": "张楚/蚂蚁蚂蚁.mp3", - "is_dir": False, - "is_file": True, - "size": 1234567, - "md5": "1eec826501e9db88f8cd8145aaf8cce4", - }, - ], - ) + self.remote_files = [ + { + "path": "张楚", + "is_dir": True, + "is_file": False, + "size": 0, + "md5": None, + }, + { + "path": "张楚/孤独的人是可耻的.mp3", + "is_dir": False, + "is_file": True, + "size": 9518361, + "md5": "6d5bea8001e9db88f8cd8145aaf8cce4", + }, + { + "path": "张楚/蚂蚁蚂蚁.mp3", + "is_dir": False, + "is_file": True, + "size": 1234567, + "md5": "1eec826501e9db88f8cd8145aaf8cce4", + }, + ] + self.task.set_files(self.remote_files) self.task.save() def test_create_task(self): @@ -216,7 +215,47 @@ def test_resume_not_failed_task_do_nothing(self): assert task.message == "" def test_files(self): + response = self.client.get( + reverse("task-files", args=[self.task.id]), + {}, + format="json", + ) + + assert response.json() == self.remote_files assert self.task.total_files == 2 assert self.task.total_size == 10752928 assert self.task.largest_file == "张楚/孤独的人是可耻的.mp3" assert self.task.largest_file_size == 9518361 + + def test_local_files(self): + response = self.client.get( + reverse("task-local-files", args=[self.task.id]), + {}, + format="json", + ) + + assert response.json() == [] + + @patch("task.views.get_baidupcs_client") + def test_delete_remote_files(self, mock_get_baidupcs_client): + mock_get_baidupcs_client.return_value = Mock() + + id = self.task.id + response = self.client.delete(reverse("task-files", args=[id])) + + assert response.json() == {str(id): "remote files deleted"} + + def test_delete_local_files(self): + id = self.task.id + response = self.client.delete(reverse("task-local-files", args=[id])) + + assert response.json() == {str(id): "local files deleted"} + + @patch("task.views.get_baidupcs_client") + def test_erase(self, mock_get_baidupcs_client): + mock_get_baidupcs_client.return_value = Mock() + + id = self.task.id + response = self.client.delete(reverse("task-erase", args=[id])) + + assert response.json() == {str(id): "task deleted"} diff --git a/task/views.py b/task/views.py index f81a14e..928a99a 100644 --- a/task/views.py +++ b/task/views.py @@ -1,6 +1,7 @@ import logging from io import BytesIO +from baidupcs_py.baidupcs import BaiduPCSError from django.http import HttpResponse from django_filters import rest_framework as filters from rest_framework import mixins @@ -20,6 +21,23 @@ logger = logging.getLogger(__name__) +def delete_remote_files( + task_id, + remote_path, + success_message, + catch_error=True, +): + try: + client = get_baidupcs_client() + client.delete(remote_path) + except Exception as exc: + if catch_error: + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + else: + raise exc + return Response({task_id: success_message}) + + class TaskViewSet( mixins.CreateModelMixin, mixins.RetrieveModelMixin, @@ -32,22 +50,33 @@ class TaskViewSet( filter_backends = (filters.DjangoFilterBackend,) filterset_fields = ("shared_link", "shared_id", "status", "failed") - @action(detail=True) + @action(methods=["get", "delete"], detail=True, name="Remote Files") def files(self, request, pk=None): task = self.get_object() - return Response(task.load_files()) + if request.method == "GET": + return Response(task.load_files()) + if request.method == "DELETE": + return delete_remote_files( + task.id, + task.remote_path, + "remote files deleted", + ) - @action(detail=True) + @action(methods=["get", "delete"], detail=True, name="Local Files") def local_files(self, request, pk=None): task = self.get_object() - return Response(task.list_local_files()) + if request.method == "GET": + return Response(task.list_local_files()) + if request.method == "DELETE": + task.delete_files() + return Response({task.id: "local files deleted"}) - @action(detail=True) + @action(detail=True, name="Captch Image") def captcha(self, request, pk=None): task = self.get_object() return HttpResponse(BytesIO(task.captcha), content_type="image/jpeg") - @action(methods=["post"], detail=True) + @action(methods=["post"], detail=True, name="Input Captcha Code") def captcha_code(self, request, pk=None): serializer = CaptchaCodeSerializer(data=request.data) serializer.is_valid(raise_exception=True) @@ -71,7 +100,7 @@ def captcha_code(self, request, pk=None): return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) return Response(TaskSerializer(task).data) - @action(methods=["post"], detail=True) + @action(methods=["post"], detail=True, name="Approve to download whole files") def full_download_now(self, request, pk=None): serializer = FullDownloadNowSerializer(data=request.data) serializer.is_valid(raise_exception=True) @@ -80,29 +109,45 @@ def full_download_now(self, request, pk=None): task.save() return Response(TaskSerializer(task).data) - @action(methods=["post"], detail=True) + @action(methods=["post"], detail=True, name="Restart task to downloading files") def restart_downloading(self, request, pk=None): task = self.get_object() task.restart_downloading() return Response({"status": task.status}) - @action(methods=["post"], detail=True) + @action(methods=["post"], detail=True, name="Restart task from inited status") def restart(self, request, pk=None): task = self.get_object() task.restart() return Response({"status": task.status}) - @action(methods=["post"], detail=True) + @action(methods=["post"], detail=True, name="Resume failed task") def resume(self, request, pk=None): task = self.get_object() task.schedule_resume() return Response({"status": task.status}) + @action(methods=["delete"], detail=True, name="Erase task, remote and local files") + def erase(self, request, pk=None): + task = self.get_object() + task_id = task.id + task.erase() + message = "task deleted" + try: + return delete_remote_files( + task_id, + task.remote_path, + message, + catch_error=False, + ) + except BaiduPCSError: + return Response({task_id: message}) + def get_serializer(self, *args, **kwargs): if self.action == "captcha_code": return CaptchaCodeSerializer(*args, **kwargs) if self.action == "full_download_now": return FullDownloadNowSerializer(*args, **kwargs) - if self.action in ["restart", "restart_downloading", "resume"]: + if self.action in ["restart", "restart_downloading", "resume", "erase"]: return OperationSerializer(*args, **kwargs) return super().get_serializer(*args, **kwargs) diff --git a/ui/templates/ui/base.html b/ui/templates/ui/base.html index b8733dc..6a07813 100644 --- a/ui/templates/ui/base.html +++ b/ui/templates/ui/base.html @@ -14,7 +14,7 @@ - +