Merge pull request #3614 from Flamefire/downloadFileChunked

Add option to write file from file-like object and use in download_file

Author: boegel

easybuild/tools/filetools.py

@@ -217,7 +217,7 @@ def write_file(path, data, append=False, forced=False, backup=False, always_over
     overwrites current file contents without backup by default!
 
     :param path: location of file
-    :param data: contents to write to file
+    :param data: contents to write to file. Can be a file-like object of binary data
     :param append: append to existing file rather than overwrite
     :param forced: force actually writing file in (extended) dry run mode
     :param backup: back up existing file before overwriting or modifying it
@@ -246,15 +246,21 @@ def write_file(path, data, append=False, forced=False, backup=False, always_over
     # cfr. https://docs.python.org/3/library/functions.html#open
     mode = 'a' if append else 'w'
 
+    data_is_file_obj = hasattr(data, 'read')
+
     # special care must be taken with binary data in Python 3
-    if sys.version_info[0] >= 3 and isinstance(data, bytes):
+    if sys.version_info[0] >= 3 and (isinstance(data, bytes) or data_is_file_obj):
         mode += 'b'
 
     # note: we can't use try-except-finally, because Python 2.4 doesn't support it as a single block
     try:
         mkdir(os.path.dirname(path), parents=True)
         with open_file(path, mode) as fh:
-            fh.write(data)
+            if data_is_file_obj:
+                # if a file-like object was provided, use copyfileobj (which reads the file in chunks)
+                shutil.copyfileobj(data, fh)
+            else:
+                fh.write(data)
     except IOError as err:
         raise EasyBuildError("Failed to write to %s: %s", path, err)
 
@@ -710,7 +716,11 @@ def download_file(filename, url, path, forced=False):
                 url_fd = response.raw
                 url_fd.decode_content = True
             _log.debug('response code for given url %s: %s' % (url, status_code))
-            write_file(path, url_fd.read(), forced=forced, backup=True)
+            # note: we pass the file object to write_file rather than reading the file first,
+            # to ensure the data is read in chunks (which prevents problems in Python 3.9+);
+            # cfr. https://github.com/easybuilders/easybuild-framework/issues/3455
+            # and https://bugs.python.org/issue42853
+            write_file(path, url_fd, forced=forced, backup=True)
             _log.info("Downloaded file %s from url %s to %s" % (filename, url, path))
             downloaded = True
             url_fd.close()

test/framework/filetools.py

@@ -705,6 +705,26 @@ def test_read_write_file(self):
         # test use of 'mode' in read_file
         self.assertEqual(ft.read_file(foo, mode='rb'), b'bar')
 
+    def test_write_file_obj(self):
+        """Test writing from a file-like object directly"""
+        # Write a text file
+        fp = os.path.join(self.test_prefix, 'test.txt')
+        fp_out = os.path.join(self.test_prefix, 'test_out.txt')
+        ft.write_file(fp, b'Hyphen: \xe2\x80\x93\nEuro sign: \xe2\x82\xac\na with dots: \xc3\xa4')
+
+        with ft.open_file(fp, 'rb') as fh:
+            ft.write_file(fp_out, fh)
+        self.assertEqual(ft.read_file(fp_out), ft.read_file(fp))
+
+        # Write a binary file
+        fp = os.path.join(self.test_prefix, 'test.bin')
+        fp_out = os.path.join(self.test_prefix, 'test_out.bin')
+        ft.write_file(fp, b'\x00\x01'+os.urandom(42)+b'\x02\x03')
+
+        with ft.open_file(fp, 'rb') as fh:
+            ft.write_file(fp_out, fh)
+        self.assertEqual(ft.read_file(fp_out, mode='rb'), ft.read_file(fp, mode='rb'))
+
     def test_is_binary(self):
         """Test is_binary function."""