8000 Merge pull request #44 from t8y8/tdsx-support · digideskio/document-api-python@aba2a35 · GitHub
[go: up one dir, main page]
Skip to content

Commit aba2a35

Browse files
t8y8t8y8
authored
Merge pull request tableau#44 from t8y8/tdsx-support
Added TDSX support and refactored Workbooks and Datasources to share zip-handling code
2 parents 23e897a + a885f15 commit aba2a35

File tree

5 files changed

+135
-98
lines changed

5 files changed

+135
-98
lines changed

‎tableaudocumentapi/datasource.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33
# Datasource - A class for writing datasources to Tableau files
44
#
55
###############################################################################
6+
import os
7+
import zipfile
8+
69
import xml.etree.ElementTree as ET
7-
from tableaudocumentapi import Connection
10+
from tableaudocumentapi import Connection, xfile
811

912

1013
class ConnectionParser(object):
@@ -56,7 +59,11 @@ def __init__(self, dsxml, filename=None):
5659
@classmethod
5760
def from_file(cls, filename):
5861
"Initialize datasource from file (.tds)"
59-
dsxml = ET.parse(filename).getroot()
62+
63+
if zipfile.is_zipfile(filename):
64+
dsxml = xfile.get_xml_from_archive(filename).getroot()
65+
else:
66+
dsxml = ET.parse(filename).getroot()
6067
return cls(dsxml, filename)
6168

6269
def save(self):
@@ -72,7 +79,8 @@ def save(self):
7279
"""
7380

7481
# save the file
75-
self._datasourceTree.write(self._filename, encoding="utf-8", xml_declaration=True)
82+
83+
xfile._save_file(self._filename, self._datasourceTree)
7684

7785
def save_as(self, new_filename):
7886
"""
@@ -85,7 +93,7 @@ def save_as(self, new_filename):
8593
Nothing.
8694
8795
"""
88-
self._datasourceTree.write(new_filename, encoding="utf-8", xml_declaration=True)
96+
xfile._save_file(self._filename, self._datasourceTree, new_filename)
8997

9098
###########
9199
# name

‎tableaudocumentapi/workbook.py

Lines changed: 6 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,12 @@
33
# Workbook - A class for writing Tableau workbook files
44
#
55
###############################################################################
6-
import contextlib
76
import os
8-
import shutil
9-
import tempfile
107
import zipfile
118

129
import xml.etree.ElementTree as ET
1310

14-
from tableaudocumentapi import Datasource
11+
from tableaudocumentapi import Datasource, xfile
1512

1613
###########################################################################
1714
#
@@ -20,41 +17,6 @@
2017
###########################################################################
2118

2219

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-
57-
5820
class Workbook(object):
5921
"""
6022
A class for writing Tableau workbook files.
@@ -75,7 +37,8 @@ def __init__(self, filename):
7537

7638
# Determine if this is a twb or twbx and get the xml root
7739
if zipfile.is_zipfile(self._filename):
78-
self._workbookTree = get_twb_xml_from_twbx(self._filename)
40+
self._workbookTree = xfile.get_xml_from_archive(
41+
self._filename)
7942
else:
8043
self._workbookTree = ET.parse(self._filename)
8144

