diff --git a/beetsplug/export.py b/beetsplug/export.py index d783f5b93..cd91fc393 100644 --- a/beetsplug/export.py +++ b/beetsplug/export.py @@ -18,8 +18,10 @@ from __future__ import division, absolute_import, print_function import sys -import json import codecs +import json +import csv +import xml.etree.ElementTree as ET from datetime import datetime, date from beets.plugins import BeetsPlugin @@ -52,6 +54,24 @@ class ExportPlugin(BeetsPlugin): 'sort_keys': True } }, + 'csv': { + # csv module formatting options + 'formatting': { + 'ensure_ascii': False, + 'indent': 0, + 'separators': (','), + 'sort_keys': True + } + }, + 'xml': { + # xml module formatting options + 'formatting': { + 'ensure_ascii': False, + 'indent': 4, + 'separators': (','), + 'sort_keys': True + } + } # TODO: Use something like the edit plugin # 'item_fields': [] }) @@ -78,17 +98,22 @@ class ExportPlugin(BeetsPlugin): u'-o', u'--output', help=u'path for the output file. If not given, will print the data' ) + cmd.parser.add_option( + u'-f', u'--format', default='json', + help=u'specify the format of the exported data. Your options are json (deafult), csv, and xml' + ) return [cmd] def run(self, lib, opts, args): file_path = opts.output - file_format = self.config['default_format'].get(str) file_mode = 'a' if opts.append else 'w' + file_format = opts.format format_options = self.config[file_format]['formatting'].get(dict) export_format = ExportFormat.factory( - file_format, **{ + file_type=file_format, + **{ 'file_path': file_path, 'file_mode': file_mode } @@ -108,44 +133,107 @@ class ExportPlugin(BeetsPlugin): except (mediafile.UnreadableFileError, IOError) as ex: self._log.error(u'cannot read file: {0}', ex) continue - data = key_filter(data) items += [data] - export_format.export(items, **format_options) class ExportFormat(object): """The output format type""" - - @classmethod - def factory(cls, type, **kwargs): - if type == "json": - if kwargs['file_path']: - return JsonFileFormat(**kwargs) - else: - return JsonPrintFormat() - raise NotImplementedError() - - def export(self, data, **kwargs): - raise NotImplementedError() - - -class JsonPrintFormat(ExportFormat): - """Outputs to the console""" - - def export(self, data, **kwargs): - json.dump(data, sys.stdout, cls=ExportEncoder, **kwargs) - - -class JsonFileFormat(ExportFormat): - """Saves in a json file""" - def __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'): self.path = file_path self.mode = file_mode self.encoding = encoding + @classmethod + def factory(cls, file_type, **kwargs): + if file_type == "json": + return JsonFormat(**kwargs) + elif file_type == "csv": + return CSVFormat(**kwargs) + elif file_type == "xml": + return XMLFormat(**kwargs) + else: + raise NotImplementedError() + def export(self, data, **kwargs): + raise NotImplementedError() + + +class JsonFormat(ExportFormat): + """Saves in a json file""" + def __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'): + super(JsonFormat, self).__init__(file_path, file_mode, encoding) + self.export = self.export_to_file if self.path else self.export_to_terminal + + def export_to_terminal(self, data, **kwargs): + json.dump(data, sys.stdout, cls=ExportEncoder, **kwargs) + + def export_to_file(self, data, **kwargs): with codecs.open(self.path, self.mode, self.encoding) as f: json.dump(data, f, cls=ExportEncoder, **kwargs) + + +class CSVFormat(ExportFormat): + """Saves in a csv file""" + def __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'): + super(CSVFormat, self).__init__(file_path, file_mode, encoding) + self.header = [] + + def export(self, data, **kwargs): + if data and type(data) is list and len(data) > 0: + self.header = list(data[0].keys()) + if self.path: + self.export_to_file(data, **kwargs) + else: + self.export_to_terminal(data, **kwargs) + + def export_to_terminal(self, data, **kwargs): + writer = csv.DictWriter(sys.stdout, fieldnames=self.header) + writer.writeheader() + writer.writerows(data) + + def export_to_file(self, data, **kwargs): + with codecs.open(self.path, self.mode, self.encoding) as f: + writer = csv.DictWriter(f, fieldnames=self.header) + writer.writeheader() + writer.writerows(data) + + +class XMLFormat(ExportFormat): + """Saves in a xml file""" + def __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'): + super(XMLFormat, self).__init__(file_path, file_mode, encoding) + + def export(self, data, **kwargs): + # create the file structure + library = ET.Element('library') + tracks_key = ET.SubElement(library, 'key') + tracks_key.text = "Tracks" + tracks_dict = ET.SubElement(library, 'dict') + if data and type(data) is list \ + and len(data) > 0 and type(data[0]) is dict: + index = 1 + for item in data: + track_key = ET.SubElement(tracks_dict, 'key') + track_key.text = str(index) + track_dict = ET.SubElement(tracks_dict, 'dict') + track_details = ET.SubElement(track_dict, 'Track ID') + track_details.text = str(index) + index += 1 + for key, value in item.items(): + track_details = ET.SubElement(track_dict, key) + track_details.text = value + data = str(ET.tostring(library, encoding=self.encoding)) + #data = ET.dump(library) + if self.path: + self.export_to_file(data, **kwargs) + else: + self.export_to_terminal(data, **kwargs) + + def export_to_terminal(self, data, **kwargs): + print(data) + + def export_to_file(self, data, **kwargs): + with codecs.open(self.path, self.mode, self.encoding) as f: + f.write(data) \ No newline at end of file