Merge pull request #3400 from austinmm/Extended_Export_Plugin_Support

Extended export plugin support

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
@@ -44,14 +46,27 @@ def __init__(self):
         self.config.add({
             'default_format': 'json',
             'json': {
-                # json module formatting options
+                # JSON module formatting options.
                 'formatting': {
                     'ensure_ascii': False,
                     'indent': 4,
                     'separators': (',', ': '),
                     'sort_keys': True
                 }
             },
+            'csv': {
+                # CSV module formatting options.
+                'formatting': {
+                    # The delimiter used to seperate columns.
+                    'delimiter': ',',
+                    # The dialect to use when formating the file output.
+                    'dialect': 'excel'
+                }
+            },
+            'xml': {
+                # XML module formatting options.
+                'formatting': {}
+            }
             # TODO: Use something like the edit plugin
             # 'item_fields': []
         })
@@ -78,17 +93,21 @@ def commands(self):
             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"the output format: json (default), csv, or 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 or self.config['default_format'].get(str)
         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
             }
@@ -100,6 +119,7 @@ def run(self, lib, opts, args):
         included_keys = []
         for keys in opts.included_keys:
             included_keys.extend(keys.split(','))
+
         key_filter = make_key_filter(included_keys)
 
         for data_emitter in data_collector(lib, ui.decargs(args)):
@@ -117,35 +137,69 @@ def run(self, lib, opts, args):
 
 class ExportFormat(object):
     """The output format type"""
+    def __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'):
+        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 = codecs.open(self.path, self.mode, self.encoding) \
+            if self.path else sys.stdout
 
     @classmethod
-    def factory(cls, type, **kwargs):
-        if type == "json":
-            if kwargs['file_path']:
-                return JsonFileFormat(**kwargs)
-            else:
-                return JsonPrintFormat()
-        raise NotImplementedError()
+    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 JsonPrintFormat(ExportFormat):
-    """Outputs to the console"""
+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)
 
     def export(self, data, **kwargs):
-        json.dump(data, sys.stdout, cls=ExportEncoder, **kwargs)
+        json.dump(data, self.out_stream, cls=ExportEncoder, **kwargs)
 
 
-class JsonFileFormat(ExportFormat):
-    """Saves in a json file"""
+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)
 
+    def export(self, data, **kwargs):
+        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 __init__(self, file_path, file_mode=u'w', encoding=u'utf-8'):
-        self.path = file_path
-        self.mode = file_mode
-        self.encoding = encoding
+        super(XMLFormat, self).__init__(file_path, file_mode, encoding)
 
     def export(self, data, **kwargs):
-        with codecs.open(self.path, self.mode, self.encoding) as f:
-            json.dump(data, f, cls=ExportEncoder, **kwargs)
+        # Creates the XML file structure.
+        library = ET.Element(u'library')
+        tracks = ET.SubElement(library, u'tracks')
+        if data and isinstance(data[0], dict):
+            for index, item in enumerate(data):
+                track = ET.SubElement(tracks, u'track')
+                for key, value in item.items():
+                    track_details = ET.SubElement(track, key)
+                    track_details.text = value
+        # Depending on the version of python the encoding needs to change
+        try:
+            data = ET.tostring(library, encoding='unicode', **kwargs)
+        except LookupError:
+            data = ET.tostring(library, encoding='utf-8', **kwargs)
+
+        self.out_stream.write(data)

docs/changelog.rst

@@ -6,6 +6,10 @@ Changelog
 
 New features:
 
+* :doc:`/plugins/export`: Added new ``-f`` (``--format``) flag; 
+  which allows for the ability to export in json, csv and xml.
+  Thanks to :user:`austinmm`.
+  :bug:`3402`
 * :doc:`/plugins/unimported`: lets you find untracked files in your library directory.
 * We now fetch information about `works`_ from MusicBrainz.
   MusicBrainz matches provide the fields ``work`` (the title), ``mb_workid``

docs/plugins/export.rst

@@ -2,9 +2,11 @@ Export Plugin
 =============
 
 The ``export`` plugin lets you get data from the items and export the content
-as `JSON`_.
+as `JSON`_, `CSV`_, or `XML`_.
 
 .. _JSON: https://www.json.org
+.. _CSV: https://fileinfo.com/extension/csv
+.. _XML: https://fileinfo.com/extension/xml
 
 Enable the ``export`` plugin (see :ref:`using-plugins` for help). Then, type ``beet export`` followed by a :doc:`query </reference/query>` to get the data from
 your library. For example, run this::
