1from __future__ import absolute_import
2# Copyright (c) 2010-2019 openpyxl
3
4"""
5File manifest
6"""
7from mimetypes import MimeTypes
8import os.path
9
10from openpyxl.descriptors.serialisable import Serialisable
11from openpyxl.descriptors import String, Sequence
12from openpyxl.xml.functions import fromstring
13from openpyxl.xml.constants import (
14    ARC_CORE,
15    ARC_CONTENT_TYPES,
16    ARC_WORKBOOK,
17    ARC_APP,
18    ARC_THEME,
19    ARC_STYLE,
20    ARC_SHARED_STRINGS,
21    EXTERNAL_LINK,
22    THEME_TYPE,
23    STYLES_TYPE,
24    XLSX,
25    XLSM,
26    XLTM,
27    XLTX,
28    WORKSHEET_TYPE,
29    COMMENTS_TYPE,
30    SHARED_STRINGS,
31    DRAWING_TYPE,
32    CHART_TYPE,
33    CHARTSHAPE_TYPE,
34    CHARTSHEET_TYPE,
35    CONTYPES_NS,
36    ACTIVEX,
37    CTRL,
38    VBA,
39)
40from openpyxl.xml.functions import tostring
41
42# initialise mime-types
43mimetypes = MimeTypes()
44mimetypes.add_type('application/xml', ".xml")
45mimetypes.add_type('application/vnd.openxmlformats-package.relationships+xml', ".rels")
46mimetypes.add_type("application/vnd.ms-office.vbaProject", ".bin")
47mimetypes.add_type("application/vnd.openxmlformats-officedocument.vmlDrawing", ".vml")
48mimetypes.add_type("image/x-emf", ".emf")
49
50
51class FileExtension(Serialisable):
52
53    tagname = "Default"
54
55    Extension = String()
56    ContentType = String()
57
58    def __init__(self, Extension, ContentType):
59        self.Extension = Extension
60        self.ContentType = ContentType
61
62
63class Override(Serialisable):
64
65    tagname = "Override"
66
67    PartName = String()
68    ContentType = String()
69
70    def __init__(self, PartName, ContentType):
71        self.PartName = PartName
72        self.ContentType = ContentType
73
74
75DEFAULT_TYPES = [
76    FileExtension("rels", "application/vnd.openxmlformats-package.relationships+xml"),
77    FileExtension("xml", "application/xml"),
78]
79
80DEFAULT_OVERRIDE = [
81    Override("/" + ARC_STYLE, STYLES_TYPE), # Styles
82    Override("/" + ARC_THEME, THEME_TYPE), # Theme
83    Override("/docProps/core.xml", "application/vnd.openxmlformats-package.core-properties+xml"),
84    Override("/docProps/app.xml", "application/vnd.openxmlformats-officedocument.extended-properties+xml")
85]
86
87
88class Manifest(Serialisable):
89
90    tagname = "Types"
91
92    Default = Sequence(expected_type=FileExtension, unique=True)
93    Override = Sequence(expected_type=Override, unique=True)
94    path = "[Content_Types].xml"
95
96    __elements__ = ("Default", "Override")
97
98    def __init__(self,
99                 Default=(),
100                 Override=(),
101                 ):
102        if not Default:
103            Default = DEFAULT_TYPES
104        self.Default = Default
105        if not Override:
106            Override = DEFAULT_OVERRIDE
107        self.Override = Override
108
109
110    @property
111    def filenames(self):
112        return [part.PartName for part in self.Override]
113
114
115    @property
116    def extensions(self):
117        """
118        Map content types to file extensions
119        Skip parts without extensions
120        """
121        exts = set([os.path.splitext(part.PartName)[-1] for part in self.Override])
122        return [(ext[1:], mimetypes.types_map[True][ext]) for ext in sorted(exts) if ext]
123
124
125    def to_tree(self):
126        """
127        Custom serialisation method to allow setting a default namespace
128        """
129        defaults = [t.Extension for t in self.Default]
130        for ext, mime in self.extensions:
131            if ext not in defaults:
132                mime = FileExtension(ext, mime)
133                self.Default.append(mime)
134        tree = super(Manifest, self).to_tree()
135        tree.set("xmlns", CONTYPES_NS)
136        return tree
137
138
139    def __contains__(self, content_type):
140        """
141        Check whether a particular content type is contained
142        """
143        for t in self.Override:
144            if t.ContentType == content_type:
145                return True
146
147
148    def find(self, content_type):
149        """
150        Find specific content-type
151        """
152        try:
153            return next(self.findall(content_type))
154        except StopIteration:
155            return
156
157
158    def findall(self, content_type):
159        """
160        Find all elements of a specific content-type
161        """
162        for t in self.Override:
163            if t.ContentType == content_type:
164                yield t
165
166
167    def append(self, obj):
168        """
169        Add content object to the package manifest
170        # needs a contract...
171        """
172        ct = Override(PartName=obj.path, ContentType=obj.mime_type)
173        self.Override.append(ct)
174
175
176    def _write(self, archive, workbook):
177        """
178        Write manifest to the archive
179        """
180        self.append(workbook)
181        self._write_vba(workbook)
182        self._register_mimetypes(filenames=archive.namelist())
183        archive.writestr(self.path, tostring(self.to_tree()))
184
185
186    def _register_mimetypes(self, filenames):
187        """
188        Make sure that the mime type for all file extensions is registered
189        """
190        for fn in filenames:
191            ext = os.path.splitext(fn)[-1]
192            if not ext:
193                continue
194            mime = mimetypes.types_map[True][ext]
195            fe = FileExtension(ext[1:], mime)
196            self.Default.append(fe)
197
198
199    def _write_vba(self, workbook):
200        """
201        Add content types from cached workbook when keeping VBA
202        """
203        if workbook.vba_archive:
204            node = fromstring(workbook.vba_archive.read(ARC_CONTENT_TYPES))
205            mf = Manifest.from_tree(node)
206            filenames = self.filenames
207            for override in mf.Override:
208                if override.PartName not in (ACTIVEX, CTRL, VBA):
209                    continue
210                if override.PartName not in filenames:
211                    self.Override.append(override)
212