8000 Initial TDSX support. Lots of duplicated code with TWBX's · csweezy/document-api-python@9a8f0d1 · GitHub
[go: up one dir, main page]
Skip to content

Commit 9a8f0d1

Browse files
t8y8t8y8
committed
Initial TDSX support. Lots of duplicated code with TWBX's
1 parent 23e897a commit 9a8f0d1

File tree

3 files changed

+122
-4
lines changed

3 files changed

+122
-4
lines changed

‎tableaudocumentapi/datasource.py

Lines changed: 80 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,51 @@
33
# Datasource - A class for writing datasources to Tableau files
44
#
55
###############################################################################
6+
import contextlib
7+
import os
8+
import shutil
9+
import tempfile
10+
import zipfile
11+
612
import xml.etree.ElementTree as ET
713
from tableaudocumentapi import Connection
814

915

16+
@contextlib.contextmanager
17+
def temporary_directory(*args, **kwargs):
18+
d = tempfile.mkdtemp(*args, **kwargs)
19+
try:
20+
yield d
21+
finally:
22+
shutil.rmtree(d)
23+
24+
25+
def find_tds_in_zip(zip):
26+
for filename in zip.namelist():
27+
if os.path.splitext(filename)[-1].lower() == '.tds':
28+
return filename
29+
30+
31+
def get_tds_xml_from_tdsx(filename):
32+
with temporary_directory() as temp:
33+
with zipfile.ZipFile(filename) as zf:
34+
zf.extractall(temp)
35+
tds_file = find_tds_in_zip(zf)
36+
tds_xml = ET.parse(os.path.join(temp, tds_file))
37+
38+
return tds_xml
39+
40+
41+
def build_tdsx_file(tdsx_contents, zip):
42+
for root_dir, _, files in os.walk(tdsx_contents):
43+
relative_dir = os.path.relpath(root_dir, tdsx_contents)
44+
for f in files:
45+
temp_file_full_path = os.path.join(
46+
tdsx_contents, relative_dir, f)
47+
zipname = os.path.join(relative_dir, f)
48+
zip.write(temp_file_full_path, arcname=zipname)
49+
50+
1051
class ConnectionParser(object):
1152

1253
def __init__(self, datasource_xml, version):
@@ -56,9 +97,36 @@ def __init__(self, dsxml, filename=None):
5697
@classmethod
5798
def from_file(cls, filename):
5899
"Initialize datasource from file (.tds)"
59-
dsxml = ET.parse(filename).getroot()
100+
101+
if zipfile.is_zipfile(filename):
102+
dsxml = get_tds_xml_from_tdsx(filename).getroot()
103+
else:
104+
dsxml = ET.parse(filename).getroot()
60105
return cls(dsxml, filename)
61106

107+
def _save_into_tdsx(self, filename=None):
108+
# Save reuses existing filename, 'save as' takes a new one
109+
if filename is None:
110+
filename = self._filename
111+
112+
# Saving a tdsx means extracting the contents into a temp folder,
113+
# saving the changes over the tds in that folder, and then
114+
# packaging it back up into a specifically formatted zip with the correct
115+
# relative file paths
116+
117+
# Extract to temp directory
118+
with temporary_directory() as temp_path:
119+
with zipfile.ZipFile(self._filename) as zf:
120+
tds_file = find_tds_in_zip(zf)
121+
zf.extractall(temp_path)
122+
# Write the new version of the tds to the temp directory
123+
self._datasourceTree.write(os.path.join(
124+
temp_path, tds_file), encoding="utf-8", xml_declaration=True)
125+
126+
# Write the new tdsx with the contents of the temp folder
127+
with zipfile.ZipFile(filename, "w", compression=zipfile.ZIP_DEFLATED) as new_tdsx:
128+
build_tdsx_file(temp_path, new_tdsx)
129+
62130
def save(self):
63131
"""
64132
Call finalization code and save file.
@@ -72,7 +140,12 @@ def save(self):
72140
"""
73141

74142
# save the file
75-
self._datasourceTree.write(self._filename, encoding="utf-8", xml_declaration=True)
143+
144+
if zipfile.is_zipfile(self._filename):
145+
self._save_into_tdsx(self._filename)
146+
else:
147+
self._datasourceTree.write(
148+
self._filename, encoding="utf-8", xml_declaration=True)
76149

77150
def save_as(self, new_filename):
78151
"""
@@ -85,7 +158,11 @@ def save_as(self, new_filename):
85158
Nothing.
86159
87160
"""
88-
self._datasourceTree.write(new_filename, encoding="utf-8", xml_declaration=True)
161+
if zipfile.is_zipfile(self._filename):
162+
self._save_into_tdsx(new_filename)
163+
else:
164+
self._datasourceTree.write(
165+
new_filename, encoding="utf-8", xml_declaration=True)
89166

90167
###########
91168
# name

‎test/assets/TABLEAU_10_TDSX.tdsx

1.82 KB
Binary file not shown.

‎test/bvt.py

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,13 @@
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(TEST_DIR, 'assets', 'TABLEAU_10_TWBX.twbx')
2122

23+
TABLEAU_10_TDSX = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TDSX.tdsx')
24+
2225

2326
class HelperMethodTests(unittest.TestCase):
2427

@@ -144,6 +147,44 @@ def test_can_save_tds(self):
144147
self.assertEqual(new_tds.connections[0].dbname, 'newdb.test.tsi.lan')
145148

146149

150+
class DatasourceModelV10TDSXTests(unittest.TestCase):
151+
152+
def setUp(self):
153+
with open(TABLEAU_10_TDSX, 'rb') as in_file, open('test.tdsx', 'wb') as out_file:
154+
out_file.write(in_file.read())
155+
self.tdsx_file = out_file
156+
157+
def tearDown(self):
158+
self.tdsx_file.close()
159+
os.unlink(self.tdsx_file.name)
160+
161+
def test_can_open_tdsx(self):
162+
ds = Datasource.from_file(self.tdsx_file.name)
163+
self.assertTrue(ds.connections)
164+
self.assertTrue(ds.name)
165+
166+
def test_can_open_tdsx_and_save_changes(self):
167+
original_tdsx = Datasource.from_file(self.tdsx_file.name)
168+
original_tdsx.connections[0].server = 'newdb.test.tsi.lan'
169+
original_tdsx.save()
170+
171+
new_tdsx = Datasource.from_file(self.tdsx_file.name)
172+
self.assertEqual(new_tdsx.connections[
173+
0].server, 'newdb.test.tsi.lan')
174+
175+
def test_can_open_tdsx_and_save_as_changes(self):
176+
new_tdsx_filename = self.tdsx_file.name + "_TEST_SAVE_AS"
177+
original_wb = Datasource.from_file(self.tdsx_file.name)
178+
original_wb.connections[0].server = 'newdb.test.tsi.lan'
179+
original_wb.save_as(new_tdsx_filename)
180+
181+
new_wb = Datasource.from_file(new_tdsx_filename)
182+
self.assertEqual(new_wb.connections[
183+
0].server, 'newdb.test.tsi.lan')
184+
185+
os.unlink(new_tdsx_filename)
186+
187+
147188
class WorkbookModelTests(unittest.TestCase):
148189

149190
def setUp(self):

0 commit comments

Comments
 (0)
0