8000 Merge pull request #43 from t8y8/zip-support · randyoswald/document-api-python@23e897a · GitHub
[go: up one dir, main page]
Skip to content

Commit 23e897a

Browse files
author
Russell Hay
authored
Merge pull request tableau#43 from t8y8/zip-support
Basic TWBX Support - LGTM
2 parents c26ee64 + 92668e0 commit 23e897a

File tree

8 files changed

+159
-49
lines changed

8 files changed

+159
-49
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._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,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()

‎test/assets/CONNECTION.xml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
<connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection>

‎test/assets/TABLEAU_10_TDS.tds

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
<?xml version='1.0' encoding='utf-8' ?><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource>

‎test/assets/TABLEAU_10_TWB.twb

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
<?xml version='1.0' encoding='utf-8' ?><workbook source-build='0.0.0 (0000.16.0510.1300)' source-platform='mac' version='10.0' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource></datasources></workbook>

‎test/assets/TABLEAU_10_TWBX.twbx

11.9 KB
Binary file not shown.

‎test/assets/TABLEAU_93_TDS.tds

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
<?xml version='1.0' encoding='utf-8' ?><datasource formatted-name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' inline='true' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource>

‎test/assets/TABLEAU_93_TWB.twb

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
<?xml version='1.0' encoding='utf-8' ?><workbook source-build='9.3.1 (9300.16.0510.0100)' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy (TestV1)' inline='true' name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' version='9.3'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource></datasources></workbook>

‎test/bvt.py

Lines changed: 62 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,23 @@
1-
import unittest
2-
import io
31
import os
2+
import unittest
3+
44
import xml.etree.ElementTree as ET
55

66
from tableaudocumentapi import Workbook, Datasource, Connection, ConnectionParser
77

8-
# Disable the 120 line limit because of the embedded XML on these lines
9-
# TODO: Move the XML into external files and load them when needed
8+
TEST_DIR = os.path.dirname(__file__)
9+
10+
TABLEAU_93_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TWB.twb')
1011

11-
TABLEAU_93_WORKBOOK = '''<?xml version='1.0' encoding='utf-8' ?><workbook source-build='9.3.1 (9300.16.0510.0100)' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy (TestV1)' inline='true' name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' version='9.3'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource></datasources></workbook>''' # noqa
12+
TABLEAU_93_TDS = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TDS.tds')
1213

13-
TABLEAU_93_TDS = '''<?xml version='1.0' encoding='utf-8' ?><datasource formatted-name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' inline='true' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource>''' # noqa
14+
TABLEAU_10_TDS = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TDS.tds')
1415

15-
TABLEAU_10_TDS = '''<?xml version='1.0' encoding='utf-8' ?><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource>''' # noqa
16+
TABLEAU_10_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TWB.twb')
1617

17-
TABLEAU_10_WORKBOOK = '''<?xml version='1.0' encoding='utf-8' ?><workbook source-build='0.0.0 (0000.16.0510.1300)' source-platform='mac' version='10.0' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource></datasources></workbook>''' # noqa
18+
TABLEAU_CONNECTION_XML = ET.parse(os.path.join(TEST_DIR, 'assets', 'CONNECTION.xml')).getroot()
1819

19-
TABLEAU_CONNECTION_XML = ET.fromstring(
20-
'''<connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection>''') # noqa
20+
TABLEAU_10_TWBX = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TWBX.twbx')
2121

2222

2323
class HelperMethodTests(unittest.TestCase):
@@ -36,14 +36,14 @@ def test_is_valid_file_with_invalid_inputs(self):
3636
class ConnectionParserTests(unittest.TestCase):
3737

3838
def test_can_extract_legacy_connection(self):
39-
parser = ConnectionParser(ET.fromstring(TABLEAU_93_TDS), '9.2')
39+
parser = ConnectionParser(ET.parse(TABLEAU_93_TDS), '9.2')
4040
connections = parser.get_connections()
4141
self.assertIsInstance(connections, list)
4242
self.assertIsInstance(connections[0], Connection)
4343
self.assertEqual(connections[0].dbname, 'TestV1')
4444

4545
def test_can_extract_federated_connections(self):
46-
parser = ConnectionParser(ET.fromstring(TABLEAU_10_TDS), '10.0')
46+
parser = ConnectionParser(ET.parse(TABLEAU_10_TDS), '10.0')
4747
connections = parser.get_connections()
4848
self.assertIsInstance(connections, list)
4949
self.assertIsInstance(connections[0], Connection)
@@ -76,9 +76,9 @@ def test_can_write_attributes_to_connection(self):
7676
class DatasourceModelTests(unittest.TestCase):
7777

