86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
from ipytv import playlist as tv_playlist
|
|
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
import re
|
|
|
|
def parse_arguments():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("-r", "--regex", dest="regex", help="Regex to match the categories", required=True)
|
|
parser.add_argument("-i", "--input", dest="input", help="Path to m3u+ file", required=True, type=Path)
|
|
parser.add_argument("-v", "--verbose", dest="verbose", help="Enable verbose mode", action="store_true")
|
|
|
|
return parser.parse_args()
|
|
|
|
def debug(message: str):
|
|
global verbose
|
|
|
|
if verbose:
|
|
print("[M3U-CATEGORY-MATCHER:DEBUG] " + message)
|
|
|
|
def process(arguments):
|
|
debug("Processing...")
|
|
|
|
try:
|
|
playlist = tv_playlist.loadf(str(arguments.input))
|
|
except UnicodeDecodeError:
|
|
print("Your playlist is not UTF-8 encoded, or has invalid characters. Please convert it to UTF-8.")
|
|
exit(1)
|
|
except Exception as e:
|
|
print("Something went wrong while loading playlist.")
|
|
print(e)
|
|
exit(1)
|
|
|
|
debug("Playlist loaded.")
|
|
debug("Matching categories...")
|
|
|
|
try:
|
|
regex = re.compile(arguments.regex)
|
|
except re.error:
|
|
print("Invalid regex provided.")
|
|
exit(1)
|
|
|
|
if regex.groups != 1:
|
|
print("Regex must have exactly one group to match.")
|
|
exit(1)
|
|
|
|
# Example:
|
|
# Regex: ### ([a-zA-Z0-9]+) ###
|
|
# Channel 1 name: "### Movies ###"
|
|
# Channel 2 name: "1 HD"
|
|
# Channel 3 name: "2 HD"
|
|
# Channel 4 name: "### News ###"
|
|
# Channel 5 name: "3 HD"
|
|
# Add category "Movies" to channels: 2., 3.
|
|
# Add category "News" to channel 5.
|
|
|
|
current_category = None
|
|
channels_to_remove = []
|
|
|
|
for index, channel in enumerate(playlist):
|
|
match = regex.match(channel.name)
|
|
if match:
|
|
channels_to_remove.append(index)
|
|
current_category = match.group(1)
|
|
|
|
debug(f"Category matched: {current_category}")
|
|
elif current_category:
|
|
debug(f"Adding category {current_category} to channel {channel.name}")
|
|
channel.attributes["group-title"] = current_category
|
|
|
|
for index in channels_to_remove:
|
|
playlist.remove_channel(index)
|
|
|
|
# save playlist
|
|
debug("Saving playlist...")
|
|
with open(arguments.input, "w+", encoding="utf-8") as f:
|
|
f.write(playlist.to_m3u_plus_playlist())
|
|
|
|
print("All done!")
|
|
print("Thanks for using this script!")
|
|
|
|
if __name__ == "__main__":
|
|
arguments = parse_arguments()
|
|
verbose = arguments.verbose
|
|
|
|
process(arguments)
|