1# Copyright 2016 Christoph Reiter
2#
3# This program is free software; you can redistribute it and/or modify
4# it under the terms of the GNU General Public License as published by
5# the Free Software Foundation; either version 2 of the License, or
6# (at your option) any later version.
7
8"""Code for serializing AudioFile instances"""
9
10import pickle
11from senf import bytes2fsn, fsn2bytes
12
13from quodlibet.util.picklehelper import pickle_loads, pickle_dumps
14from quodlibet.util import is_windows
15from ._audio import AudioFile
16
17
18class SerializationError(Exception):
19    pass
20
21
22def _py2_to_py3(items):
23    for i in items:
24        try:
25            l = list(i.items())
26        except AttributeError:
27            raise SerializationError
28        i.clear()
29        for k, v in l:
30            if isinstance(k, bytes):
31                k = k.decode("utf-8", "replace")
32            else:
33                # strip surrogates
34                try:
35                    k.encode("utf-8")
36                except UnicodeEncodeError:
37                    k = k.encode("utf-8", "replace").decode("utf-8")
38
39            if k == "~filename" or k == "~mountpoint":
40                if isinstance(v, bytes):
41                    try:
42                        v = bytes2fsn(v, "utf-8")
43                    except ValueError:
44                        # just in case, only on Windows
45                        assert is_windows()
46                        v = v.decode("utf-8", "replace")
47            elif isinstance(v, bytes):
48                v = v.decode("utf-8", "replace")
49            elif isinstance(v, str):
50                # strip surrogates
51                try:
52                    v.encode("utf-8")
53                except UnicodeEncodeError:
54                    v = v.encode("utf-8", "replace").decode("utf-8")
55
56            i[k] = v
57
58    return items
59
60
61def _py3_to_py2(items):
62    is_win = is_windows()
63
64    new_list = []
65    for i in items:
66        inst = dict.__new__(i.__class__)
67        for key, value in i.items():
68            if key in ("~filename", "~mountpoint") and not is_win:
69                value = fsn2bytes(value, None)
70            try:
71                key = key.encode("ascii")
72            except UnicodeEncodeError:
73                pass
74            dict.__setitem__(inst, key, value)
75        new_list.append(inst)
76    return new_list
77
78
79def load_audio_files(data, process=True):
80    """unpickles the item list and if some class isn't found unpickle
81    as a dict and filter them out afterwards.
82
83    In case everything gets filtered out will raise SerializationError
84    (because then likely something larger went wrong)
85
86    Args:
87        data (bytes)
88        process (bool): if the dict key/value types should be converted,
89            either to be usable from py3 or to convert to newer types
90    Returns:
91        List[AudioFile]
92    Raises:
93        SerializationError
94    """
95
96    dummy = type("dummy", (dict,), {})
97    error_occured = []
98    temp_type_cache = {}
99
100    def lookup_func(base, module, name):
101        try:
102            real_type = base(module, name)
103        except (ImportError, AttributeError):
104            error_occured.append(True)
105            return dummy
106
107        if module.split(".")[0] not in ("quodlibet", "tests"):
108            return real_type
109
110        # return a straight dict subclass so that unpickle doesn't call
111        # our __setitem__. Further down we simply change the __class__
112        # to our real type.
113        if not real_type in temp_type_cache:
114            new_type = type(name, (dict,), {"real_type": real_type})
115            temp_type_cache[real_type] = new_type
116
117        return temp_type_cache[real_type]
118
119    try:
120        items = pickle_loads(data, lookup_func)
121    except pickle.UnpicklingError as e:
122        raise SerializationError(e)
123
124    if error_occured:
125        items = [i for i in items if not isinstance(i, dummy)]
126
127        if not items:
128            raise SerializationError(
129                "all class lookups failed. something is wrong")
130
131    if process:
132        items = _py2_to_py3(items)
133
134    try:
135        for i in items:
136            i.__class__ = i.real_type
137    except AttributeError as e:
138        raise SerializationError(e)
139
140    return items
141
142
143def dump_audio_files(item_list, process=True):
144    """Pickles a list of AudioFiles
145
146    Returns:
147        bytes
148    Raises:
149        SerializationError
150    """
151
152    assert isinstance(item_list, list)
153    assert not item_list or isinstance(item_list[0], AudioFile)
154
155    if process:
156        item_list = _py3_to_py2(item_list)
157
158    try:
159        return pickle_dumps(item_list, 2)
160    except pickle.PicklingError as e:
161        raise SerializationError(e)
162