Initial stab at TWBX support

.gitignore

@@ -63,3 +63,4 @@ target/
 
 #Other things
 .DS_Store
+.idea

.travis.yml

@@ -12,9 +12,9 @@ install:
 # command to run tests
 script:
   # Tests
-  - python test.py
+  - python setup.py test
   # pep8
-  - pep8 --ignore=E501 .
+  - pep8 .
   # Examples
   - (cd "Examples/Replicate Workbook" && python replicateWorkbook.py)
   - (cd "Examples/List TDS Info" && python listTDSInfo.py)

setup.cfg

@@ -0,0 +1,10 @@
+[wheel]
+universal = 1
+
+[pycodestyle]
+select = 
+max_line_length = 120
+
+[pep8]
+max_line_length = 120
+

setup.py

@@ -5,12 +5,12 @@
 
 setup(
     name='tableaudocumentapi',
-    version='0.0.1',
-    summary='A Python module for working with Tableau files.',
+    version='0.1.0-dev',
     author='Tableau Software',
     author_email='[email protected]',
     url='https://github.com/tableau/document-api-python',
     py_modules=['tableaudocumentapi'],
     license='MIT',
-    description='A Python module for working with Tableau files.'
+    description='A Python module for working with Tableau files.',
+    test_suite='test'
 )

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()

test/assets/CONNECTION.xml

@@ -0,0 +1 @@
+<connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection>

test/assets/TABLEAU_10_TDS.tds

@@ -0,0 +1 @@
+<?xml version='1.0' encoding='utf-8' ?><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource>

test/assets/TABLEAU_10_TWB.twb

@@ -0,0 +1 @@
+<?xml version='1.0' encoding='utf-8' ?><workbook source-build='0.0.0 (0000.16.0510.1300)' source-platform='mac' version='10.0' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource></datasources></workbook>

test/assets/TABLEAU_93_TDS.tds

@@ -0,0 +1 @@
+<?xml version='1.0' encoding='utf-8' ?><datasource formatted-name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' inline='true' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource>

test/assets/TABLEAU_93_TWB.twb

@@ -0,0 +1 @@
+<?xml version='1.0' encoding='utf-8' ?><workbook source-build='9.3.1 (9300.16.0510.0100)' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy (TestV1)' inline='true' name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' version='9.3'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource></datasources></workbook>

test.py → test/bvt.py

@@ -1,21 +1,22 @@
-import unittest
-import io
 import os
+import unittest
+
 import xml.etree.ElementTree as ET
 
 from tableaudocumentapi import Workbook, Datasource, Connection, ConnectionParser
 
 
-TABLEAU_93_WORKBOOK = '''<?xml version='1.0' encoding='utf-8' ?><workbook source-build='9.3.1 (9300.16.0510.0100)' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy (TestV1)' inline='true' name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' version='9.3'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource></datasources></workbook>'''
+TABLEAU_93_TWB = 'test/assets/TABLEAU_93_TWB.twb'
+
+TABLEAU_93_TDS = 'test/assets/TABLEAU_93_TDS.tds'
 
-TABLEAU_93_TDS = '''<?xml version='1.0' encoding='utf-8' ?><datasource formatted-name='sqlserver.17u3bqc16tjtxn14e2hxh19tyvpo' inline='true' source-platform='mac' version='9.3' xmlns:user='http://www.tableausoftware.com/xml/user'><connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection></datasource>'''
+TABLEAU_10_TDS = 'test/assets/TABLEAU_10_TDS.tds'
 
-TABLEAU_10_TDS = '''<?xml version='1.0' encoding='utf-8' ?><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource>'''
+TABLEAU_10_TWB = 'test/assets/TABLEAU_10_TWB.twb'
 
-TABLEAU_10_WORKBOOK = '''<?xml version='1.0' encoding='utf-8' ?><workbook source-build='0.0.0 (0000.16.0510.1300)' source-platform='mac' version='10.0' xmlns:user='http://www.tableausoftware.com/xml/user'><datasources><datasource caption='xy+ (Multiple Connections)' inline='true' name='federated.1s4nxn20cywkdv13ql0yk0g1mpdx' version='10.0'><connection class='federated'><named-connections><named-connection caption='mysql55.test.tsi.lan' name='mysql.1ewmkrw0mtgsev1dnurma1blii4x'><connection class='mysql' dbname='testv1' odbc-native-protocol='yes' port='3306' server='mysql55.test.tsi.lan' source-charset='' username='test' /></named-connection><named-connection caption='mssql2012.test.tsi.lan' name='sqlserver.1erdwp01uqynlb14ul78p0haai2r'><connection authentication='sqlserver' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username='test' /></named-connection></named-connections></connection></datasource></datasources></workbook>'''
+TABLEAU_CONNECTION_XML = ET.parse('test/assets/CONNECTION.xml').getroot()
 
