FilesExpand file tree

master

/

export.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

265 lines (223 loc) · 8.13 KB

master

/

export.py

Copy path

Top

File metadata and controls

• Code
• Blame

265 lines (223 loc) · 8.13 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

"""Exports data from beets"""

from __future__ import annotations

import csv

import json

import os

import sys

from datetime import date, datetime

from typing import TYPE_CHECKING, Any, Literal, Protocol, get_args

from xml.etree import ElementTree

import mediafile

from beets import ui

from beets.dbcore.types import BasePathType

from beets.library.fields import TYPE_BY_FIELD

from beets.plugins import BeetsPlugin

from beetsplug.info import library_data, tag_data

if TYPE_CHECKING:

from collections.abc import Iterable, Iterator, Sequence

from beets.library import Library

from ._typing import JSONDict

Format = Literal["json", "jsonlines", "csv", "xml"]

VALID_FORMATS = get_args(Format)

class ExportCLIOpts(Protocol):

library: bool | None

album: bool

append: bool

included_keys: Sequence[str]

limit: int | None

output: str | None

format: Format | None

class ExportEncoder(json.JSONEncoder):

"""Deals with dates because JSON doesn't have a standard"""

def default(self, o: object) -> Any:

if isinstance(o, (datetime, date)):

return o.isoformat()

return json.JSONEncoder.default(self, o)

class ExportPlugin(BeetsPlugin):

default_format: Format

def __init__(self) -> None:

super().__init__()

self.config.add(

{

"default_format": "json",

"json": {

# JSON module formatting options.

"formatting": {

"ensure_ascii": False,

"indent": 4,

"separators": (",", ": "),

"sort_keys": True,

}

},

"jsonlines": {

# JSON Lines formatting options.

"formatting": {

"ensure_ascii": False,

"separators": (",", ": "),

"sort_keys": True,

}

},

"csv": {

# CSV module formatting options.

"formatting": {

# The delimiter used to separate columns.

"delimiter": ",",

# The dialect to use when formatting the file output.

"dialect": "excel",

}

},

"xml": {

# XML module formatting options.

"formatting": {}

},

# TODO: Use something like the edit plugin

# 'item_fields': []

}

)

self.default_format = self.config["default_format"].as_choice(

VALID_FORMATS

)

def commands(self) -> list[ui.Subcommand]:

cmd = ui.Subcommand("export", help="export data from beets")

cmd.func = self.run

cmd.parser.add_album_option()

cmd.parser.add_option(

"-l",

"--library",

action="store_true",

help="show library fields instead of tags",

)

cmd.parser.add_option(

"--append",

action="store_true",

default=False,

help="if should append data to the file",

)

cmd.parser.add_option(

"-i",

"--include-keys",

default=[],

action="append",

dest="included_keys",

help="comma separated list of keys to show",

)

cmd.parser.add_option(

"-o",

"--output",

help="path for the output file. If not given, will print the data",

)

cmd.parser.add_option(

"-f",

"--format",

type="choice",

choices=VALID_FORMATS,

default=self.config["default_format"].get(),

help="the output format: json|jsonlines|csv|xml",

)

cmd.parser.add_limit_option(flags=("--limit",))

return [cmd]

def run(

self, lib: Library, opts: ExportCLIOpts, args: Sequence[str]

) -> None:

file_path = opts.output

file_mode = "a" if opts.append else "w"

file_format = opts.format or self.default_format

format_options = self.config[file_format]["formatting"].get(dict)

_format = (

CSVFormat

if file_format == "csv"

else XMLFormat

if file_format == "xml"

else JsonFormat

)

export_format = _format(

file_format=file_format, file_path=file_path, file_mode=file_mode

)

if opts.library or opts.album:

data_collector = library_data

else:

data_collector = tag_data

included_keys = []

for keys in opts.included_keys:

included_keys.extend(keys.split(","))

byte_fields = [

k for k, v in TYPE_BY_FIELD.items() if isinstance(v, BasePathType)

]

def collect_data() -> Iterator[JSONDict]:

for data_emitter in data_collector(

lib, args, album=opts.album, limit=opts.limit

):

try:

data, _ = data_emitter(included_keys or "*")

except (mediafile.UnreadableFileError, OSError) as ex:

self._log.error("cannot read file: {}", ex)

continue

else:

yield data

def stringify_bytes(data: JSONDict) -> JSONDict:

for field in byte_fields:

if (value := data.get(field)) is not None:

data[field] = os.fsdecode(value)

return data

export_format.export(

map(stringify_bytes, collect_data()), **format_options

)

class ExportFormat:

"""The output format type"""

def __init__(

self,

file_format: str,

file_path: str | None,

file_mode: str = "w",

encoding: str = "utf-8",

) -> None:

self.file_format = file_format

self.path = file_path

self.mode = file_mode

self.encoding = encoding

# creates a file object to write/append or sets to stdout

self.out_stream = (

open(self.path, self.mode, encoding=self.encoding, newline="")

if self.path

else sys.stdout

)

def export(self, data_iter: Iterable[JSONDict], **kwargs) -> None:

raise NotImplementedError()

class JsonFormat(ExportFormat):

"""Saves in a json file"""

def _print_json(self, data: Any, **kwargs) -> None:

json.dump(data, self.out_stream, cls=ExportEncoder, **kwargs)

self.out_stream.write("\n")

def export(self, data_iter: Iterable[JSONDict], **kwargs) -> None:

if self.file_format == "json":

self._print_json(list(data_iter), **kwargs)

else:

for item in data_iter:

self._print_json(item, **kwargs)

class CSVFormat(ExportFormat):

"""Saves in a csv file"""

def export(self, data_iter: Iterable[JSONDict], **kwargs) -> None:

data = list(data_iter)

header = list(data[0].keys()) if data else []

writer = csv.DictWriter(self.out_stream, fieldnames=header, **kwargs)

writer.writeheader()

writer.writerows(data)

class XMLFormat(ExportFormat):

"""Saves in a xml file"""

def export(self, data_iter: Iterable[JSONDict], **kwargs) -> None:

# Creates the XML file structure.

data = list(data_iter)

library = ElementTree.Element("library")

tracks = ElementTree.SubElement(library, "tracks")

if data and isinstance(data[0], dict):

for index, item in enumerate(data):

track = ElementTree.SubElement(tracks, "track")

for key, value in item.items():

track_details = ElementTree.SubElement(track, key)

track_details.text = value

# Depending on the version of python the encoding needs to change

try:

string = ElementTree.tostring(library, encoding="unicode", **kwargs)

except LookupError:

string = ElementTree.tostring(library, encoding="utf-8", **kwargs)

self.out_stream.write(string)