Refactor all the archive manipulation logic into one module

Author: t8y8

tableaudocumentapi/archivefile.py

@@ -0,0 +1,67 @@
+import contextlib
+import os
+import shutil
+import tempfile
+import zipfile
+
+import xml.etree.ElementTree as ET
+
+
+@contextlib.contextmanager
+def temporary_directory(*args, **kwargs):
+    d = tempfile.mkdtemp(*args, **kwargs)
+    try:
+        yield d
+    finally:
+        shutil.rmtree(d)
+
+
+def find_file_in_zip(zip, ext):
+    for filename in zip.namelist():
+        if os.path.splitext(filename)[-1].lower() == ext[:-1]:
+            return filename
+
+
+def get_xml_from_archive(filename):
+    file_type = os.path.splitext(filename)[-1].lower()
+    with temporary_directory() as temp:
+        with zipfile.ZipFile(filename) as zf:
+            zf.extractall(temp)
+            xml_file = find_file_in_zip(zf, file_type)
+            xml_tree = ET.parse(os.path.join(temp, xml_file))
+
+    return xml_tree
+
+
+def build_archive_file(archive_contents, zip):
+    for root_dir, _, files in os.walk(archive_contents):
+        relative_dir = os.path.relpath(root_dir, archive_contents)
+        for f in files:
+            temp_file_full_path = os.path.join(
+                archive_contents, relative_dir, f)
+            zipname = os.path.join(relative_dir, f)
+            zip.write(temp_file_full_path, arcname=zipname)
+
+
+def save_into_archive(xml_tree, filename, new_filename=None):
+    # Saving a archive means extracting the contents into a temp folder,
+    # saving the changes over the twb in that folder, and then
+    # packaging it back up into a specifically formatted zip with the correct
+    # relative file paths
+
+    if new_filename is None:
+        new_filename = filename
+
+    # Extract to temp directory
+    with temporary_directory() as temp_path:
+        file_type = os.path.splitext(filename)[-1].lower()
+        with zipfile.ZipFile(filename) as zf:
+            twb_file = find_file_in_zip(zf, file_type)
+            zf.extractall(temp_path)
+        # Write the new version of the twb to the temp directory
+        xml_tree.write(os.path.join(
+            temp_path, twb_file), encoding="utf-8", xml_declaration=True)
+
+        # Write the new archive with the contents of the temp folder
+        with zipfile.ZipFile(new_filename, "w", compression=zipfile.ZIP_DEFLATED) as new_archive:
+            build_archive_file(temp_path, new_archive)

tableaudocumentapi/datasource.py

@@ -3,49 +3,11 @@
 # Datasource - A class for writing datasources to Tableau files
 #
 ###############################################################################
-import contextlib
 import os
-import shutil
-import tempfile
 import zipfile
 
 import xml.etree.ElementTree as ET
-from tableaudocumentapi import Connection
-
-
-@contextlib.contextmanager
-def temporary_directory(*args, **kwargs):
-    d = tempfile.mkdtemp(*args, **kwargs)
-    try:
-        yield d
-    finally:
-        shutil.rmtree(d)
-
-
-def find_tds_in_zip(zip):
-    for filename in zip.namelist():
-        if os.path.splitext(filename)[-1].lower() == '.tds':
-            return filename
-
-
-def get_tds_xml_from_tdsx(filename):
-    with temporary_directory() as temp:
-        with zipfile.ZipFile(filename) as zf:
-            zf.extractall(temp)
-            tds_file = find_tds_in_zip(zf)
-            tds_xml = ET.parse(os.path.join(temp, tds_file))
-
-    return tds_xml
-
-
-def build_tdsx_file(tdsx_contents, zip):
-    for root_dir, _, files in os.walk(tdsx_contents):
-        relative_dir = os.path.relpath(root_dir, tdsx_contents)
-        for f in files:
-            temp_file_full_path = os.path.join(
-                tdsx_contents, relative_dir, f)
-            zipname = os.path.join(relative_dir, f)
-            zip.write(temp_file_full_path, arcname=zipname)
+from tableaudocumentapi import Connection, archivefile
 
 
 class ConnectionParser(object):