7878
def setUp(self):
79-
self.tds_file = io.FileIO('test.tds', 'w')
80-
self.tds_file.write(TABLEAU_93_TDS.encode('utf8'))
81-
self.tds_file.seek(0)
79+
with open(TABLEAU_93_TDS, 'rb') as in_file, open('test.tds', 'wb') as out_file:
80+
out_file.write(in_file.read())
81+
self.tds_file = out_file
8282

8383
def tearDown(self):
8484
self.tds_file.close()
@@ -117,9 +117,9 @@ def test_save_has_xml_declaration(self):
117117
class DatasourceModelV10Tests(unittest.TestCase):
118118

119119
def setUp(self):
120-
self.tds_file = io.FileIO('test10.tds', 'w')
121-
self.tds_file.write(TABLEAU_10_TDS.encode('utf8'))
122-
self.tds_file.seek(0)
120+
with open(TABLEAU_10_TDS, 'rb') as in_file, open('test.twb', 'wb') as out_file:
121+
out_file.write(in_file.read())
122+
self.tds_file = out_file
123123

124124
def tearDown(self):
125125
self.tds_file.close()
@@ -147,9 +147,9 @@ def test_can_save_tds(self):
147147
class WorkbookModelTests(unittest.TestCase):
148148

149149
def setUp(self):
150-
self.workbook_file = io.FileIO('test.twb', 'w')
151-
self.workbook_file.write(TABLEAU_93_WORKBOOK.encode('utf8'))
152-
self.workbook_file.seek(0)
150+
with open(TABLEAU_93_TWB, 'rb') as in_file, open('test.twb', 'wb') as out_file:
151+
out_file.write(in_file.read())
152+
self.workbook_file = out_file
153153

154154
def tearDown(self):
155155
self.workbook_file.close()
@@ -175,9 +175,9 @@ def test_can_update_datasource_connection_and_save(self):
175175
class WorkbookModelV10Tests(unittest.TestCase):
176176

177177
def setUp(self):
178-
self.workbook_file = io.FileIO('testv10.twb', 'w')
179-
self.workbook_file.write(TABLEAU_10_WORKBOOK.encode('utf8'))
180-
self.workbook_file.seek(0)
178+
with open(TABLEAU_10_TWB, 'rb') as in_file, open('test.twb', 'wb') as out_file:
179+
out_file.write(in_file.read())
180+
self.workbook_file = out_file
181181

182182
def tearDown(self):
183183
self.workbook_file.close()
@@ -213,5 +213,43 @@ def test_save_has_xml_declaration(self):
213213
self.assertEqual(
214214
first_line, "<?xml version='1.0' encoding='utf-8'?>")
215215

216+
217+
class WorkbookModelV10TWBXTests(unittest.TestCase):
218+
219+
def setUp(self):
220+
with open(TABLEAU_10_TWBX, 'rb') as in_file, open('test.twbx', 'wb') as out_file:
221+
out_file.write(in_file.read())
222+
self.workbook_file = out_file
223+
224+
def tearDown(self):
225+
self.workbook_file.close()
226+
os.unlink(self.workbook_file.name)
227+
228+
def test_can_open_twbx(self):
229+
wb = Workbook(self.workbook_file.name)
230+
self.assertTrue(wb.datasources)
231+
self.assertTrue(wb.datasources[0].connections)
232+
233+
def test_can_open_twbx_and_save_changes(self):
234+
original_wb = Workbook(self.workbook_file.name)
235+
original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
236+
original_wb.save()
237+
238+
new_wb = Workbook(self.workbook_file.name)
239+
self.assertEqual(new_wb.datasources[0].connections[
240+
0].server, 'newdb.test.tsi.lan')
241+
242+
def test_can_open_twbx_and_save_as_changes(self):
243+
new_twbx_filename = self.workbook_file.name + "_TEST_SAVE_AS"
244+
original_wb = Workbook(self.workbook_file.name)
245+
original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
246+
original_wb.save_as(new_twbx_filename)
247+
248+
new_wb = Workbook(new_twbx_filename)
249+
self.assertEqual(new_wb.datasources[0].connections[
250+
0].server, 'newdb.test.tsi.lan')
251+
252+
os.unlink(new_twbx_filename)
253+
216254
if __name__ == '__main__':
217255
unittest.main()

0 commit comments

Comments
 (0)
0