diff --git a/tableaudocumentapi/datasource.py b/tableaudocumentapi/datasource.py index 3f64145..09860b4 100644 --- a/tableaudocumentapi/datasource.py +++ b/tableaudocumentapi/datasource.py @@ -113,7 +113,7 @@ def __init__(self, dsxml, filename=None): def from_file(cls, filename): """Initialize datasource from file (.tds)""" - dsxml = xml_open(filename).getroot() + dsxml = xml_open(filename, cls.__name__.lower()).getroot() return cls(dsxml, filename) def save(self): diff --git a/tableaudocumentapi/workbook.py b/tableaudocumentapi/workbook.py index 28ddd03..1359356 100644 --- a/tableaudocumentapi/workbook.py +++ b/tableaudocumentapi/workbook.py @@ -32,7 +32,7 @@ def __init__(self, filename): self._filename = filename - self._workbookTree = xml_open(self._filename) + self._workbookTree = xml_open(self._filename, self.__class__.__name__.lower()) self._workbookRoot = self._workbookTree.getroot() # prepare our datasource objects diff --git a/tableaudocumentapi/xfile.py b/tableaudocumentapi/xfile.py index a0cd62e..8f9ffd1 100644 --- a/tableaudocumentapi/xfile.py +++ b/tableaudocumentapi/xfile.py @@ -17,15 +17,28 @@ class TableauVersionNotSupportedException(Exception): pass -def xml_open(filename): - # Determine if this is a twb or twbx and get the xml root +class TableauInvalidFileException(Exception): + pass + + +def xml_open(filename, expected_root=None): + if zipfile.is_zipfile(filename): tree = get_xml_from_archive(filename) else: tree = ET.parse(filename) - file_version = Version(tree.getroot().attrib.get('version', '0.0')) + + tree_root = tree.getroot() + + file_version = Version(tree_root.attrib.get('version', '0.0')) + if file_version < MIN_SUPPORTED_VERSION: raise TableauVersionNotSupportedException(file_version) + + if expected_root and (expected_root != tree_root.tag): + raise TableauInvalidFileException( + "'{}'' is not a valid '{}' file".format(filename, expected_root)) + return tree @@ -40,14 +53,13 @@ def temporary_directory(*args, **kwargs): def find_file_in_zip(zip_file): for filename in zip_file.namelist(): - try: - with zip_file.open(filename) as xml_candidate: - ET.parse(xml_candidate).getroot().tag in ( - 'workbook', 'datasource') + with zip_file.open(filename) as xml_candidate: + try: + ET.parse(xml_candidate) return filename - except ET.ParseError: - # That's not an XML file by gosh - pass + except ET.ParseError: + # That's not an XML file by gosh + pass def get_xml_from_archive(filename): diff --git a/test/assets/TABLEAU_82_TWB.twb b/test/assets/TABLEAU_82_TWB.twb new file mode 100644 index 0000000..2ab236d --- /dev/null +++ b/test/assets/TABLEAU_82_TWB.twb @@ -0,0 +1 @@ + diff --git a/test/bvt.py b/test/bvt.py index 6a7cdf8..27d3703 100644 --- a/test/bvt.py +++ b/test/bvt.py @@ -4,9 +4,12 @@ import xml.etree.ElementTree as ET from tableaudocumentapi import Workbook, Datasource, Connection, ConnectionParser +from tableaudocumentapi.xfile import TableauInvalidFileException, TableauVersionNotSupportedException TEST_DIR = os.path.dirname(__file__) +TABLEAU_82_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_82_TWB.twb') + TABLEAU_93_TWB = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TWB.twb') TABLEAU_93_TDS = os.path.join(TEST_DIR, 'assets', 'TABLEAU_93_TDS.tds') @@ -282,10 +285,28 @@ def test_can_open_twbx_and_save_as_changes(self): class EmptyWorkbookWillLoad(unittest.TestCase): + def test_no_exceptions_thrown(self): wb = Workbook(EMPTY_WORKBOOK) self.assertIsNotNone(wb) +class LoadOnlyValidFileTypes(unittest.TestCase): + + def test_exception_when_workbook_given_tdsx(self): + with self.assertRaises(TableauInvalidFileException): + wb = Workbook(TABLEAU_10_TDSX) + + def test_exception_when_datasource_given_twbx(self): + with self.assertRaises(TableauInvalidFileException): + ds = Datasource.from_file(TABLEAU_10_TWBX) + + +class SupportedWorkbookVersions(unittest.TestCase): + + def test_82_workbook_throws_exception(self): + with self.assertRaises(TableauVersionNotSupportedException): + wb = Workbook(TABLEAU_82_TWB) + if __name__ == '__main__': unittest.main()