8000 Refactor all the archive manipulation logic into one module · digideskio/document-api-python@bd6d3c9 · GitHub
[go: up one dir, main page]
Skip to content

Commit bd6d3c9

Browse files
t8y8t8y8
committed
Refactor all the archive manipulation logic into one module
1 parent 9a8f0d1 commit bd6d3c9

File tree

4 files changed

+81
-133
lines changed

4 files changed

+81
-133
lines changed

‎tableaudocumentapi/archivefile.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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, ext):
20+
for filename in zip.namelist():
21+
if os.path.splitext(filename)[-1].lower() == ext[:-1]:
22+
return filename
23+
24+
25+
def get_xml_from_archive(filename):
26+
file_type = os.path.splitext(filename)[-1].lower()
27+
with temporary_directory() as temp:
28+
with zipfile.ZipFile(filename) as zf:
29+
zf.extractall(temp)
30+
xml_file = find_file_in_zip(zf, file_type)
31+
xml_tree = ET.parse(os.path.join(temp, xml_file))
32+
33+
return xml_tree
34+
35+
36+
def build_archive_file(archive_contents, zip):
37+
for root_dir, _, files in os.walk(archive_contents):
38+
relative_dir = os.path.relpath(root_dir, archive_contents)
39+
for f in files:
40+
temp_file_full_path = os.path.join(
41+
archive_contents, relative_dir, f)
42+
zipname = os.path.join(relative_dir, f)
43+
zip.write(temp_file_full_path, arcname=zipname)
44+
45+
46+
def save_into_archive(xml_tree, filename, new_filename=None):
47+
# Saving a archive means extracting the contents into a temp folder,
48+
# saving the changes over the twb in that folder, and then
49+
# packaging it back up into a specifically formatted zip with the correct
50+
# relative file paths
51+
52+
if new_filename is None:
53+
new_filename = filename
54+
55+
# Extract to temp directory
56+
with temporary_directory() as temp_path:
57+
file_type = os.path.splitext(filename)[-1].lower()
58+
with zipfile.ZipFile(filename) as zf:
59+
twb_file = find_file_in_zip(zf, file_type)
60+
zf.extractall(temp_path)
61+
# Write the new version of the twb to the temp directory
62+
xml_tree.write(os.path.join(
63+
temp_path, twb_file), encoding="utf-8", xml_declaration=True)
64+
65+
# Write the new archive with the contents of the temp folder
66+
with zipfile.ZipFile(new_filename, "w", compression=zipfile.ZIP_DEFLATED) as new_archive:
67+
build_archive_file(temp_path, new_archive)

‎tableaudocumentapi/datasource.py

Lines changed: 5 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -3,49 +3,11 @@
33
# Datasource - A class for writing datasources to Tableau 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
13-
from tableaudocumentapi import Connection
14-
15-
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)
10+
from tableaudocumentapi import Connection, archivefile
4911

5012

5113
class ConnectionParser(object):
@@ -99,34 +61,11 @@ def from_file(cls, filename):
9961
"Initialize datasource from file (.tds)"
10062

10163
if zipfile.is_zipfile(filename):
102-
dsxml = get_tds_xml_from_tdsx(filename).getroot()
64+
dsxml = archivefile.get_xml_from_archive(filename).getroot()
10365
else:
10466
dsxml = ET.parse(filename).getroot()
10567
return cls(dsxml, filename)
10668

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-
13069
def save(self):
13170
"""
13271
Call finalization code and save file.
@@ -142,7 +81,7 @@ def save(self):
14281
# save the file
14382

14483
if zipfile.is_zipfile(self._filename):
145-
self._save_into_tdsx(self._filename)
84+
archivefile.save_into_archive(self._datasourceTree, self._filename)
14685
else:
14786
self._datasourceTree.write(
14887
self._filename, encoding="utf-8", xml_declaration=True)
@@ -159,7 +98,8 @@ def save_as(self, new_filename):
15998
16099
"""
161100
if zipfile.is_zipfile(self._filename):
162-
self._save_into_tdsx(new_filename)
101+
archivefile.save_into_archive(
102+
self._datasourceTree, self._filename, new_filename)
163103
else:
164104
self._datasourceTree.write(
165105
new_filename, encoding="utf-8", xml_declaration=True)

‎tableaudocumentapi/workbook.py

Lines changed: 7 additions & 65 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, archivefile
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 = archivefile.get_xml_from_archive(
41+
self._filename)
7942
else:
8043
self._workbookTree = ET.parse(self._filename)
8144

@@ -113,7 +76,8 @@ def save(self):
11376
# save the file
11477

11578
if zipfile.is_zipfile(self._filename):
116-
self._save_into_twbx(self._filename)
79+
archivefile.save_into_archive(
80+
self._workbookTree, filename=self._filename)
11781
else:
11882
self._workbookTree.write(
11983
self._filename, encoding="utf-8", xml_declaration=True)
@@ -131,7 +95,8 @@ def save_as(self, new_filename):
13195
"""
13296

13397
if zipfile.is_zipfile(self._filename):
134-
self._save_into_twbx(new_filename)
98+
archivefile.save_into_archive(
99+
self._workbookTree, self._filename, new_filename)
135100
else:
136101
self._workbookTree.write(
137102
new_filename, encoding="utf-8", xml_declaration=True)
@@ -151,29 +116,6 @@ def _prepare_datasources(self, xmlRoot):
151116

152117
return datasources
153118

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-
177119
@staticmethod
178120
def _is_valid_file(filename):
179121
fileExtension = os.path.splitext(filename)[-1].lower()

‎test/bvt.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,15 +173,14 @@ def test_can_open_tdsx_and_save_changes(self):
173173
0].server, 'newdb.test.tsi.lan')
174174

175175
def test_can_open_tdsx_and_save_as_changes(self):
176-
new_tdsx_filename = self.tdsx_file.name + "_TEST_SAVE_AS"
176+
new_tdsx_filename = 'newtdsx.tdsx'
177177
original_wb = Datasource.from_file(self.tdsx_file.name)
178178
original_wb.connections[0].server = 'newdb.test.tsi.lan'
179179
original_wb.save_as(new_tdsx_filename)
180180

181181
new_wb = Datasource.from_file(new_tdsx_filename)
182182
self.assertEqual(new_wb.connections[
183183
0].server, 'newdb.test.tsi.lan')
184-
185184
os.unlink(new_tdsx_filename)
186185

187186

@@ -281,7 +280,7 @@ def test_can_open_twbx_and_save_changes(self):
281280
0].server, 'newdb.test.tsi.lan')
282281

283282
def test_can_open_twbx_and_save_as_changes(self):
284-
new_twbx_filename = self.workbook_file.name + "_TEST_SAVE_AS"
283+
new_twbx_filename = 'newtwbx.twbx'
285284
original_wb = Workbook(self.workbook_file.name)
286285
original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
287286
original_wb.save_as(new_twbx_filename)

0 commit comments

Comments
 (0)
0