@@ -99,34 +61,11 @@ def from_file(cls, filename):
         "Initialize datasource from file (.tds)"
 
         if zipfile.is_zipfile(filename):
-            dsxml = get_tds_xml_from_tdsx(filename).getroot()
+            dsxml = archivefile.get_xml_from_archive(filename).getroot()
         else:
             dsxml = ET.parse(filename).getroot()
         return cls(dsxml, filename)
 
-    def _save_into_tdsx(self, filename=None):
-        # Save reuses existing filename, 'save as' takes a new one
-        if filename is None:
-            filename = self._filename
-
-        # Saving a tdsx means extracting the contents into a temp folder,
-        # saving the changes over the tds in that folder, and then
-        # packaging it back up into a specifically formatted zip with the correct
-        # relative file paths
-
-        # Extract to temp directory
-        with temporary_directory() as temp_path:
-            with zipfile.ZipFile(self._filename) as zf:
-                tds_file = find_tds_in_zip(zf)
-                zf.extractall(temp_path)
-            # Write the new version of the tds to the temp directory
-            self._datasourceTree.write(os.path.join(
-                temp_path, tds_file), encoding="utf-8", xml_declaration=True)
-
-            # Write the new tdsx with the contents of the temp folder
-            with zipfile.ZipFile(filename, "w", compression=zipfile.ZIP_DEFLATED) as new_tdsx:
-                build_tdsx_file(temp_path, new_tdsx)
-
     def save(self):
         """
         Call finalization code and save file.
@@ -142,7 +81,7 @@ def save(self):
         # save the file
 
         if zipfile.is_zipfile(self._filename):
-            self._save_into_tdsx(self._filename)
+            archivefile.save_into_archive(self._datasourceTree, self._filename)
         else:
             self._datasourceTree.write(
                 self._filename, encoding="utf-8", xml_declaration=True)
@@ -159,7 +98,8 @@ def save_as(self, new_filename):
 
         """
         if zipfile.is_zipfile(self._filename):
-            self._save_into_tdsx(new_filename)
+            archivefile.save_into_archive(
+                self._datasourceTree, self._filename, new_filename)
         else:
             self._datasourceTree.write(
                 new_filename, encoding="utf-8", xml_declaration=True)

tableaudocumentapi/workbook.py

@@ -3,15 +3,12 @@
 # Workbook - A class for writing Tableau workbook files
 #
 ###############################################################################
-import contextlib
 import os
-import shutil
-import tempfile
 import zipfile
 
 import xml.etree.ElementTree as ET
 
-from tableaudocumentapi import Datasource
+from tableaudocumentapi import Datasource, archivefile
 
 ###########################################################################
 #
@@ -20,41 +17,6 @@
 ###########################################################################
 
 