@@ -13,6 +15,7 @@ your library. For example, run this::
 
 to print a JSON file containing information about your Beatles tracks.
 
+
 Command-Line Options
 --------------------
 
@@ -36,30 +39,42 @@ The ``export`` command has these command-line options:
 
 * ``--append``: Appends the data to the file instead of writing.
 
+* ``--format`` or ``-f``: Specifies the format the data will be exported as. If not informed, JSON will be used by default. The format options include csv, json and xml.
+
 Configuration
 -------------
 
 To configure the plugin, make a ``export:`` section in your configuration
-file. Under the ``json`` key, these options are available:
+file.
+For JSON export, these options are available under the ``json`` key:
 
 - **ensure_ascii**: Escape non-ASCII characters with ``\uXXXX`` entities.
-
 - **indent**: The number of spaces for indentation.
-
 - **separators**: A ``[item_separator, dict_separator]`` tuple.
-
 - **sort_keys**: Sorts the keys in JSON dictionaries.
 
-These options match the options from the `Python json module`_.
+Those options match the options from the `Python json module`_.
+Similarly, these options are available for the CSV format under the ``csv``
+key:
+
+- **delimiter**: Used as the separating character between fields. The default value is a comma (,).
+- **dialect**: The kind of CSV file to produce. The default is `excel`.
+
+These options match the options from the `Python csv module`_.
 
 .. _Python json module: https://docs.python.org/2/library/json.html#basic-usage
+.. _Python csv module: https://docs.python.org/3/library/csv.html#csv-fmt-params
 
 The default options look like this::
 
     export:
         json:
             formatting:
-                ensure_ascii: False
+                ensure_ascii: false
                 indent: 4
                 separators: [',' , ': ']
                 sort_keys: true
+        csv:
+            formatting:
+                delimiter: ','
+                dialect: excel

test/test_export.py

@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+# This file is part of beets.
+# Copyright 2019, Carl Suster
+#
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+#
+# The above copyright notice and this permission notice shall be
+# included in all copies or substantial portions of the Software.
+
+"""Test the beets.export utilities associated with the export plugin.
+"""
+
+from __future__ import division, absolute_import, print_function
+
+import unittest
+from test.helper import TestHelper
+import re  # used to test csv format
+import json
+from xml.etree.ElementTree import Element
+import xml.etree.ElementTree as ET
+
+
+class ExportPluginTest(unittest.TestCase, TestHelper):
+    def setUp(self):
+        self.setup_beets()
+        self.load_plugins('export')
+        self.test_values = {'title': 'xtitle', 'album': 'xalbum'}
+
+    def tearDown(self):
+        self.unload_plugins()
+        self.teardown_beets()
+
+    def execute_command(self, format_type, artist):
+        query = ','.join(self.test_values.keys())
+        out = self.run_with_output(
+            'export',
+            '-f', format_type,
+            '-i', query,
+            artist
+        )
+        return out
+
+    def create_item(self):
+        item, = self.add_item_fixtures()
+        item.artist = 'xartist'
+        item.title = self.test_values['title']
+        item.album = self.test_values['album']
+        item.write()
+        item.store()
+        return item
+
+    def test_json_output(self):
+        item1 = self.create_item()
+        out = self.execute_command(
+            format_type='json',
+            artist=item1.artist
+        )
+        json_data = json.loads(out)[0]
+        for key, val in self.test_values.items():
+            self.assertTrue(key in json_data)
+            self.assertEqual(val, json_data[key])
+
+    def test_csv_output(self):
+        item1 = self.create_item()
+        out = self.execute_command(
+            format_type='csv',
+            artist=item1.artist
+        )
+        csv_list = re.split('\r', re.sub('\n', '', out))
+        head = re.split(',', csv_list[0])
+        vals = re.split(',|\r', csv_list[1])
+        for index, column in enumerate(head):
+            self.assertTrue(self.test_values.get(column, None) is not None)
+            self.assertEqual(vals[index], self.test_values[column])
+
+    def test_xml_output(self):
+        item1 = self.create_item()
+        out = self.execute_command(
+            format_type='xml',
+            artist=item1.artist
+        )
+        library = ET.fromstring(out)
+        self.assertIsInstance(library, Element)
+        for track in library[0]:
+            for details in track:
+                tag = details.tag
+                txt = details.text
+                self.assertTrue(tag in self.test_values, msg=tag)
+                self.assertEqual(self.test_values[tag], txt, msg=txt)
+
+
+def suite():
+    return unittest.TestLoader().loadTestsFromName(__name__)
+
+if __name__ == '__main__':
+    unittest.main(defaultTest='suite')