@@ -111,12 +74,7 @@ def save(self):
11174
"""
11275

11376
# save the file
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)
77+
xfile._save_file(self._filename, self._workbookTree)
12078

12179
def save_as(self, new_filename):
12280
"""
@@ -129,12 +87,8 @@ def save_as(self, new_filename):
12987
Nothing.
13088
13189
"""
132-
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)
90+
xfile._save_file(
91+
self._filename, self._workbookTree, new_filename)
13892

13993
###########################################################################
14094
#
@@ -150,31 +104,3 @@ def _prepare_datasources(self, xmlRoot):
150104
datasources.append(ds)
151105

152106
return datasources
153-
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-
177-
@staticmethod
178-
def _is_valid_file(filename):
179-
fileExtension = os.path.splitext(filename)[-1].lower()
180-
return fileExtension in ('.twb', '.tds')

‎tableaudocumentapi/xfile.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import contextlib
2+
import os
3+
import shutil
4+
import tempfile
5+
import zipfile
6+
7+
import xml.etree.ElementTree as ET
8+
9+
10+
@contextlib.contextmanager
11+
def temporary_directory(*args, **kwargs):
12+
d = tempfile.mkdtemp(*args, **kwargs)
13+
try:
14+
yield d
15+
finally:
16+
shutil.rmtree(d)
17+
18+
19+
def find_file_in_zip(zip):
20+
for filename in zip.namelist():
21+
try:
22+
with zip.open(filename) as xml_candidate:
23+
ET.parse(xml_candidate).getroot().tag in (
24+
'workbook', 'datasource')
25+
return filename
26+
except ET.ParseError:
27+
# That's not an XML file by gosh
28+
pass
29+
30+
31+
def get_xml_from_archive(filename):
32+
with zipfile.ZipFile(filename) as zf:
33+
with zf.open(find_file_in_zip(zf)) as xml_file:
34+
xml_tree = ET.parse(xml_file)
35+
36+
return xml_tree
37+
38+
39+
def build_archive_file(archive_contents, zip):
40+
for root_dir, _, files in os.walk(archive_contents):
41+
relative_dir = os.path.relpath(root_dir, archive_contents)
42+
for f in files:
43+
temp_file_full_path = os.path.join(
44+
archive_contents, relative_dir, f)
45+
zipname = os.path.join(relative_dir, f)
46+
zip.write(temp_file_full_path, arcname=zipname)
47+
48+
49+
def save_into_archive(xml_tree, filename, new_filename=None):
50+
# Saving a archive means extracting the contents into a temp folder,
51+
# saving the changes over the twb/tds in that folder, and then
52+
# packaging it back up into a specifically formatted zip with the correct
53+
# relative file paths
54+
55+
if new_filename is None:
56+
new_filename = filename
57+
58+
# Extract to temp directory
59+
with temporary_directory() as temp_path:
60+
with zipfile.ZipFile(filename) as zf:
61+
xml_file = find_file_in_zip(zf)
62+
zf.extractall(temp_path)
63+
# Write the new version of the file to the temp directory
64+
xml_tree.write(os.path.join(
65+
temp_path, xml_file), encoding="utf-8", xml_declaration=True)
66+
67+
# Write the new archive with the contents of the temp folder
68+
with zipfile.ZipFile(new_filename, "w", compression=zipfile.ZIP_DEFLATED) as new_archive:
69+
build_archive_file(temp_path, new_archive)
70+
71+
72+
def _save_file(container_file, xml_tree, new_filename=None):
73+
if zipfile.is_zipfile(container_file):
74+
save_into_archive(xml_tree, container_file, new_filename)
75+
else:
76+
xml_tree.write(container_file, encoding="utf-8", xml_declaration=True)

‎test/assets/TABLEAU_10_TDSX.tdsx

1.82 KB
Binary file not shown.

‎test/bvt.py

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -15,22 +15,12 @@
1515

1616
TABLEAU_10_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TWB.twb')
1717

18-
TABLEAU_CONNECTION_XML = ET.parse(os.path.join(TEST_DIR, 'assets', 'CONNECTION.xml')).getroot()
18+
TABLEAU_CONNECTION_XML = ET.parse(os.path.join(
19+
TEST_DIR, 'assets', 'CONNECTION.xml')).getroot()
1920

2021
TABLEAU_10_TWBX = os.path.join( 10000 TEST_DIR, 'assets', 'TABLEAU_10_TWBX.twbx')
2122

22-
23-
class HelperMethodTests(unittest.TestCase):
24-
25-
def test_is_valid_file_with_valid_inputs(self):
26-
self.assertTrue(Workbook._is_valid_file('file1.tds'))
27-
self.assertTrue(Workbook._is_valid_file('file2.twb'))
28-
self.assertTrue(Workbook._is_valid_file('tds.twb'))
29-
30-
def test_is_valid_file_with_invalid_inputs(self):
31-
self.assertFalse(Workbook._is_valid_file(''))
32-
self.assertFalse(Workbook._is_valid_file('file1.tds2'))
33-
self.assertFalse(Workbook._is_valid_file('file2.twb3'))
23+
TABLEAU_10_TDSX = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TDSX.tdsx')
3424

3525

3626
class ConnectionParserTests(unittest.TestCase):
@@ -144,6 +134,43 @@ def test_can_save_tds(self):
144134
self.assertEqual(new_tds.connections[0].dbname, 'newdb.test.tsi.lan')
145135

146136

137+
class DatasourceModelV10TDSXTests(unittest.TestCase):
138+
139+
def setUp(self):
140+
with open(TABLEAU_10_TDSX, 'rb') as in_file, open('test.tdsx', 'wb') as out_file:
141+
out_file.write(in_file.read())
142+
self.tdsx_file = out_file
143+
144+
def tearDown(self):
145+
self.tdsx_file.close()
146+
os.unlink(self.tdsx_file.name)
147+
148+
def test_can_open_tdsx(self):
149+
ds = Datasource.from_file(self.tdsx_file.name)
150+
self.assertTrue(ds.connections)
151+
self.assertTrue(ds.name)
152+
153+
def test_can_open_tdsx_and_save_changes(self):
154+
original_tdsx = Datasource.from_file(self.tdsx_file.name)
155+
original_tdsx.connections[0].server = 'newdb.test.tsi.lan'
156+
original_tdsx.save()
157+
158+
new_tdsx = Datasource.from_file(self.tdsx_file.name)
159+
self.assertEqual(new_tdsx.connections[
160+
0].server, 'newdb.test.tsi.lan')
161+
162+
def test_can_open_tdsx_and_save_as_changes(self):
163+
new_tdsx_filename = 'newtdsx.tdsx'
164+
original_wb = Datasource.from_file(self.tdsx_file.name)
165+
original_wb.connections[0].server = 'newdb.test.tsi.lan'
166+
original_wb.save_as(new_tdsx_filename)
167+
168+
new_wb = Datasource.from_file(new_tdsx_filename)
169+
self.assertEqual(new_wb.connections[
170+
0].server, 'newdb.test.tsi.lan')
171+
os.unlink(new_tdsx_filename)
172+
173+
147174
class WorkbookModelTests(unittest.TestCase):
148175

149176
def setUp(self):
@@ -240,7 +267,7 @@ def test_can_open_twbx_and_save_changes(self):
240267
0].server, 'newdb.test.tsi.lan')
241268

242269
def test_can_open_twbx_and_save_as_changes(self):
243-
new_twbx_filename = self.workbook_file.name + "_TEST_SAVE_AS"
270+
new_twbx_filename = 'newtwbx.twbx'
244271
original_wb = Workbook(self.workbook_file.name)
245272
original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
246273
original_wb.save_as(new_twbx_filename)

0 commit comments

Comments
 (0)
0