Make sure artifact temporary files names are valid file names

allegroai · Mar 12, 2020 · 5b29aa1 · 5b29aa1
1 parent 84a3442
commit 5b29aa1
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/trains/binding/artifacts.py b/trains/binding/artifacts.py
@@ -2,6 +2,7 @@
 import json
 import mimetypes
 import os
+from six.moves.urllib.parse import quote
 from copy import deepcopy
 from datetime import datetime
 from multiprocessing import RLock, Event
@@ -299,7 +300,7 @@ def upload_artifact(self, name, artifact_object=None, metadata=None, delete_afte
             artifact_type_data.preview = str(artifact_object.__repr__())
             override_filename_ext_in_uri = '.npz'
             override_filename_in_uri = name+override_filename_ext_in_uri
-            fd, local_filename = mkstemp(prefix=name+'.', suffix=override_filename_ext_in_uri)
+            fd, local_filename = mkstemp(prefix=quote(name, safe="")+'.', suffix=override_filename_ext_in_uri)
             os.close(fd)
             np.savez_compressed(local_filename, **{name: artifact_object})
             delete_after_upload = True
@@ -309,7 +310,7 @@ def upload_artifact(self, name, artifact_object=None, metadata=None, delete_afte
             artifact_type_data.preview = str(artifact_object.__repr__())
             override_filename_ext_in_uri = self._save_format
             override_filename_in_uri = name
-            fd, local_filename = mkstemp(prefix=name+'.', suffix=override_filename_ext_in_uri)
+            fd, local_filename = mkstemp(prefix=quote(name, safe="")+'.', suffix=override_filename_ext_in_uri)
             os.close(fd)
             artifact_object.to_csv(local_filename, compression=self._compression)
             delete_after_upload = True
@@ -320,7 +321,7 @@ def upload_artifact(self, name, artifact_object=None, metadata=None, delete_afte
             artifact_type_data.preview = desc[1:desc.find(' at ')]
             override_filename_ext_in_uri = '.png'
             override_filename_in_uri = name + override_filename_ext_in_uri
-            fd, local_filename = mkstemp(prefix=name+'.', suffix=override_filename_ext_in_uri)
+            fd, local_filename = mkstemp(prefix=quote(name, safe="")+'.', suffix=override_filename_ext_in_uri)
             os.close(fd)
             artifact_object.save(local_filename)
             delete_after_upload = True
@@ -330,7 +331,7 @@ def upload_artifact(self, name, artifact_object=None, metadata=None, delete_afte
             preview = json.dumps(artifact_object, sort_keys=True, indent=4)
             override_filename_ext_in_uri = '.json'
             override_filename_in_uri = name + override_filename_ext_in_uri
-            fd, local_filename = mkstemp(prefix=name+'.', suffix=override_filename_ext_in_uri)
+            fd, local_filename = mkstemp(prefix=quote(name, safe="")+'.', suffix=override_filename_ext_in_uri)
             os.write(fd, bytes(preview.encode()))
             os.close(fd)
             artifact_type_data.preview = preview
@@ -365,7 +366,7 @@ def upload_artifact(self, name, artifact_object=None, metadata=None, delete_afte
                 files = list(Path(folder).rglob(wildcard))
                 override_filename_ext_in_uri = '.zip'
                 override_filename_in_uri = folder.parts[-1] + override_filename_ext_in_uri
-                fd, zip_file = mkstemp(prefix=folder.parts[-1]+'.', suffix=override_filename_ext_in_uri)
+                fd, zip_file = mkstemp(prefix=quote(folder.parts[-1], safe="")+'.', suffix=override_filename_ext_in_uri)
                 try:
                     artifact_type_data.content_type = 'application/zip'
                     artifact_type_data.preview = 'Archive content {}:\n'.format(artifact_object.as_posix())
@@ -510,7 +511,7 @@ def _upload_data_audit_artifacts(self, name):
 
         override_filename_ext_in_uri = self._save_format
         override_filename_in_uri = name
-        fd, local_csv = mkstemp(prefix=name + '.', suffix=override_filename_ext_in_uri)
+        fd, local_csv = mkstemp(prefix=quote(name, safe="") + '.', suffix=override_filename_ext_in_uri)
         os.close(fd)
         local_csv = Path(local_csv)
         pd_artifact.to_csv(local_csv.as_posix(), index=False, compression=self._compression)