tableau · t8y8 · Oct 31, 2019 · Oct 29, 2019 · Oct 29, 2019
diff --git a/setup.py b/setup.py
@@ -30,6 +30,7 @@
     ],
     tests_require=[
         'requests-mock>=1.0,<2.0',
-        'pytest'
+        'pytest',
+        'mock'
     ]
 )
diff --git a/tableauserverclient/filesys_helpers.py b/tableauserverclient/filesys_helpers.py
@@ -1,6 +1,22 @@
+import os
 ALLOWED_SPECIAL = (' ', '.', '_', '-')
 
 
 def to_filename(string_to_sanitize):
     sanitized = (c for c in string_to_sanitize if c.isalnum() or c in ALLOWED_SPECIAL)
     return "".join(sanitized)
+
+
+def make_download_path(filepath, filename):
+    download_path = None
+
+    if filepath is None:
+        download_path = filename
+
+    elif os.path.isdir(filepath):
+        download_path = os.path.join(filepath, filename)
+
+    else:
+        download_path = filepath + os.path.splitext(filename)[1]
+
+    return download_path
diff --git a/tableauserverclient/server/endpoint/datasources_endpoint.py b/tableauserverclient/server/endpoint/datasources_endpoint.py
@@ -6,7 +6,7 @@
 from .fileuploads_endpoint import Fileuploads
 from .resource_tagger import _ResourceTagger
 from .. import RequestFactory, DatasourceItem, PaginationItem, ConnectionItem
-from ...filesys_helpers import to_filename
+from ...filesys_helpers import to_filename, make_download_path
 from ...models.tag_item import TagItem
 from ...models.job_item import JobItem
 import os
@@ -104,17 +104,15 @@ def download(self, datasource_id, filepath=None, include_extract=True, no_extrac
         with closing(self.get_request(url, parameters={'stream': True})) as server_response:
             _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
             filename = to_filename(os.path.basename(params['filename']))
-            if filepath is None:
-                filepath = filename
-            elif os.path.isdir(filepath):
-                filepath = os.path.join(filepath, filename)
 
-            with open(filepath, 'wb') as f:
+            download_path = make_download_path(filepath, filename)
+
+            with open(download_path, 'wb') as f:
                 for chunk in server_response.iter_content(1024):  # 1KB
                     f.write(chunk)
 
-        logger.info('Downloaded datasource to {0} (ID: {1})'.format(filepath, datasource_id))
-        return os.path.abspath(filepath)
+        logger.info('Downloaded datasource to {0} (ID: {1})'.format(download_path, datasource_id))
+        return os.path.abspath(download_path)
 
     # Update datasource
     @api(version="2.0")

diff --git a/tableauserverclient/server/endpoint/flows_endpoint.py b/tableauserverclient/server/endpoint/flows_endpoint.py
@@ -6,7 +6,7 @@
 from .fileuploads_endpoint import Fileuploads
 from .resource_tagger import _ResourceTagger
 from .. import RequestFactory, FlowItem, PaginationItem, ConnectionItem
-from ...filesys_helpers import to_filename
+from ...filesys_helpers import to_filename, make_download_path
 from ...models.tag_item import TagItem
 from ...models.job_item import JobItem
 import os
@@ -94,17 +94,15 @@ def download(self, flow_id, filepath=None):
         with closing(self.get_request(url, parameters={'stream': True})) as server_response:
             _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
             filename = to_filename(os.path.basename(params['filename']))
-            if filepath is None:
-                filepath = filename
-            elif os.path.isdir(filepath):
-                filepath = os.path.join(filepath, filename)
 
-            with open(filepath, 'wb') as f:
+            download_path = make_download_path(filepath, filename)
+
+            with open(download_path, 'wb') as f:
                 for chunk in server_response.iter_content(1024):  # 1KB
                     f.write(chunk)
 
-        logger.info('Downloaded flow to {0} (ID: {1})'.format(filepath, flow_id))
-        return os.path.abspath(filepath)
+        logger.info('Downloaded flow to {0} (ID: {1})'.format(download_path, flow_id))
+        return os.path.abspath(download_path)
 
     # Update flow
     @api(version="3.3")

diff --git a/tableauserverclient/server/endpoint/workbooks_endpoint.py b/tableauserverclient/server/endpoint/workbooks_endpoint.py
@@ -8,7 +8,7 @@
 from .. import RequestFactory, WorkbookItem, ConnectionItem, ViewItem, PaginationItem
 from ...models.tag_item import TagItem
 from ...models.job_item import JobItem
-from ...filesys_helpers import to_filename
+from ...filesys_helpers import to_filename, make_download_path
 
 import os
 import logging
@@ -129,16 +129,14 @@ def download(self, workbook_id, filepath=None, include_extract=True, no_extract=
         with closing(self.get_request(url, parameters={"stream": True})) as server_response:
             _, params = cgi.parse_header(server_response.headers['Content-Disposition'])
             filename = to_filename(os.path.basename(params['filename']))
-            if filepath is None:
-                filepath = filename
-            elif os.path.isdir(filepath):
-                filepath = os.path.join(filepath, filename)
 
-            with open(filepath, 'wb') as f:
+            download_path = make_download_path(filepath, filename)
+
+            with open(download_path, 'wb') as f:
                 for chunk in server_response.iter_content(1024):  # 1KB
                     f.write(chunk)
-        logger.info('Downloaded workbook to {0} (ID: {1})'.format(filepath, workbook_id))
-        return os.path.abspath(filepath)
+        logger.info('Downloaded workbook to {0} (ID: {1})'.format(download_path, workbook_id))
+        return os.path.abspath(download_path)
 
     # Get all views of workbook
     @api(version="2.0")

diff --git a/test/test_regression_tests.py b/test/test_regression_tests.py
@@ -1,6 +1,13 @@
 import unittest
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
 import tableauserverclient.server.request_factory as factory
 from tableauserverclient.server.endpoint import Endpoint
+from tableauserverclient.filesys_helpers import to_filename, make_download_path
 
 
 class BugFix257(unittest.TestCase):
@@ -21,3 +28,36 @@ class FakeResponse(object):
         server_response = FakeResponse()
 
         self.assertEqual(Endpoint._safe_to_log(server_response), '[Truncated File Contents]')
+
+
+class FileSysHelpers(unittest.TestCase):
+    def test_to_filename(self):
+        invalid = [
+            "23brhafbjrjhkbbea.txt",
+            'a_b_C.txt',
+            'windows space.txt',
+            'abc#def.txt',
+            't@bL3A()',
+        ]
+
+        valid = [
+            "23brhafbjrjhkbbea.txt",
+            'a_b_C.txt',
+            'windows space.txt',
+            'abcdef.txt',
+            'tbL3A',
+        ]
+
+        self.assertTrue(all([(to_filename(i) == v) for i, v in zip(invalid, valid)]))
+
+    def test_make_download_path(self):
+        no_file_path = (None, 'file.ext')
+        has_file_path_folder = ('/root/folder/', 'file.ext')
+        has_file_path_file = ('out', 'file.ext')
+
+        self.assertEquals('file.ext', make_download_path(*no_file_path))
+        self.assertEquals('out.ext', make_download_path(*has_file_path_file))
+
+        with mock.patch('os.path.isdir') as mocked_isdir:
+            mocked_isdir.return_value = True
+            self.assertEquals('/root/folder/file.ext', make_download_path(*has_file_path_folder))