1# -*- coding: utf-8 -*-
2# -----------------------------------------------------------------------------
3# Name:         musicxml/archiveTools.py
4# Purpose:      Tools for compressing and decompressing MusicXML files
5#
6# Authors:      Christopher Ariza
7#               Michael Scott Cuthbert
8#
9# Copyright:    Copyright © 2009, 2017 Michael Scott Cuthbert and the music21 Project
10# License:      BSD, see license.txt
11# -----------------------------------------------------------------------------
12'''
13Tools for compressing and decompressing musicxml files.
14'''
15import os
16import pathlib
17import zipfile
18from typing import Union
19
20from music21 import common
21from music21 import environment
22_MOD = 'musicxml.archiveTools'
23environLocal = environment.Environment(_MOD)
24
25
26# -----------------------------------------------------------------------------
27# compression
28
29
30def compressAllXMLFiles(*, deleteOriginal=False):
31    '''
32    Takes all filenames in corpus.paths and runs
33    :meth:`music21.musicxml.archiveTools.compressXML` on each.  If the musicXML files are
34    compressed, the originals are deleted from the system.
35    '''
36    from music21.corpus.corpora import CoreCorpus
37    environLocal.warn("Compressing musicXML files...")
38
39    # this gets all .xml, .musicxml, .mxl etc.
40    for filename in CoreCorpus().getPaths(fileExtensions=('.xml',)):
41        compressXML(filename, deleteOriginal=deleteOriginal)
42    environLocal.warn(
43        'Compression complete. '
44        'Run the main test suite, fix bugs if necessary,'
45        'and then commit modified directories in corpus.'
46    )
47
48
49def compressXML(filename: Union[str, pathlib.Path],
50                *,
51                deleteOriginal=False,
52                silent=False,
53                strictMxlCheck=True) -> bool:
54    '''
55    Takes a filename, and if the filename corresponds to a musicXML file with
56    an .musicxml or .xml extension, creates a corresponding compressed .mxl file in the same
57    directory.
58
59    If deleteOriginal is set to True, the original musicXML file is deleted
60    from the system.
61
62    If strictMxlCheck is False then any suffix will do.
63
64    Returns bool if successful.
65    '''
66    filename = str(filename)
67    if strictMxlCheck and not filename.endswith('.xml') and not filename.endswith('.musicxml'):
68        return False  # not a musicXML file
69    fp = common.pathTools.cleanpath(filename, returnPathlib=True)
70    if not silent:  # pragma: no cover
71        environLocal.warn(f"Updating file: {fp}")
72    newFilename = str(fp.with_suffix('.mxl'))
73
74
75    # contents of container.xml file in META-INF folder
76    container = f'''<?xml version="1.0" encoding="UTF-8"?>
77<container>
78  <rootfiles>
79    <rootfile full-path="{fp.name}"/>
80  </rootfiles>
81</container>
82    '''
83    # Export container and original xml file to system as a compressed XML.
84    with zipfile.ZipFile(
85            newFilename,
86            'w',
87            compression=zipfile.ZIP_DEFLATED,
88    ) as myZip:
89        myZip.write(filename, fp.name)
90        myZip.writestr(
91            'META-INF' + os.path.sep + 'container.xml',
92            container,
93        )
94    # Delete uncompressed xml file from system
95    if deleteOriginal:
96        fp.unlink()
97
98    return True
99
100
101def uncompressMXL(filename: Union[str, pathlib.Path],
102                  *,
103                  deleteOriginal=False,
104                  strictMxlCheck=True) -> bool:
105    '''
106    Takes a filename, and if the filename corresponds to a compressed musicXML
107    file with an .mxl extension, creates a corresponding uncompressed .musicxml file
108    in the same directory.
109
110    If deleteOriginal is set to True, the original compressed musicXML file is
111    deleted from the system.
112
113    If strictMxlCheck is False then any type of file will attempt to be extracted.
114
115    Returns bool if successful.
116    '''
117    filename = str(filename)
118    if not filename.endswith('.mxl') and strictMxlCheck:
119        return  # not a compressed musicXML file
120
121    fp: pathlib.Path = common.pathTools.cleanpath(filename, returnPathlib=True)
122    environLocal.warn(f"Updating file: {fp}")
123    extractPath = str(fp.parent)
124    unarchivedName = fp.with_suffix('.musicxml').name
125    # Export container and original xml file to system as a compressed XML.
126    with zipfile.ZipFile(filename, 'r', compression=zipfile.ZIP_DEFLATED) as myZip:
127        try:
128            myZip.extract(member=unarchivedName, path=extractPath)
129        except KeyError:
130            try:
131                unarchivedName = unarchivedName.replace('.xml', '.musicxml')
132                myZip.extract(member=unarchivedName, path=extractPath)
133            except KeyError:
134                found_one_file = False
135                for storedName in myZip.namelist():
136                    if 'META-INF' in storedName:
137                        continue
138                    myZip.extract(member=storedName, path=extractPath)
139                    if not found_one_file:
140                        # only rename one file...hope it is the right one.
141                        extractPath_pathlib = pathlib.Path(extractPath)
142                        wrongName = extractPath_pathlib / storedName
143                        correctName = extractPath_pathlib / unarchivedName
144                        wrongName.rename(correctName)
145                        found_one_file = True
146
147
148    # Delete uncompressed xml file from system
149    if deleteOriginal:
150        fp.unlink()
151
152
153if __name__ == '__main__':
154    import sys
155    if len(sys.argv) >= 1:
156        for xmlName in sys.argv[1:]:
157            if xmlName.endswith('.xml'):
158                compressXML(xmlName)
159            elif xmlName.endswith('.musicxml'):
160                compressXML(xmlName)
161            elif xmlName.endswith('.mxl'):
162                uncompressMXL(xmlName)
163