Coverage for toktokkie/commands/SuperCutCommand.py : 19%
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 sys
22import argparse
23from typing import List, Dict
24from datetime import datetime
25from subprocess import Popen, check_output
26from puffotter.prompt import yn_prompt
27from toktokkie.commands.Command import Command
28from toktokkie.Directory import Directory
29from toktokkie.enums import MediaType
30from toktokkie.metadata.tv.Tv import Tv
33class SuperCutCommand(Command):
34 """
35 Class that encapsulates behaviour of the supercut command
36 """
38 @classmethod
39 def name(cls) -> str:
40 """
41 :return: The command name
42 """
43 return "supercut"
45 @classmethod
46 def help(cls) -> str:
47 """
48 :return: The help message for the command
49 """
50 return "Allows creation of supercuts of tv shows"
52 @classmethod
53 def prepare_parser(cls, parser: argparse.ArgumentParser):
54 """
55 Prepares an argumentparser for this command
56 :param parser: The parser to prepare
57 :return: None
58 """
59 parser.add_argument("directory",
60 help="The directory for which to create a "
61 "supercut")
62 parser.add_argument("--create", action="store_true",
63 help="If this flag is set, "
64 "will generate new supercut instructions")
66 def execute(self):
67 """
68 Executes the commands
69 :return: None
70 """
71 for command in ["ffmpeg", "mkvmerge", "mkvextract", "mkvpropedit"]:
72 if not self.is_installed(command):
73 self.logger.warning(
74 "{} is not installed. Can not continue.".format(command)
75 )
76 sys.exit(1)
78 directory = Directory(self.args.directory)
79 supercut_dir = os.path.join(directory.path, ".supercut")
80 supercut_file = os.path.join(supercut_dir, "supercut.txt")
82 if directory.metadata.media_type() != MediaType.TV_SERIES:
83 self.logger.warning("Only TV Series support supercut instructions")
84 elif not self.args.create and not os.path.isfile(supercut_file):
85 self.logger.warning("Couldn't find supercut instructions. "
86 "Run with --create to generate an instruction "
87 "file")
88 else:
89 if self.args.create:
91 if not os.path.isdir(supercut_dir):
92 os.makedirs(supercut_dir)
94 if os.path.isfile(supercut_file):
95 resp = yn_prompt("Instructions exist. "
96 "Do you want to overwrite the "
97 "instructions file?")
99 if resp:
100 self.generate_instructions(directory)
101 else:
102 self.generate_instructions(directory)
103 else:
104 self.create_supercut(directory)
106 @staticmethod
107 def generate_instructions(directory: Directory):
108 """
109 Generates a supercut instructions file based on the the existing
110 season metadata and episode files
111 :param directory: The directory for which to generate the instructions
112 :return: None
113 """
114 metadata = directory.metadata # type: Tv # type: ignore
115 supercut_dir = os.path.join(directory.path, ".supercut")
116 supercut_file = os.path.join(supercut_dir, "supercut.txt")
118 instructions = "# Supercut Instructions for {}".format(metadata.name)
119 instructions += "\n# Example for entry:"
120 instructions += "\n# Episode <path-to-episode>:"
121 instructions += "\n# 00:15:35-00:17:45; First appearance of char X\n"
123 seasons = list(filter(lambda x: x.season_number > 0, metadata.seasons))
124 seasons.sort(key=lambda x: x.season_number)
126 for season in seasons:
127 for episode in sorted(os.listdir(season.path)):
128 episode_path = os.path.join(season.name, episode)
129 instructions += "\nEpisode {}:".format(episode_path)
130 instructions += "\n#" + ("-" * 80)
132 with open(supercut_file, "w") as f:
133 f.write(instructions)
135 def create_supercut(self, directory: Directory):
136 """
137 Creates the supercut file
138 :param directory: The directory for which to create a supercut
139 :return: None
140 """
141 metadata = directory.metadata # type: Tv # type: ignore
143 supercut_dir = os.path.join(directory.path, ".supercut")
144 supercut_file = os.path.join(supercut_dir, "supercut.txt")
145 instructions = self.read_supercut_instructions(
146 metadata, supercut_file
147 )
149 chapter_info = []
150 for instruction in instructions:
151 chapter_info.append(self.cut_part(supercut_dir, instruction))
153 files = list(map(lambda x: x["file"], chapter_info))
154 chapters = list(map(lambda x: x["title"], chapter_info))
156 supercut_result = self.merge_parts(supercut_dir, files)
157 self.adjust_chapter_names(supercut_result, chapters)
159 def read_supercut_instructions(
160 self,
161 metadata: Tv,
162 supercut_file: str
163 ) -> List[Dict[str, str]]:
164 """
165 Reads the supercut instructions file and generates an easily
166 processable configuration list
167 :param metadata: The metadata used for additional information
168 :param supercut_file: The supercut file to parse
169 :return: A list of dicitonaries modelling the individual parts of the
170 supercut in order
171 """
172 with open(supercut_file, "r") as f:
173 instructions = f.read()
175 config = []
176 episode_file = None
177 part_count = 0
179 for line in instructions.split("\n"):
180 line = line.strip()
182 if line.startswith("#") or line == "":
183 pass
184 elif line.lower().startswith("episode "):
185 episode_file = line.split(" ", 1)[1].rsplit(":", 1)[0]
186 episode_file = os.path.join(
187 metadata.directory_path, episode_file
188 )
189 elif episode_file is None:
190 raise ValueError("Instructions Syntax Error")
191 elif not episode_file.endswith(".mkv"):
192 self.logger.warning("ONLY MKV FILES ARE SUPPORTED")
193 else:
194 part_count += 1
195 config.append({
196 "start": line.split("-")[0].strip(),
197 "end": line.split("-")[1].split(";")[0].strip(),
198 "title": line.split(";")[1].strip(),
199 "file": episode_file,
200 "part": str(part_count).zfill(4)
201 })
203 return config
205 def cut_part(self, supercut_dir: str, instruction: Dict[str, str]) \
206 -> Dict[str, str]:
207 """
208 Cuts a part of a video file based on a supercut instruction using
209 ffmpeg
210 :param supercut_dir: The supercut directory in which to
211 store the output
212 :param instruction: The instruction to use
213 :return: A dictionary detailing the path to the cut clip and its title
214 """
216 input_file = instruction["file"]
217 output_file = os.path.join(
218 supercut_dir,
219 "{} - {}.mkv".format(instruction["part"], instruction["title"])
220 )
222 start_time = instruction["start"]
223 end_time = instruction["end"]
224 delta = self.calculate_clip_duration(start_time, end_time)
226 if not os.path.isfile(output_file):
227 Popen([
228 "ffmpeg",
229 "-ss", start_time,
230 "-i", input_file,
231 "-t", str(delta),
232 output_file
233 ]).wait()
235 return {"file": output_file, "title": instruction["title"]}
237 @staticmethod
238 def merge_parts(supercut_dir: str, files: List[str]) -> str:
239 """
240 Merges a list of MKV files and stores them in supercut.mkv
241 :param supercut_dir: The supercut directory
242 :param files: The files to combine
243 :return: The path to the combined file
244 """
245 dest = os.path.join(supercut_dir, "supercut.mkv")
246 mkvmerge_com = [
247 "mkvmerge",
248 "-o", dest,
249 "--generate-chapters", "when-appending",
250 files.pop(0)
251 ]
252 mkvmerge_com += list(map(lambda x: "+" + x, files))
253 Popen(mkvmerge_com).wait()
254 return dest
256 @staticmethod
257 def adjust_chapter_names(supercut_result: str, chapters: List[str]):
258 """
259 Changes the chapter names of a file to a given list of chapter names
260 :param supercut_result: The file of which the chapters should be
261 renamed
262 :param chapters: The chapter names
263 :return: None
264 """
265 chapters_info = check_output([
266 "mkvextract", "chapters", supercut_result
267 ]).decode("utf-8")
269 adjusted = []
271 for line in chapters_info.split("\n"):
272 if line.strip().startswith("<ChapterString>"):
273 adjusted.append("<ChapterString>{}</ChapterString>".format(
274 chapters.pop(0)
275 ))
276 else:
277 adjusted.append(line)
279 with open("chapters.xml", "w") as f:
280 f.write("\n".join(adjusted))
282 Popen([
283 "mkvpropedit", supercut_result, "--chapters", "chapters.xml"
284 ]).wait()
286 os.remove("chapters.xml")
288 @staticmethod
289 def calculate_clip_duration(start: str, end: str) -> int:
290 """
291 Calculates the tme delta in seconds between two timestamps
292 :param start: The starting timestamp
293 :param end: The ending timestamp
294 :return: The time delta in seconds
295 """
296 start_hour, start_minute, start_second = \
297 list(map(lambda x: int(x), start.split(":")))
298 end_hour, end_minute, end_second = \
299 list(map(lambda x: int(x), end.split(":")))
301 _start = datetime(1, 1, 1, start_hour, start_minute, start_second)
302 _end = datetime(1, 1, 1, end_hour, end_minute, end_second)
303 return (_end - _start).seconds
305 @staticmethod
306 def is_installed(command: str) -> bool:
307 """
308 Checks whether or not a command is installed/in the path
309 :param command: THe command to check
310 :return: True if installed, else False
311 """
312 for path in os.environ["PATH"].split(":"):
313 if os.path.isfile(os.path.join(path, command)):
314 return True
315 return False