-TABLEAU_CONNECTION_XML = ET.fromstring(
-    '''<connection authentication='sspi' class='sqlserver' dbname='TestV1' odbc-native-protocol='yes' one-time-sql='' server='mssql2012.test.tsi.lan' username=''></connection>''')
+TABLEAU_10_TWBX = 'test/assets/TABLEAU_10_TWBX.twbx'
 
 
 class HelperMethodTests(unittest.TestCase):
@@ -34,14 +35,14 @@ def test_is_valid_file_with_invalid_inputs(self):
 class ConnectionParserTests(unittest.TestCase):
 
     def test_can_extract_legacy_connection(self):
-        parser = ConnectionParser(ET.fromstring(TABLEAU_93_TDS), '9.2')
+        parser = ConnectionParser(ET.parse(TABLEAU_93_TDS), '9.2')
         connections = parser.get_connections()
         self.assertIsInstance(connections, list)
         self.assertIsInstance(connections[0], Connection)
         self.assertEqual(connections[0].dbname, 'TestV1')
 
     def test_can_extract_federated_connections(self):
-        parser = ConnectionParser(ET.fromstring(TABLEAU_10_TDS), '10.0')
+        parser = ConnectionParser(ET.parse(TABLEAU_10_TDS), '10.0')
         connections = parser.get_connections()
         self.assertIsInstance(connections, list)
         self.assertIsInstance(connections[0], Connection)
@@ -74,9 +75,9 @@ def test_can_write_attributes_to_connection(self):
 class DatasourceModelTests(unittest.TestCase):
 
     def setUp(self):
-        self.tds_file = io.FileIO('test.tds', 'w')
-        self.tds_file.write(TABLEAU_93_TDS.encode('utf8'))
-        self.tds_file.seek(0)
+        with open(TABLEAU_93_TDS, 'rb') as in_file, open('test.tds', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.tds_file = out_file
 
     def tearDown(self):
         self.tds_file.close()
@@ -115,9 +116,9 @@ def test_save_has_xml_declaration(self):
 class DatasourceModelV10Tests(unittest.TestCase):
 
     def setUp(self):
-        self.tds_file = io.FileIO('test10.tds', 'w')
-        self.tds_file.write(TABLEAU_10_TDS.encode('utf8'))
-        self.tds_file.seek(0)
+        with open(TABLEAU_10_TDS, 'rb') as in_file, open('test.twb', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.tds_file = out_file
 
     def tearDown(self):
         self.tds_file.close()
@@ -145,9 +146,9 @@ def test_can_save_tds(self):
 class WorkbookModelTests(unittest.TestCase):
 
     def setUp(self):
-        self.workbook_file = io.FileIO('test.twb', 'w')
-        self.workbook_file.write(TABLEAU_93_WORKBOOK.encode('utf8'))
-        self.workbook_file.seek(0)
+        with open(TABLEAU_93_TWB, 'rb') as in_file, open('test.twb', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.workbook_file = out_file
 
     def tearDown(self):
         self.workbook_file.close()
@@ -173,9 +174,9 @@ def test_can_update_datasource_connection_and_save(self):
 class WorkbookModelV10Tests(unittest.TestCase):
 
     def setUp(self):
-        self.workbook_file = io.FileIO('testv10.twb', 'w')
-        self.workbook_file.write(TABLEAU_10_WORKBOOK.encode('utf8'))
-        self.workbook_file.seek(0)
+        with open(TABLEAU_10_TWB, 'rb') as in_file, open('test.twb', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.workbook_file = out_file
 
     def tearDown(self):
         self.workbook_file.close()
@@ -211,5 +212,43 @@ def test_save_has_xml_declaration(self):
             self.assertEqual(
                 first_line, "<?xml version='1.0' encoding='utf-8'?>")
 
+
+class WorkbookModelV10TWBXTests(unittest.TestCase):
+
+    def setUp(self):
+        with open(TABLEAU_10_TWBX, 'rb') as in_file, open('test.twbx', 'wb') as out_file:
+            out_file.write(in_file.read())
+            self.workbook_file = out_file
+
+    def tearDown(self):
+        self.workbook_file.close()
+        os.unlink(self.workbook_file.name)
+
+    def test_can_open_twbx(self):
+        wb = Workbook(self.workbook_file.name)
+        self.assertTrue(wb.datasources)
+        self.assertTrue(wb.datasources[0].connections)
+
+    def test_can_open_twbx_and_save_changes(self):
+        original_wb = Workbook(self.workbook_file.name)
+        original_wb.datasources[0].connections[0].server = 'newdb.test.tsi.lan'
+        original_wb.save()
+
+        new_wb = Workbook(self.workbook_file.name)
+        self.assertEqual(new_wb.datasources[0].connections[
+                         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"
+        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)
+
+        new_wb = Workbook(new_twbx_filename)
+        self.assertEqual(new_wb.datasources[0].connections[
+                         0].server, 'newdb.test.tsi.lan')
+
+        os.unlink(new_twbx_filename)
+
 if __name__ == '__main__':
     unittest.main()