postprocessor.py 6.5 KB
Newer Older
echel0n's avatar
echel0n committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# ##############################################################################
#  Author: echel0n <[email protected]>
#  URL: https://sickrage.ca/
#  Git: https://git.sickrage.ca/SiCKRAGE/sickrage.git
#  -
#  This file is part of SiCKRAGE.
#  -
#  SiCKRAGE is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#  -
#  SiCKRAGE is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  -
#  You should have received a copy of the GNU General Public License
#  along with SiCKRAGE.  If not, see <http://www.gnu.org/licenses/>.
# ##############################################################################
echel0n's avatar
echel0n committed
21
22
import os
import traceback
23
from enum import Enum
echel0n's avatar
echel0n committed
24
25

import sickrage
echel0n's avatar
echel0n committed
26
from sickrage.core.process_tv import ProcessResult
27
from sickrage.core.queues import Queue, Task, TaskPriority
echel0n's avatar
echel0n committed
28
29


30
31
32
class PostProcessorTaskActions(Enum):
    AUTO = 'Auto'
    MANUAL = 'Manual'
echel0n's avatar
echel0n committed
33
34


35
class PostProcessorQueue(Queue):
echel0n's avatar
echel0n committed
36
    def __init__(self):
37
        Queue.__init__(self, "POSTPROCESSORQUEUE")
echel0n's avatar
echel0n committed
38
39
40
41
42
43
44
45
46
        self._output = []

    @property
    def output(self):
        return '\n'.join(self._output)

    def log(self, message, level=None):
        sickrage.app.log.log(level or sickrage.app.log.INFO, message)
        self._output.append(message)
echel0n's avatar
echel0n committed
47

48
49
50
    def clear_log(self):
        self._output = []

echel0n's avatar
echel0n committed
51
52
    def find_in_queue(self, dirName, proc_type):
        """
53
54
55
        Finds any item in the queue with the given dirName and proc_type pair
        :param dirName: directory to be processed by the task
        :param proc_type: processing type, auto/manual
56
        :return: instance of PostProcessorTask or None
echel0n's avatar
echel0n committed
57
        """
58
        for task in self.tasks.copy().values():
59
            if isinstance(task, PostProcessorTask) and task.dirName == dirName and task.proc_type == proc_type:
60
                return True
61

62
63
        return False

echel0n's avatar
echel0n committed
64
65
66
67
68
69
70
71
    @property
    def queue_length(self):
        """
        Returns a dict showing how many auto and manual tasks are in the queue
        :return: dict
        """
        length = {'auto': 0, 'manual': 0}

72
        for task in self.tasks.copy().values():
73
74
            if isinstance(task, PostProcessorTask):
                if task.proc_type == 'auto':
echel0n's avatar
echel0n committed
75
76
77
78
79
80
                    length['auto'] += 1
                else:
                    length['manual'] += 1

        return length

81
    def put(self, dirName, nzbName=None, process_method=None, force=False, is_priority=None, delete_on=False,
82
            failed=False, proc_type="auto", force_next=False, **kwargs):
echel0n's avatar
echel0n committed
83
84
85
86
87
88
89
90
91
92
        """
        Adds an item to post-processing queue
        :param dirName: directory to process
        :param nzbName: release/nzb name if available
        :param process_method: processing method, copy/move/symlink/link
        :param force: force overwriting of existing files regardless of quality
        :param is_priority: whether to replace the file even if it exists at higher quality
        :param delete_on: delete files and folders after they are processed (always happens with move and auto combination)
        :param failed: mark downloads as failed if they fail to process
        :param proc_type: processing type: auto/manual
93
        :param force_next: wait until the current item in the queue is finished then process this item next
echel0n's avatar
echel0n committed
94
95
96
        :return: string indicating success or failure
        """

97
98
        self.clear_log()

echel0n's avatar
echel0n committed
99
        if not dirName:
100
            self.log("{} post-processing attempted but directory is not set".format(proc_type.title()), sickrage.app.log.WARNING)
echel0n's avatar
echel0n committed
101
            return self.output
echel0n's avatar
echel0n committed
102
103

        if not os.path.isabs(dirName):
echel0n's avatar
echel0n committed
104
105
106
107
            self.log("{} post-processing attempted but directory is relative (and probably not "
                     "what you really want to process): {}".format(proc_type.title(), dirName),
                     sickrage.app.log.WARNING)
            return self.output
echel0n's avatar
echel0n committed
108
109

        if not delete_on:
echel0n's avatar
echel0n committed
110
            delete_on = (False, (not sickrage.app.config.no_delete, True)[process_method == "move"])[proc_type == "auto"]
echel0n's avatar
echel0n committed
111

112
        if self.find_in_queue(dirName, proc_type):
echel0n's avatar
echel0n committed
113
114
            self.log("An item with directory {} is already being processed in the queue".format(dirName))
            return self.output
echel0n's avatar
echel0n committed
115
        else:
116
117
            task_id = super(PostProcessorQueue, self).put(
                PostProcessorTask(dirName, nzbName, process_method, force, is_priority, delete_on, failed, proc_type))
118

echel0n's avatar
echel0n committed
119
            if force_next:
120
                result = self.get_result(task_id) or ""
121
122
                return result

echel0n's avatar
echel0n committed
123
            self.log("{} post-processing job for {} has been added to the queue".format(proc_type.title(), dirName))
124
            return self.output + "<p><span class='hidden'>Processing succeeded</span></p>"
echel0n's avatar
echel0n committed
125
126


127
class PostProcessorTask(Task):
128
    def __init__(self, dirName, nzbName=None, process_method=None, force=False, is_priority=None, delete_on=False, failed=False, proc_type="auto"):
129
130
        action = (PostProcessorTaskActions.MANUAL, PostProcessorTaskActions.AUTO)[proc_type == "auto"]
        super(PostProcessorTask, self).__init__(action.value, action)
echel0n's avatar
echel0n committed
131
132
133
134
135
136
137
138
139
140

        self.dirName = dirName
        self.nzbName = nzbName
        self.process_method = process_method
        self.force = force
        self.is_priority = is_priority
        self.delete_on = delete_on
        self.failed = failed
        self.proc_type = proc_type

141
        self.priority = (TaskPriority.HIGH, TaskPriority.NORMAL)[proc_type == 'auto']
echel0n's avatar
echel0n committed
142

143
144
        self.auto_remove = (False, True)[proc_type == "auto"]

echel0n's avatar
echel0n committed
145
146
147
148
149
150
151
    def run(self):
        """
        Runs the task
        :return: None
        """

        try:
152
            sickrage.app.log.info("Started {} post-processing job for: {}".format(self.proc_type, self.dirName))
echel0n's avatar
echel0n committed
153

154
            result = ProcessResult(self.dirName, self.process_method, self.proc_type).process(
echel0n's avatar
echel0n committed
155
156
157
158
                nzbName=self.nzbName,
                force=self.force,
                is_priority=self.is_priority,
                delete_on=self.delete_on,
echel0n's avatar
echel0n committed
159
160
                failed=self.failed
            )
echel0n's avatar
echel0n committed
161

162
            sickrage.app.log.info("Finished {} post-processing job for: {}".format(self.proc_type, self.dirName))
echel0n's avatar
echel0n committed
163
        except Exception:
164
            sickrage.app.log.debug(traceback.format_exc())
165
166
            result = '{}'.format(traceback.format_exc())
            result += 'Processing Failed'
167

168
        return result