From 93a47fec2b048874cff669dd77c5cd0ccb08f9d4 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Wed, 15 May 2019 07:19:56 +0200 Subject: [PATCH] Use binary mode for local copy list in `upload_calculation` (#2748) The file content of the files in the node repository is not necessarily encoded but could be binary, so the best thing is to not guess the encoding but transfer the raw bytes. --- aiida/backends/tests/orm/data/test_singlefile.py | 6 +++--- aiida/common/folders.py | 4 ++++ aiida/engine/daemon/execmanager.py | 6 ++++-- aiida/orm/nodes/data/singlefile.py | 2 +- aiida/orm/nodes/node.py | 4 ++-- aiida/orm/utils/repository.py | 5 +++-- 6 files changed, 17 insertions(+), 10 deletions(-) diff --git a/aiida/backends/tests/orm/data/test_singlefile.py b/aiida/backends/tests/orm/data/test_singlefile.py index 5d1e2ec24e..269ad33422 100644 --- a/aiida/backends/tests/orm/data/test_singlefile.py +++ b/aiida/backends/tests/orm/data/test_singlefile.py @@ -66,9 +66,9 @@ def test_construct_from_filelike(self): """Test constructing an instance from filelike instead of filepath.""" content_original = u'some testing text\nwith a newline' - with tempfile.NamedTemporaryFile(mode='w+') as handle: + with tempfile.NamedTemporaryFile(mode='wb+') as handle: basename = os.path.basename(handle.name) - handle.write(content_original) + handle.write(content_original.encode('utf-8')) handle.flush() handle.seek(0) node = SinglefileData(file=handle) @@ -91,7 +91,7 @@ def test_construct_from_string(self): """Test constructing an instance from a string.""" content_original = u'some testing text\nwith a newline' - with io.StringIO(content_original) as handle: + with io.BytesIO(content_original.encode('utf-8')) as handle: node = SinglefileData(file=handle) with node.open() as handle: diff --git a/aiida/common/folders.py b/aiida/common/folders.py index 4f9268afa6..e72a6a7759 100644 --- a/aiida/common/folders.py +++ b/aiida/common/folders.py @@ -230,6 +230,9 @@ def create_file_from_filelike(self, filelike, filename, mode='wb', encoding=None filename = six.text_type(filename) filepath = self.get_abs_path(filename) + if 'b' in mode: + encoding = None + with io.open(filepath, mode=mode, encoding=encoding) as handle: # In python 2 a string literal can either be of unicode or string (bytes) type. Since we do not know what @@ -243,6 +246,7 @@ def create_file_from_filelike(self, filelike, filename, mode='wb', encoding=None try: shutil.copyfileobj(utf8reader(filelike), handle) except (UnicodeDecodeError, UnicodeEncodeError): + filelike.seek(0) shutil.copyfileobj(filelike, handle) else: shutil.copyfileobj(filelike, handle) diff --git a/aiida/engine/daemon/execmanager.py b/aiida/engine/daemon/execmanager.py index 67406b6473..fb21ac743a 100644 --- a/aiida/engine/daemon/execmanager.py +++ b/aiida/engine/daemon/execmanager.py @@ -161,9 +161,11 @@ def upload_calculation(node, transport, calc_info, script_filename, dry_run=Fals # Note, once #2579 is implemented, use the `node.open` method instead of the named temporary file in # combination with the new `Transport.put_object_from_filelike` - with NamedTemporaryFile(mode='w+') as handle: - handle.write(data_node.get_object_content(filename)) + # Since the content of the node could potentially be binary, we read the raw bytes and pass them on + with NamedTemporaryFile(mode='wb+') as handle: + handle.write(data_node.get_object_content(filename, mode='rb')) handle.flush() + handle.seek(0) transport.put(handle.name, target) if dry_run: diff --git a/aiida/orm/nodes/data/singlefile.py b/aiida/orm/nodes/data/singlefile.py index d73a8d2a12..cbcb0e6c23 100644 --- a/aiida/orm/nodes/data/singlefile.py +++ b/aiida/orm/nodes/data/singlefile.py @@ -97,7 +97,7 @@ def set_file(self, file): pass if is_filelike: - self.put_object_from_filelike(file, key) + self.put_object_from_filelike(file, key, mode='wb') else: self.put_object_from_file(file, key) diff --git a/aiida/orm/nodes/node.py b/aiida/orm/nodes/node.py index 251b865aaf..e5bf9e7c6d 100644 --- a/aiida/orm/nodes/node.py +++ b/aiida/orm/nodes/node.py @@ -682,12 +682,12 @@ def get_object(self, key): """ return self._repository.get_object(key) - def get_object_content(self, key): + def get_object_content(self, key, mode='r'): """Return the content of a object identified by key. :param key: fully qualified identifier for the object within the repository """ - return self._repository.get_object_content(key) + return self._repository.get_object_content(key, mode) def put_object_from_tree(self, path, key=None, contents_only=True, force=False): """Store a new object under `key` with the contents of the directory located at `path` on this file system. diff --git a/aiida/orm/utils/repository.py b/aiida/orm/utils/repository.py index eb61f4dff1..84044d5cf6 100644 --- a/aiida/orm/utils/repository.py +++ b/aiida/orm/utils/repository.py @@ -126,12 +126,13 @@ def get_object(self, key): return File(filename, FileType.FILE) - def get_object_content(self, key): + def get_object_content(self, key, mode='r'): """Return the content of a object identified by key. :param key: fully qualified identifier for the object within the repository + :param mode: the mode under which to open the handle """ - with self.open(key) as handle: + with self.open(key, mode=mode) as handle: return handle.read() def put_object_from_tree(self, path, key=None, contents_only=True, force=False):