Merge branch 'zip-support' into development

Author: t8y8

tableaudocumentapi/workbook.py

@@ -3,10 +3,57 @@
 # 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
 
+###########################################################################
+#
+# Utility Functions
+#
+###########################################################################
+
+
+@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):
     """
@@ -24,30 +71,18 @@ def __init__(self, filename):
         Constructor.
 
         """
-        # We have a valid type of input file
-        if self._is_valid_file(filename):
-            # set our filename, open .twb, initialize things
-            self._filename = filename
-            self._workbookTree = ET.parse(filename)
-            self._workbookRoot = self._workbookTree.getroot()
-
-            # prepare our datasource objects
-            self._datasources = self._prepare_datasources(
-                self._workbookRoot)  # self.workbookRoot.find('datasources')
-        else:
-            print('Invalid file type. Must be .twb or .tds.')
-            raise Exception()
-
-    @classmethod
-    def from_file(cls, filename):
-        "Initialize datasource from file (.tds)"
-        if self._is_valid_file(filename):
-            self._filename = filename
-            dsxml = ET.parse(filename).getroot()
-            return cls(dsxml)
+        self._filename = 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)
         else:
-            print('Invalid file type. Must be .twb or .tds.')
-            raise Exception()
+            self._workbookTree = ET.parse(self._filename)
+
+        self._workbookRoot = self._workbookTree.getroot()
+        # prepare our datasource objects
+        self._datasources = self._prepare_datasources(
+            self._workbookRoot)  # self.workbookRoot.find('datasources')
 
     ###########
     # datasources
@@ -76,7 +111,12 @@ def save(self):
         """
 
         # save the file
-        self._workbookTree.write(self._filename, encoding="utf-8", xml_declaration=True)
+
+        if zipfile.is_zipfile(self._filename):
+            self._save_into_twbx(self._filename)
+        else:
+            self._workbookTree.write(
+                self._filename, encoding="utf-8", xml_declaration=True)
 
     def save_as(self, new_filename):
         """
@@ -90,7 +130,11 @@ def save_as(self, new_filename):
 
         """
 
-        self._workbookTree.write(new_filename, encoding="utf-8", xml_declaration=True)
+        if zipfile.is_zipfile(self._filename):
+            self._save_into_twbx(new_filename)
+        else:
+            self._workbookTree.write(
+                new_filename, encoding="utf-8", xml_declaration=True)
 
     ###########################################################################
     #
@@ -107,6 +151,29 @@ 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()