Coverage for toktokkie/metadata/tv/TvRenamer.py : 100%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""LICENSE
2Copyright 2015 Hermann Krumrey <hermann@krumreyh.com>
4This file is part of toktokkie.
6toktokkie is free software: you can redistribute it and/or modify
7it under the terms of the GNU General Public License as published by
8the Free Software Foundation, either version 3 of the License, or
9(at your option) any later version.
11toktokkie is distributed in the hope that it will be useful,
12but WITHOUT ANY WARRANTY; without even the implied warranty of
13MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14GNU General Public License for more details.
16You should have received a copy of the GNU General Public License
17along with toktokkie. If not, see <http://www.gnu.org/licenses/>.
18LICENSE"""
20import os
21import tvdb_api
22from abc import ABC
23from typing import List, Optional, Dict
24from puffotter.os import get_ext, listdir
25from tvdb_api import tvdb_episodenotfound, tvdb_seasonnotfound, \
26 tvdb_shownotfound
27from toktokkie.utils.ImdbCache import ImdbCache
28from toktokkie.metadata.base.Renamer import Renamer
29from toktokkie.metadata.base.components.RenameOperation import RenameOperation
30from toktokkie.enums import IdType
31from toktokkie.metadata.tv.TvExtras import TvExtras
34class TvRenamer(Renamer, TvExtras, ABC):
35 """
36 Class that handles renaming operations for tv series
37 """
39 @property
40 def selected_renaming_id_type(self) -> Optional[IdType]:
41 """
42 :return: The preferred ID type that's available for renaming
43 """
44 priority = [IdType.IMDB, IdType.TVDB]
45 for id_type in priority:
46 if len(self.ids[id_type]) > 0:
47 return id_type
48 return None
50 # noinspection PyMethodMayBeStatic
51 def create_rename_operations(self) -> List[RenameOperation]:
52 """
53 Performs rename operations on the content referenced by
54 this metadata object
55 :return: The rename operations for this metadata
56 """
57 operations = []
59 id_type = self.selected_renaming_id_type
61 if id_type is None:
62 excluded, multis, start_overrides = {}, {}, {}
63 else:
64 excluded = self.excludes.get(id_type, {})
65 multis = self.multi_episodes.get(id_type, {})
66 start_overrides = self.season_start_overrides.get(id_type, {})
68 content_info = self.get_episode_files(id_type)
70 for service_id, season_data in content_info.items():
72 if id_type is None:
73 service_ids = []
74 else:
75 service_ids = self.ids.get(id_type, [])
77 is_spinoff = id_type is not None and service_ids[0] != service_id
79 if is_spinoff:
80 sample_episode = season_data[list(season_data)[0]][0]
81 location = os.path.dirname(sample_episode)
82 series_name = os.path.basename(location)
83 else:
84 series_name = self.name
86 for _season_number, episodes in season_data.items():
87 season_number = _season_number if not is_spinoff else 1
89 season_excluded = excluded.get(season_number, [])
90 season_multis = multis.get(season_number, {})
91 episode_number = start_overrides.get(season_number, 1)
93 for episode_file in episodes:
95 while episode_number in season_excluded:
96 episode_number += 1
98 if episode_number not in season_multis:
99 end = None # type: Optional[int]
100 else:
101 end = season_multis[episode_number]
103 episode_name = self.load_episode_name(
104 service_id, id_type, season_number, episode_number, end
105 )
107 new_name = self.generate_tv_episode_filename(
108 episode_file,
109 series_name,
110 season_number,
111 episode_number,
112 episode_name,
113 end
114 )
116 if end is not None:
117 episode_number = end
119 operations.append(RenameOperation(episode_file, new_name))
120 episode_number += 1
122 return operations
124 @staticmethod
125 def generate_tv_episode_filename(
126 original_file: str,
127 series_name: str,
128 season_number: int,
129 episode_number: int,
130 episode_name: str,
131 multi_end: Optional[int] = None
132 ):
133 """
134 Generates an episode name for a given episode
135 :param original_file: The original file. Used to get the file extension
136 :param series_name: The name of the series
137 :param season_number: The season number
138 :param episode_name: The episode name
139 :param episode_number: The episode number
140 :param multi_end: Can be provided to create a multi-episode range
141 :return: The generated episode name
142 """
143 ext = get_ext(original_file)
144 if ext is not None:
145 ext = "." + ext
146 else:
147 ext = ""
149 if multi_end is None:
150 return "{} - S{}E{} - {}{}".format(
151 series_name,
152 str(season_number).zfill(2),
153 str(episode_number).zfill(2),
154 episode_name,
155 ext
156 )
157 else:
158 return "{} - S{}E{}-E{} - {}{}".format(
159 series_name,
160 str(season_number).zfill(2),
161 str(episode_number).zfill(2),
162 str(multi_end).zfill(2),
163 episode_name,
164 ext
165 )
167 @staticmethod
168 def load_episode_name(
169 service_id: str,
170 id_type: Optional[IdType],
171 season_number: int,
172 episode_number: int,
173 multi_end: Optional[int] = None
174 ) -> str:
175 """
176 Loads an episode name from external sources, if available
177 :param service_id: The external service ID for the episode's series
178 :param id_type: The external ID type
179 :param season_number: The season number
180 :param episode_number: The episode number
181 :param multi_end: If provided,
182 will generate a name for a range of episodes
183 :return: The episode name
184 """
185 default = "Episode " + str(episode_number)
187 if id_type is None or service_id == "0":
188 return default
190 if multi_end is not None:
191 episode_names = []
192 for episode in range(episode_number, multi_end + 1):
193 episode_names.append(TvRenamer.load_episode_name(
194 service_id,
195 id_type,
196 season_number,
197 episode
198 ))
199 return " | ".join(episode_names)
201 if id_type == IdType.IMDB:
202 return ImdbCache.load_episode_name(
203 service_id, season_number, episode_number
204 )
205 elif id_type == IdType.TVDB: # pragma: no cover
206 # Due to tvdb's new paid model, this will no longer be able
207 # to be tested in unit tests.
208 # IMDB will now be the de-facto default
209 try:
210 tvdb = tvdb_api.Tvdb()
211 info = tvdb[int(service_id)]
212 return info[season_number][episode_number]["episodeName"]
214 except (tvdb_episodenotfound, tvdb_seasonnotfound,
215 tvdb_shownotfound, ConnectionError, KeyError,
216 AttributeError, ValueError) as e:
217 logger = TvRenamer.logger
218 # If not found, or other error, just return generic name
219 if str(e) == "cache_location": # pragma: no cover
220 logger.warning("TheTVDB.com is down!")
221 elif str(e).startswith("apikey argument is now required"):
222 logger.warning("TheTVDB now requires an API key")
223 return default
224 else:
225 return default
227 def get_episode_files(self, id_type: Optional[IdType]) \
228 -> Dict[str, Dict[int, List[str]]]:
229 """
230 Generates a dictionary categorizing internal episode files for further
231 processing.
232 A current limitation is, that only a single service ID per season is
233 supported. It's currently not planned to lift this limitation,
234 as no valid use case for more than one ID per season has come up.
235 The episode lists are sorted by their episode name.
236 :param id_type: The ID type for which to group the episodes
237 :return: The generated dictionary. It will have the following form:
238 {service_id: {season_number: [episode_files]}}
239 """
240 content_info = {} # type: Dict[str, Dict[int, List[str]]]
242 for season_name, season_path in listdir(
243 self.directory_path, no_files=True, no_dot=True
244 ):
245 season_metadata = self.get_season(season_name)
246 if id_type is None:
247 service_id = "0"
248 else:
249 service_id = season_metadata.ids.get(id_type, ["0"])[0]
251 if service_id not in content_info:
252 content_info[service_id] = {}
254 season_number = season_metadata.season_number
255 if season_metadata.is_spinoff():
256 season_number = 1
258 if season_number not in content_info[service_id]:
259 content_info[service_id][season_number] = []
261 for episode, episode_path in listdir(
262 season_metadata.path, no_dirs=True, no_dot=True
263 ):
264 content_info[service_id][season_number].append(episode_path)
266 # Sort the episode lists
267 for service_id in content_info:
268 for season in content_info[service_id]:
269 content_info[service_id][season].sort(
270 key=lambda x: os.path.basename(x)
271 )
273 return content_info
275 def resolve_title_name(self) -> str:
276 """
277 If possible, will fetch the appropriate name for the
278 metadata based on IDs, falling back to the
279 directory name if this is not possible or supported.
280 """
281 return self.load_title_and_year([
282 IdType.ANILIST,
283 IdType.IMDB
284 ])[0]