improve mp3 support (id3v1 and id3v2)

1 files changed, 96 insertions, 20 deletions

diff --git a/formats/mp3.py b/formats/mp3.py
index 0c91911..c668b47 100644
--- a/formats/mp3.py
+++ b/formats/mp3.py
@@ -2,18 +2,92 @@ from formats.__base__ import *
 from utils import to_bytes
 
 import subprocess
+import struct
+import array
+
+def synchsafe(num):
+	if num <= 0x7f:
+		return num
+
+	return synchsafe(num >> 7) << 8 | num & 0x7f
+
+class ID3Tagger:
+	# id3v2 frame mapping
+	__mapping = {
+		"album":	"TALB",
+		"artist":	"TPE1",
+		"composer":	"TCOM",
+		"date":		"TDRC",
+		"title":	"TIT2",
+		"tracknumber":	"TRCK",
+	}
+
+	# id3v1 offsets
+	__offset = {
+		"title": 3,
+		"artist": 33,
+		"album": 63,
+	}
+
+	@staticmethod
+	def header(size):
+		return struct.pack(">3s3BI", b"ID3", 4, 0, 0, synchsafe(size))
+
+	@staticmethod
+	def frame(name, data):
+		size = len(data) + 1
+		hdr = struct.pack(">4sIHB", name, size, 0, 3)
+		return hdr + data
+
+	def __init__(self):
+		self.frames = []
+
+		self.v1 = array.array("B", b"\x00" * 128)
+		struct.pack_into("3s", self.v1, 0, b"TAG")
+		struct.pack_into("B", self.v1, 127, 0xff)
+
+	def frame_size(self):
+		return sum(map(len, self.frames))
+
+	def add(self, tag, value):
+		value = to_bytes(value)
+
+		if tag in self.__mapping:
+			key = to_bytes(self.__mapping[tag])
+			self.frames.append(self.frame(key, value))
+
+		off = self.__offset.get(tag)
+		if off:
+			struct.pack_into("30s", self.v1, off, value)
+		elif tag is "date":
+			struct.pack_into("4s", self.v1, 93, value)
+		elif tag is "tracknumber":
+			number = int(value.partition(b"/")[0])
+			struct.pack_into("B", self.v1, 126, number)
+
+	def write(self, path):
+		fp = open(path, "r+b")
+		data = fp.read()
+
+		fp.seek(0)
+		fp.truncate(0)
+
+		# save id3v2
+		fp.write(self.header(self.frame_size()))
+		for frame in self.frames:
+			fp.write(frame)
+
+		fp.write(data)
+
+		# save id3v1
+		self.v1.tofile(fp)
+
+		fp.close()
 
 class Mp3Handler(BaseHandler):
 	name = "mp3"
 	ext = "mp3"
 
-	__tag_opts = {
-		"album": "-a",
-		"artist": "-A",
-		"date": "-y",
-		"title": "-t"
-	}
-
 	def encode(self, opt, info):
 		self.add("cust ext=%s sox -" % self.ext)
 
@@ -25,25 +99,27 @@ class Mp3Handler(BaseHandler):
 
 		return self.build()
 
-	def tag(self, path, tags):
-		self.add("id3v2", "--id3v1-only")
+	def do_tag(self, path, tags):
+		tagger = ID3Tagger()
 
 		for k, v in tags.items():
-			if k in self.__tag_opts and v:
-				self.add(self.__tag_opts[k])
-				self.add(v)
+			if v and k not in ("tracknumber", "tracktotal"):
+				tagger.add(k, v)
 
-		self.add("-T", "%d/%d" % (tags["tracknumber"], tags["tracktotal"]))
-		self.add(path)
+		number = "%d/%d" % (tags["tracknumber"], tags["tracktotal"])
+		tagger.add("tracknumber", number)
 
-		self.log("Tag [%s] : ", path)
+		tagger.write(path)
+		return True
 
-		if subprocess.call(self.build(False)):
-			self.log("FAILED\n")
-			return False
+	def tag(self, path, tags):
+		self.log("Tag [%s] : ", path)
+		if self.do_tag(path, tags):
+			self.log("OK\n")
+			return True
 
-		self.log("OK\n")
-		return True
+		self.log("FAILED\n")
+		return False
 
 def init():
 	return Mp3Handler