Initial TDSX support. Lots of duplicated code with TWBX's

Author: t8y8

tableaudocumentapi/datasource.py

@@ -3,10 +3,51 @@
 # 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)
+
+
 class ConnectionParser(object):
 
     def __init__(self, datasource_xml, version):
@@ -56,9 +97,36 @@ def __init__(self, dsxml, filename=None):
     @classmethod
     def from_file(cls, filename):
         "Initialize datasource from file (.tds)"
-        dsxml = ET.parse(filename).getroot()
+
+        if zipfile.is_zipfile(filename):
+            dsxml = get_tds_xml_from_tdsx(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.
@@ -72,7 +140,12 @@ def save(self):
         """
 
         # save the file
-        self._datasourceTree.write(self._filename, encoding="utf-8", xml_declaration=True)
+
+        if zipfile.is_zipfile(self._filename):
+            self._save_into_tdsx(self._filename)
+        else:
+            self._datasourceTree.write(
+                self._filename, encoding="utf-8", xml_declaration=True)
 
     def save_as(self, new_filename):
         """
@@ -85,7 +158,11 @@ def save_as(self, new_filename):
             Nothing.
 
         """
-        self._datasourceTree.write(new_filename, encoding="utf-8", xml_declaration=True)
+        if zipfile.is_zipfile(self._filename):
+            self._save_into_tdsx(new_filename)
+        else:
+            self._datasourceTree.write(
+                new_filename, encoding="utf-8", xml_declaration=True)
 
     ###########
     # name

test/assets/TABLEAU_10_TDSX.tdsx

@@ -15,10 +15,13 @@
 
 TABLEAU_10_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TWB.twb')
 
-TABLEAU_CONNECTION_XML = ET.parse(os.path.join(TEST_DIR, 'assets', 'CONNECTION.xml')).getroot()
+TABLEAU_CONNECTION_XML = ET.parse(os.path.join(
+    TEST_DIR, 'assets', 'CONNECTION.xml')).getroot()
 
 TABLEAU_10_TWBX = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TWBX.twbx')
 
+TABLEAU_10_TDSX = os.path.join(TEST_DIR, 'assets', 'TABLEAU_10_TDSX.tdsx')
+
 
 class HelperMethodTests(unittest.TestCase):
 
@@ -144,6 +147,44 @@ def test_can_save_tds(self):
         self.assertEqual(new_tds.connections[0].dbname, 'newdb.test.tsi.lan')
 
 
+class DatasourceModelV10TDSXTests(unittest.TestCase):
+
+    def setUp(self):
+        with open(TABLEAU_10_TDSX, 'rb') as in_file, open('test.tdsx', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.tdsx_file = out_file
+
+    def tearDown(self):
+        self.tdsx_file.close()
+        os.unlink(self.tdsx_file.name)
+
+    def test_can_open_tdsx(self):
+        ds = Datasource.from_file(self.tdsx_file.name)
+        self.assertTrue(ds.connections)
+        self.assertTrue(ds.name)
+
+    def test_can_open_tdsx_and_save_changes(self):
+        original_tdsx = Datasource.from_file(self.tdsx_file.name)
+        original_tdsx.connections[0].server = 'newdb.test.tsi.lan'
+        original_tdsx.save()
+
+        new_tdsx = Datasource.from_file(self.tdsx_file.name)
+        self.assertEqual(new_tdsx.connections[
+                         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"
+        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)
+
+
 class WorkbookModelTests(unittest.TestCase):
 
     def setUp(self):