Test Failed
Push — master ( 66dac0...cbd2dd )
by Steffen
01:49
created

Worker.get_category()   B

Complexity

Conditions 6

Size

Total Lines 24

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 6
c 1
b 0
f 0
dl 0
loc 24
rs 7.6129
1
#!/usr/bin/python
2
# -*- coding: utf-8 -*-
3
import time
4
5
try:
6
    from titlesearch import get_similar_titles
7
except ImportError:
8
    get_similar_titles = None
9
10
from saucenao import SauceNao, FileHandler
11
12
13
class Worker(SauceNao):
14
    """
15
    Worker class for checking a list of files
16
    """
17
18
    def __init__(self, files, *args, **kwargs):
19
        """
20
        initializing function
21
22
        :type files: list|tuple|Generator
23
        :param args:
24
        :param kwargs:
25
        """
26
        super().__init__(*args, **kwargs)
27
        self.complete_file_list = files
28
29
    def run(self):
30
        """Check all files with SauceNao and execute the specified tasks
31
32
        :return:
33
        """
34
        for file_name in self.files:
35
            start_time = time.time()
36
37
            filtered_results = self.check_file(file_name)
38
39
            if not filtered_results:
40
                self.logger.info('No results found for image: {0:s}'.format(file_name))
41
                continue
42
43
            if self._move_to_categories:
44
                self.move_to_categories(file_name=file_name, results=filtered_results)
45
            else:
46
                yield {
47
                    'filename': file_name,
48
                    'results': filtered_results
49
                }
50
51
            duration = time.time() - start_time
52
            if duration < (30 / SauceNao.LIMIT_30_SECONDS):
53
                self.logger.debug("sleeping '{:.2f}' seconds".format((30 / SauceNao.LIMIT_30_SECONDS) - duration))
54
                time.sleep((30 / SauceNao.LIMIT_30_SECONDS) - duration)
55
56
    @property
57
    def excludes(self):
58
        """Property for excludes
59
60
        :return:
61
        """
62
        if self._exclude_categories:
63
            return [l.lower() for l in self._exclude_categories.split(",")]
64
        else:
65
            return []
66
67
    @property
68
    def files(self):
69
        """Property for files
70
71
        :return:
72
        """
73
        if self._start_file:
74
            # change files from generator to list
75
            files = list(self.complete_file_list)
76
            try:
77
                return files[files.index(self._start_file):]
78
            except ValueError:
79
                return self.complete_file_list
80
        return self.complete_file_list
81
82
    def get_category(self, results):
83
        """retrieve the category of the checked image based which can be either
84
        the content of the image or the author of the image
85
86
        :param results:
87
        :return: str
88
        """
89
        if self._use_author_as_category:
90
            categories = self.get_title_value(results, SauceNao.CONTENT_AUTHOR_KEY)
91
        else:
92
            categories = self.get_content_value(results, SauceNao.CONTENT_CATEGORY_KEY)
93
94
        if not categories:
95
            return ''
96
97
        self.logger.debug('categories: {0:s}'.format(', '.join(categories)))
98
99
        # since many pictures are tagged as original and with a proper category
100
        # we remove the original category if we have more than 1 category
101
        if not self._use_author_as_category and len(categories) > 1 and 'original' in categories:
102
            categories.remove('original')
103
104
        # take the first category
105
        return categories[0]
106
107
    def move_to_categories(self, file_name: str, results):
108
        """Check the file for categories and move it to the corresponding folder
109
110
        :type file_name: str
111
        :type results: list|tuple|Generator
112
        :return: bool
113
        """
114
        category = self.get_category(results)
115
        if not category:
116
            self.logger.info("no categories found for file: {0:s}".format(file_name))
117
            return False
118
119
        if not self._use_author_as_category:
120
            category = self.get_similar_title(category)
121
122
        # sub categories we don't want to move like original etc
123
        if category.lower() in self.excludes:
124
            self.logger.info("skipping excluded category: {0:s} ({1:s})".format(category, file_name))
125
            return False
126
127
        self.logger.info("moving {0:s} to category: {1:s}".format(file_name, category))
128
        FileHandler.move_to_category(file_name, category, base_directory=self._directory)
129
        return True
130
131
    def get_similar_title(self, category: str):
132
        """Check for a similar title of the category using my TitleSearch project which you can find here:
133
        https://github.com/DaRealFreak/TitleSearch
134
135
        :param category:
136
        :return:
137
        """
138
        if get_similar_titles:
139
            similar_titles = get_similar_titles(category)
140
141
            if similar_titles and similar_titles[0]['similarity'] * 100 >= self._title_minimum_similarity:
142
                self.logger.info(
143
                    "Similar title found: {0:s}, {1:s} ({2:.2f}%)".format(
144
                        category, similar_titles[0]['title'], similar_titles[0]['similarity'] * 100))
145
                return similar_titles[0]['title']
146
147
        return category
148