AudioFile.__setitem__: do key type checking; convert keys to ascii st…

…r on py2 if possible

quodlibet/quodlibet/formats/_audio.py

@@ -170,6 +170,27 @@ def __setstate__(self, state):
         pass
 
     def __setitem__(self, key, value):
+        # validate key
+        if PY3:
+            if not isinstance(key, text_type):
+                raise TypeError("key has to be str")
+        else:
+            if isinstance(key, text_type):
+                # we try to save keys as encoded ASCII to save memory
+                # under PY2. Everything else besides ASCII combined with
+                # unicode breaks hashing even if the default encoding
+                # it utf-8.
+                try:
+                    key = key.encode("ascii")
+                except UnicodeEncodeError:
+                    pass
+            elif isinstance(key, bytes):
+                # make sure we set ascii keys only
+                key.decode("ascii")
+            else:
+                raise TypeError("key needs to be unicode or ASCII str")
+
+        # validate value
         if key.startswith("~#"):
             if not isinstance(value, number_types):
                 raise TypeError

quodlibet/tests/test_formats__audio.py

@@ -10,7 +10,7 @@
 from senf import fsnative, fsn2text, bytes2fsn
 
 from quodlibet import config
-from quodlibet.compat import PY2, text_type, long
+from quodlibet.compat import PY2, text_type, long, listkeys, PY3
 from quodlibet.formats import AudioFile, types as format_types, AudioFileError
 from quodlibet.formats._audio import NUMERIC_ZERO_DEFAULT
 from quodlibet.formats import decode_value, MusicFile
@@ -93,6 +93,26 @@ def test_trackdisc(self):
         self.failIf(bar_1_2("~#discs"))
         self.failIf(bar_2_1("~#tracks"))
 
+    def test_setitem_keys(self):
+        af = AudioFile()
+        af[u"foo"] = u"bar"
+        assert "foo" in af
+        assert isinstance(listkeys(af)[0], str)
+        af.clear()
+        af[u"öäü"] = u"bar"
+        assert u"öäü" in af
+        assert isinstance(listkeys(af)[0], text_type)
+
+        with self.assertRaises(TypeError):
+            af[42] = u"foo"
+
+        if PY3:
+            with self.assertRaises(TypeError):
+                af[b"foo"] = u"bar"
+        else:
+            with self.assertRaises(ValueError):
+                af[b"\xff"] = u"bar"
+
     def test_call(self):
         # real keys should lookup the same
         for key in bar_1_1.realkeys():