8000 Here's an intermediate stab at twbx support. Does not include TDSX ye… · onware/document-api-python@6663c87 · GitHub
[go: up one dir, main page]
Skip to content

Commit 6663c87

Browse files
t8y8t8y8
committed
Here's an intermediate stab at twbx support. Does not include TDSX yet. I'm not sure where a lot of the zip stuff should go.
1 parent 07aad95 commit 6663c87

File tree

2 files changed

+137
-30
lines changed

2 files changed

+137
-30
lines changed

‎tableaudocumentapi/workbook.py

Copy file name to clipboard
Lines changed: 89 additions & 26 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._filename = 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,11 @@ def save(self):
76111
"""
77112

78113
# save the file
79-
self._workbookTree.write(self._filename)
114+
115+
if zipfile.is_zipfile(self._filename):
116+
self._save_into_twbx(self._filename)
117+
else:
118+
self._workbookTree.write(self._filename)
80119

81120
def save_as(self, new_filename):
82121
"""
@@ -89,8 +128,10 @@ def save_as(self, new_filename):
89128
Nothing.
90129
91130
"""
92-
93-
self._workbookTree.write(new_filename)
131+
if zipfile.is_zipfile(self._filename):
132+
self._save_into_twbx(new_filename)
133+
else:
134+
self._workbookTree.write(new_filename)
94135

95136
###########################################################################
96137
#
@@ -107,6 +148,28 @@ def _prepare_datasources(self, xmlRoot):
107148

108149
return datasources
109150

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

0 commit comments

Comments
 (0)
0