8000 Merge branch 'zip-support' into development · fossabot/document-api-python@76344b4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 76344b4

Browse files
t8y8t8y8
committed
Merge branch 'zip-support' into development
2 parents 807fa9d + cbb596b commit 76344b4

File tree

2 files changed

+135
-26
lines changed

2 files changed

+135
-26
lines changed

‎tableaudocumentapi/workbook.py

Lines changed: 92 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,57 @@
33
# Workbook - A class for writing Tableau workbook files
44
#
55
###############################################################################
6+
import contextlib
67
import os
8+
import shutil
9+
import tempfile
10+
import zipfile
11+
712
import xml.etree.ElementTree as ET
13+
814
from tableaudocumentapi import Datasource
915

16+
###########################################################################
17+
#
18+
# Utility Functions
19+
#
20+
###########################################################################
21+
22+
23+
@contextlib.contextmanager
24+
def temporary_directory(*args, **kwargs):
25+
d = tempfile.mkdtemp(*args, **kwargs)
26+
try:
27+
yield d
28+
finally:
29+
shutil.rmtree(d)
30+
31+
32+
def find_twb_in_zip(zip):
33+
for filename in zip.namelist():
34+
if os.path.splitext(filename)[-1].lower() == '.twb':
35+
return filename
36+
37+
38+
def get_twb_xml_from_twbx(filename):
39+
with temporary_directory() as temp:
40+
with zipfile.ZipFile(filename) as zf:
41+
zf.extractall(temp)
42+
twb_file = find_twb_in_zip(zf)
43+
twb_xml = ET.parse(os.path.join(temp, twb_file))
44+
45+
return twb_xml
46+
47+
48+
def build_twbx_file(twbx_contents, zip):
49+
for root_dir, _, files in os.walk(twbx_contents):
50+
relative_dir = os.path.relpath(root_dir, twbx_contents)
51+
for f in files:
52+
temp_file_full_path = os.path.join(
53+
twbx_contents, relative_dir, f)
54+
zipname = os.path.join(relative_dir, f)
55+
zip.write(temp_file_full_path, arcname=zipname)
56+
1057

1158
class Workbook(object):
1259
"""
@@ -24,30 +71,18 @@ def __init__(self, filename):
2471
Constructor.
2572
2673
"""
27-
# We have a valid type of input file
28-
if self._is_valid_file(filename):
29-
# set our filename, open .twb, initialize things
30-
self._filename = filename
31-
self._workbookTree = ET.parse(filename)
32-
self._workbookRoot = self._workbookTree.getroot()
33-
34-
# prepare our datasource objects
35-
self._datasources = self._prepare_datasources(
36-
self._workbookRoot) # self.workbookRoot.find('datasources')
37-
else:
38-
print('Invalid file type. Must be .twb or .tds.')
39-
raise Exception()
40-
41-
@classmethod
42-
def from_file(cls, filename):
43-
"Initialize datasource from file (.tds)"
44-
if self._is_valid_file(filename):
45-
self._filename = filename
46-
dsxml = ET.parse(filename).getroot()
47-
return cls(dsxml)
74+
self._fi EDBE lename = filename
75+
76+
# Determine if this is a twb or twbx and get the xml root
77+
if zipfile.is_zipfile(self._filename):
78+
self._workbookTree = get_twb_xml_from_twbx(self._filename)
4879
else:
49-
print('Invalid file type. Must be .twb or .tds.')
50-
raise Exception()
80+
self._workbookTree = ET.parse(self._filename)
81+
82+
self._workbookRoot = self._workbookTree.getroot()
83+
# prepare our datasource objects
84+
self._datasources = self._prepare_datasources(
85+
self._workbookRoot) # self.workbookRoot.find('datasources')
5186

5287
###########
5388
# datasources
@@ -76,7 +111,12 @@ def save(self):
76111
"""
77112

78113
# save the file
79-
self._workbookTree.write(self._filename, encoding="utf-8", xml_declaration=True)
114+
115+
if zipfile.is_zipfile(self._filename):
116+
self._save_into_twbx(self._filename)
117+
else:
118+
self._workbookTree.write(
119+
self._filename, encoding="utf-8", xml_declaration=True)
80120

81121
def save_as(self, new_filename):
82122
"""
@@ -90,7 +130,11 @@ def save_as(self, new_filename):
90130
91131
"""
92132

93-
self._workbookTree.write(new_filename, encoding="utf-8", xml_declaration=True)
133+
if zipfile.is_zipfile(self._filename):
134+
self._save_into_twbx(new_filename)
135+
else:
136+
self._workbookTree.write(
137+
new_filename, encoding="utf-8", xml_declaration=True)
94138

95139
###########################################################################
96140
#
@@ -107,6 +151,29 @@ def _prepare_datasources(self, xmlRoot):
107151

108152
return datasources
109153

154+
def _save_into_twbx(self, filename=None):
155+
# Save reuses existing filename, 'save as' takes a new one
156+
if filename is None:
157+
filename = self._filename
158+
159+
# Saving a twbx means extracting the contents into a temp folder,
160+
# saving the changes over the twb in that folder, and then
161+
# packaging it back up into a specifically formatted zip with the correct
162+
# relative file paths
163+
164+
# Extract to temp directory
165+
with temporary_directory() as temp_path:
166+
with zipfile.ZipFile(self._filename) as zf:
167+
twb_file = find_twb_in_zip(zf)
168+
zf.extractall(temp_path)
169+
# Write the new version of the twb to the temp directory
170+
self._workbookTree.write(os.path.join(
171+
temp_path, twb_file), encoding="utf-8", xml_declaration=True)
172+
173+
# Write the new twbx with the contents of the temp folder
174+
with zipfile.ZipFile(filename, "w", compression=zipfile.ZIP_DEFLATED) as new_twbx:
175+
build_twbx_file(temp_path, new_twbx)
176+
110177
@staticmethod
111178
def _is_valid_file(filename):
112179
fileExtension = os.path.splitext(filename)[-1].lower()

0 commit comments

Comments
 (0)
0