-@contextlib.contextmanager
-def temporary_directory(*args, **kwargs):
-    d = tempfile.mkdtemp(*args, **kwargs)
-    try:
-        yield d
-    finally:
-        shutil.rmtree(d)
-
-
-def find_twb_in_zip(zip):
-    for filename in zip.namelist():
-        if os.path.splitext(filename)[-1].lower() == '.twb':
-            return filename
-
-
-def get_twb_xml_from_twbx(filename):
-    with temporary_directory() as temp:
-        with zipfile.ZipFile(filename) as zf:
-            zf.extractall(temp)
-            twb_file = find_twb_in_zip(zf)
-            twb_xml = ET.parse(os.path.join(temp, twb_file))
-
-    return twb_xml
-
-
-def build_twbx_file(twbx_contents, zip):
-    for root_dir, _, files in os.walk(twbx_contents):
-        relative_dir = os.path.relpath(root_dir, twbx_contents)
-        for f in files:
-            temp_file_full_path = os.path.join(
-                twbx_contents, relative_dir, f)
-            zipname = os.path.join(relative_dir, f)
-            zip.write(temp_file_full_path, arcname=zipname)
-
-
 class Workbook(object):
     """
     A class for writing Tableau workbook files.
@@ -75,7 +37,8 @@ def __init__(self, filename):
 
         # Determine if this is a twb or twbx and get the xml root
         if zipfile.is_zipfile(self._filename):
-            self._workbookTree = get_twb_xml_from_twbx(self._filename)
+            self._workbookTree = archivefile.get_xml_from_archive(
+                self._filename)
         else:
             self._workbookTree = ET.parse(self._filename)
 
@@ -113,7 +76,8 @@ def save(self):
         # save the file
 
         if zipfile.is_zipfile(self._filename):
-            self._save_into_twbx(self._filename)
+            archivefile.save_into_archive(
+                self._workbookTree, filename=self._filename)
         else:
             self._workbookTree.write(
                 self._filename, encoding="utf-8", xml_declaration=True)
@@ -131,7 +95,8 @@ def save_as(self, new_filename):
         """
 
         if zipfile.is_zipfile(self._filename):
-            self._save_into_twbx(new_filename)
+            archivefile.save_into_archive(
+                self._workbookTree, self._filename, new_filename)
         else:
             self._workbookTree.write(
                 new_filename, encoding="utf-8", xml_declaration=True)
@@ -151,29 +116,6 @@ def _prepare_datasources(self, xmlRoot):
 
         return datasources
 
-    def _save_into_twbx(self, filename=None):
-        # Save reuses existing filename, 'save as' takes a new one
-        if filename is None:
-            filename = self._filename
-
-        # Saving a twbx means extracting the contents into a temp folder,
-        # saving the changes over the twb in that folder, and then
-        # packaging it back up into a specifically formatted zip with the correct
-        # relative file paths
-
-        # Extract to temp directory
-        with temporary_directory() as temp_path:
-            with zipfile.ZipFile(self._filename) as zf:
-                twb_file = find_twb_in_zip(zf)
-                zf.extractall(temp_path)
-            # Write the new version of the twb to the temp directory
-            self._workbookTree.write(os.path.join(
-                temp_path, twb_file), encoding="utf-8", xml_declaration=True)
-
-            # Write the new twbx with the contents of the temp folder
-            with zipfile.ZipFile(filename, "w", compression=zipfile.ZIP_DEFLATED) as new_twbx:
-                build_twbx_file(temp_path, new_twbx)
-
     @staticmethod
     def _is_valid_file(filename):
         fileExtension = os.path.splitext(filename)[-1].lower()

test/bvt.py

@@ -173,15 +173,14 @@ def test_can_open_tdsx_and_save_changes(self):
                          0].server, 'newdb.test.tsi.lan')
 
     def test_can_open_tdsx_and_save_as_changes(self):
-        new_tdsx_filename = self.tdsx_file.name + "_TEST_SAVE_AS"
+        new_tdsx_filename = 'newtdsx.tdsx'
         original_wb = Datasource.from_file(self.tdsx_file.name)
         original_wb.connections[0].server = 'newdb.test.tsi.lan'
         original_wb.save_as(new_tdsx_filename)
 
         new_wb = Datasource.from_file(new_tdsx_filename)
         self.assertEqual(new_wb.connections[
                          0].server, 'newdb.test.tsi.lan')
-
         os.unlink(new_tdsx_filename)
 
 
@@ -281,7 +280,7 @@ def test_can_open_twbx_and_save_changes(self):
                          0].server, 'newdb.test.tsi.lan')
 
     def test_can_open_twbx_and_save_as_changes(self):
-        new_twbx_filename = self.workbook_file.name + "_TEST_SAVE_AS"
+        new_twbx_filename = 'newtwbx.twbx'
         original_wb = Workbook(self.workbook_file.name)
         original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
         original_wb.save_as(new_twbx_filename)