From 16ddf2118c00d4d1e6fa4e39f6fc5e8443afeb1c Mon Sep 17 00:00:00 2001
From: Lynne <lynne@bune.city>
Date: Mon, 26 Oct 2020 23:34:04 +1000
Subject: [PATCH] wav support!

---
 bcao.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 46 insertions(+), 10 deletions(-)

diff --git a/bcao.py b/bcao.py
index 8ef3863..8b6d818 100755
--- a/bcao.py
+++ b/bcao.py
@@ -29,7 +29,7 @@ from mutagen.oggvorbis import OggVorbis
 from mutagen.mp3 import MP3
 from mutagen.mp4 import MP4, MP4Cover
 # noinspection PyProtectedMember
-from mutagen.id3 import APIC, PictureType
+from mutagen.id3 import APIC, PictureType, Frame, TRCK, TPE1, TIT2, TALB, TPE2
 
 from PIL import Image
 
@@ -37,7 +37,7 @@ fully_supported: List[str] = ["ogg", "flac", "mp3", "m4a", "wav"]
 MutagenFile = Union[MP3, FLAC, OggVorbis, mutagen.FileType]
 MutagenTags = Union[mutagen.id3.ID3Tags, mutagen.mp4.Tags, mutagen.oggvorbis.OggVCommentDict]
 args: argparse.Namespace
-tmp_dir: tempfile.TemporaryDirectory[Any]
+tmp_dir: tempfile.TemporaryDirectory # type: ignore
 
 class SongInfo:
 	tag_lookup: Dict[str, Dict[str, str]] = {
@@ -49,15 +49,17 @@ class SongInfo:
 	}
 
 	format_lookup: Dict[str, str] = {
-		"mp3": "id3",
-		"m4a": "m4a",
-		"ogg": "vorbis",
-		"flac": "vorbis"
+		"mp3":  "id3",
+		"m4a":  "m4a",
+		"ogg":  "vorbis",
+		"flac": "vorbis",
+		"wav":  "id3"
 	}
 
 	def __init__(self, file_name: Path):
 		self.m_file: MutagenFile = mutagen.File(file_name)
 		self.m_tags: MutagenTags = self.m_file.tags
+
 		self.file_name = str(file_name.name)
 		self.format = path.splitext(file_name)[1][1:]
 		self.fallback = False
@@ -73,7 +75,7 @@ class SongInfo:
 
 		# set default values for the tags, in case the file is missing any (or all!) of them
 		self.tags: Dict[str, str] = {
-			"track": fallbacks.group("track"),
+			"track": str(int(fallbacks.group("track"))), # convert to int and str again to turn e.g. "01" into "1"
 			"artist": fallbacks.group("artist"),
 			"title": fallbacks.group("title"),
 			"album": fallbacks.group("album"),
@@ -84,8 +86,19 @@ class SongInfo:
 		self.list_tags: Dict[str, List[str]] = dict((x[0], [x[1]]) for x in self.tags.items())
 
 		if self.m_tags is None:
+			# file has no tags
+			# generate empty tags
+			self.m_file.add_tags()
+			self.m_tags = self.m_file.tags
 			self.fallback = True
 
+			# write fallback tags to file
+			for standard_name, tag_set in self.tag_lookup.items():
+				tag = tag_set[self.format_lookup[self.format]]
+				self.m_tags[tag] = self.new_id3_tag(standard_name, self.tags[standard_name])
+
+			self.m_file.save()
+
 		else:
 			for standard_name, tag_set in self.tag_lookup.items():
 				tag = tag_set[self.format_lookup[self.format]]
@@ -112,6 +125,26 @@ class SongInfo:
 				self.tags[standard_name] = value_list[0]
 				self.list_tags[standard_name] = value_list
 
+	@staticmethod
+	def new_id3_tag(tag: str, value: str) -> Frame:
+		if tag == "track":
+			return TRCK(encoding=3, text=value)
+
+		elif tag == "artist":
+			return TPE1(encoding=3, text=value)
+
+		elif tag == "title":
+			return TIT2(encoding=3, text=value)
+
+		elif tag == "album":
+			return TALB(encoding=3, text=value)
+
+		elif tag == "album_artist":
+			return TPE2(encoding=3, text=value)
+
+		else:
+			raise ValueError(f"Unknown tag type {tag}!")
+
 	def get_target_name(self, zeroes: int) -> str:
 		return f"{self.tags['track'].zfill(zeroes)} {self.tags['title']}.{self.format}"
 
@@ -122,7 +155,7 @@ class SongInfo:
 		if self.format == "flac":
 			return len(self.m_file.pictures) != 0
 
-		if self.format == "mp3":
+		if self.format == "mp3" or self.format == "wav":
 			apics: List[APIC] = self.m_tags.getall("APIC")
 			for apic in apics:
 				if apic.type == PictureType.COVER_FRONT:
@@ -141,7 +174,7 @@ class SongInfo:
 		elif self.format == "flac":
 			self.m_file.clear_pictures()
 			self.m_file.add_picture(to_embed)
-		elif self.format == "mp3":
+		elif self.format in ["mp3", "wav"]:
 			self.m_tags.add(to_embed)
 		elif self.format == "m4a":
 			self.m_tags['covr'] = [to_embed]
@@ -157,6 +190,9 @@ def log(message: str, importance: int = 0) -> None:
 
 def die(message: str, code: int = 1) -> None:
 	print(message)
+	if tmp_dir is not None:
+		tmp_dir.cleanup()
+
 	sys.exit(code)
 
 def sanitise(in_str: str) -> str:
@@ -279,7 +315,7 @@ def main() -> None:
 				embed_cover.height = image.size[1]
 				embed_cover.depth = image.bits
 
-		elif song_format == "mp3":
+		elif song_format in ["mp3", "wav"]:
 			# apparently APIC files get compressed on save if they are "large":
 			# https://mutagen.readthedocs.io/en/latest/api/id3_frames.html#mutagen.id3.APIC
 			# i don't know what that means (lossless text compression? automatic JPEG conversion?) and i don't know if or how