Clean up file I/O operations

Author: aakallam

src/clusterfuzz/_internal/bot/tasks/setup.py

@@ -841,11 +841,14 @@ def archive_testcase_and_dependencies_in_gcs(resource_list, testcase_path: str,
 
   logs.info(f'Testcase and related files :\n{resource_list}')
 
+  upload_file_path = testcase_path
   if len(resource_list) <= 1:
     # If this does not have any resources, just save the testcase.
     # TODO(flowerhack): Update this when we teach CF how to download testcases.
     try:
-      file_handle = open(testcase_path, 'rb')
+      with open(testcase_path, 'rb') as _:
+        # Verify file can be opened.
+        pass
     except OSError:
       logs.error(f'Unable to open testcase {testcase_path}.')
       return None, None, None
@@ -879,31 +882,37 @@ def archive_testcase_and_dependencies_in_gcs(resource_list, testcase_path: str,
 
     # Create the archive.
     zip_path = os.path.join(environment.get_value('INPUT_DIR'), zip_filename)
-    zip_file = zipfile.ZipFile(zip_path, 'w')
-    for file_name in resource_list:
-      if os.path.exists(file_name):
-        relative_filename = file_name[base_len:]
-        zip_file.write(file_name, relative_filename, zipfile.ZIP_DEFLATED)
-    zip_file.close()
+    with zipfile.ZipFile(zip_path, 'w') as zip_file:
+      for file_name in resource_list:
+        if os.path.exists(file_name):
+          relative_filename = file_name[base_len:]
+          zip_file.write(file_name, relative_filename, zipfile.ZIP_DEFLATED)
 
     try:
-      file_handle = open(zip_path, 'rb')
+      with open(zip_path, 'rb') as _:
+        # Verify file can be opened.
+        pass
+      upload_file_path = zip_path
     except OSError:
       logs.error(f'Unable to open testcase archive {zip_path}.')
       return None, None, None
+    finally:
+      if zip_path:
+        shell.remove_file(zip_path)
 
     archived = True
     absolute_filename = testcase_path[base_len:]
 
-  if not storage.upload_signed_url(file_handle, upload_url):
-    logs.error('Failed to upload testcase.')
-    return None, None, None
-
-  file_handle.close()
+  try:
+    with open(upload_file_path, 'rb') as file_handle:
+      if not storage.upload_signed_url(file_handle, upload_url):
+        logs.error('Failed to upload testcase.')
+        return None, None, None
 
-  # Don't need the archive after writing testcase to blobstore.
-  if zip_path:
-    shell.remove_file(zip_path)
+  finally:
+    # Don't need the archive after writing testcase to blobstore.
+    if zip_path:
+      shell.remove_file(zip_path)
 
   return archived, absolute_filename, zip_filename
 

src/clusterfuzz/_internal/bot/tasks/unpack_task.py

@@ -89,9 +89,8 @@ def execute_task(metadata_id, job_type):
       continue
 
     try:
-      file_handle = open(absolute_file_path, 'rb')
-      blob_key = blobs.write_blob(file_handle)
-      file_handle.close()
+      with open(absolute_file_path, 'rb') as file_handle:
+        blob_key = blobs.write_blob(file_handle)
     except:
       blob_key = None
 

src/clusterfuzz/_internal/bot/tasks/utasks/minimize_task.py

@@ -1028,45 +1028,42 @@ def store_minimized_testcase(
 
   # Prepare the file.
   zip_path = None
+  testcase_filepath = None
   if testcase.archive_state:
     if len(file_list) > 1:
       testcase.archive_state |= data_types.ArchiveStatus.MINIMIZED
       minimize_task_output.archive_state = testcase.archive_state
       zip_path = os.path.join(
           environment.get_value('INPUT_DIR'), '%d.zip' % testcase.key.id())
-      zip_file = zipfile.ZipFile(zip_path, 'w')
-      count = 0
-      filtered_file_list = []
-      for file_name in file_list:
-        absolute_filename = os.path.join(base_directory, file_name)
-        is_file = os.path.isfile(absolute_filename)
-        if file_to_run_data and is_file and os.path.getsize(
-            absolute_filename) == 0 and (os.path.basename(
-                absolute_filename).encode('utf-8') not in file_to_run_data):
-          continue
-        if not os.path.exists(absolute_filename):
-          continue
-        zip_file.write(absolute_filename, file_name, zipfile.ZIP_DEFLATED)
-        if is_file:
-          count += 1
-          filtered_file_list.append(absolute_filename)
-
-      zip_file.close()
+      with zipfile.ZipFile(zip_path, 'w') as zip_file:
+        count = 0
+        filtered_file_list = []
+        for file_name in file_list:
+          absolute_filename = os.path.join(base_directory, file_name)
+          is_file = os.path.isfile(absolute_filename)
+          if file_to_run_data and is_file and os.path.getsize(
+              absolute_filename) == 0 and (os.path.basename(
+                  absolute_filename).encode('utf-8') not in file_to_run_data):
+            continue
+          if not os.path.exists(absolute_filename):
+            continue
+          zip_file.write(absolute_filename, file_name, zipfile.ZIP_DEFLATED)
+          if is_file:
+            count += 1
+            filtered_file_list.append(absolute_filename)
+
       try:
         if count > 1:
-          file_handle = open(zip_path, 'rb')
+          testcase_filepath = zip_path
         else:
           if not filtered_file_list:
             # We minimized everything. The only thing needed to reproduce is the
             # interaction gesture.
-            file_path = file_list[0]
-            file_handle = open(file_path, 'wb')
-            file_handle.close()
+            testcase_filepath = file_list[0]
           else:
-            file_path = filtered_file_list[0]
-          file_handle = open(file_path, 'rb')
-          testcase.absolute_path = os.path.join(base_directory,
-                                                os.path.basename(file_path))
+            testcase_filepath = filtered_file_list[0]
+          testcase.absolute_path = os.path.join(
+              base_directory, os.path.basename(testcase_filepath))
           minimize_task_output.absolute_path = testcase.absolute_path
           testcase.archive_state &= ~data_types.ArchiveStatus.MINIMIZED
           minimize_task_output.archive_state = testcase.archive_state
@@ -1075,20 +1072,20 @@ def store_minimized_testcase(
         return
     else:
       absolute_filename = os.path.join(base_directory, file_list[0])
-      file_handle = open(absolute_filename, 'rb')
+      testcase_filepath = absolute_filename
       testcase.archive_state &= ~data_types.ArchiveStatus.MINIMIZED
       minimize_task_output.archive_state = testcase.archive_state
   else:
-    file_handle = open(file_list[0], 'rb')
+    testcase_filepath = file_list[0]
     testcase.archive_state &= ~data_types.ArchiveStatus.MINIMIZED
     minimize_task_output.archive_state = testcase.archive_state
 
   # Store the testcase.
-  data = file_handle.read()
-  storage.upload_signed_url(data, minimize_task_input.testcase_upload_url)
-  minimized_keys = minimize_task_input.testcase_blob_name
-  file_handle.close()
+  with open(testcase_filepath, 'rb') as file_handle:
+    data = file_handle.read()
+    storage.upload_signed_url(data, minimize_task_input.testcase_upload_url)
 
+  minimized_keys = minimize_task_input.testcase_blob_name
   testcase.minimized_keys = minimized_keys
   minimize_task_output.minimized_keys = minimized_keys
 
@@ -1318,9 +1315,8 @@ def combine_tokens(tokens):
       return b''
 
     # TODO(mbarbella): Allow token combining functions to write files directly.
-    handle = open(combined_file_path, 'rb')
-    result = handle.read()
-    handle.close()
+    with open(combined_file_path, 'rb') as handle:
+      result = handle.read()
 
     shell.remove_file(combined_file_path)
     return result

src/clusterfuzz/_internal/platforms/android/device.py

@@ -214,22 +214,20 @@ def configure_system_build_properties():
 
   # Write new build.prop.
   new_build_prop_path = os.path.join(bot_tmp_directory, 'new.prop')
-  old_build_prop_file_content = open(old_build_prop_path)
-  new_build_prop_file_content = open(new_build_prop_path, 'w')
-  new_content_notification = '### CHANGED OR ADDED PROPERTIES ###'
-  for line in old_build_prop_file_content:
-    property_name = line.split('=')[0].strip()
-    if property_name in BUILD_PROPERTIES:
-      continue
-    if new_content_notification in line:
-      continue
-    new_build_prop_file_content.write(line)
-
-  new_build_prop_file_content.write(new_content_notification + '\n')
-  for flag, value in BUILD_PROPERTIES.items():
-    new_build_prop_file_content.write('%s=%s\n' % (flag, value))
-  old_build_prop_file_content.close()
-  new_build_prop_file_content.close()
+  with open(old_build_prop_path) as old_build_prop_file_content, \
+      open(new_build_prop_path, 'w') as new_build_prop_file_content:
+    new_content_notification = '### CHANGED OR ADDED PROPERTIES ###'
+    for line in old_build_prop_file_content:
+      property_name = line.split('=')[0].strip()
+      if property_name in BUILD_PROPERTIES:
+        continue
+      if new_content_notification in line:
+        continue
+      new_build_prop_file_content.write(line)
+
+    new_build_prop_file_content.write(new_content_notification + '\n')
+    for flag, value in BUILD_PROPERTIES.items():
+      new_build_prop_file_content.write('%s=%s\n' % (flag, value))
 
   # Keep verified boot disabled for M and higher releases. This makes it easy
   # to modify system's app_process to load asan libraries.

src/clusterfuzz/_internal/tests/core/base/utils_test.py

@@ -618,8 +618,9 @@ def test_file_exists_valid_content(self):
 
   def test_file_exists_empty_content(self):
     """Test reading empty manifest file."""
-    f = open(self.test_path, 'w')
-    f.close()
+    with open(self.test_path, 'w') as _:
+      # Opens file with 'w' to clear it.
+      pass
     self.assertEqual(utils.current_source_version(), '')
 
   def test_file_exists_invalid_content(self):

src/clusterfuzz/_internal/tests/core/bot/minimizer/base_minimizer_tester.py

@@ -24,7 +24,8 @@ def setUp(self):
 
   def _mock_test_function(self, data_file):
     """Mock test function to reduce time minimizer takes and simplify tests."""
-    data = open(data_file, 'rb').read()
+    with open(data_file, 'rb') as f:
+      data = f.read()
     if b'error' in data:
       return False
     return True

src/clusterfuzz/_internal/tests/core/bot/minimizer/unicode_minimizer_test.py

@@ -50,7 +50,8 @@ def setUp(self):
         self._mock_test_function)
 
   def _mock_test_function(self, data_file):
-    data = open(data_file, 'rb').read()
+    with open(data_file, 'rb') as f:
+      data = f.read()
     if data in self.crash_programs:
       return False
     if data in self.no_crash_programs:

src/clusterfuzz/_internal/tests/core/bot/tasks/update_task_test.py

@@ -248,23 +248,23 @@ def test_all_files_are_correctly_unpacked(self):
     zipfile_path = os.path.join(self._make_temp_dir(), 'linux.zip')
     storage.copy_file_from('gs://clusterfuzz-deployment/linux-3.zip',
                            zipfile_path)
-    archive = zipfile.ZipFile(zipfile_path)
-    for file in archive.namelist():
-      if os.path.basename(file) == 'adb':
-        continue
-      self.assertTrue(os.path.exists(os.path.join(self.temp_directory, file)))
+    with zipfile.ZipFile(zipfile_path) as archive:
+      for file in archive.namelist():
+        if os.path.basename(file) == 'adb':
+          continue
+        self.assertTrue(os.path.exists(os.path.join(self.temp_directory, file)))
 
   def test_archive_execute_permission_is_respected(self):
     """Tests that the exectuable bit is correctly propagated to source files."""
     update_task.update_source_code()
     zipfile_path = os.path.join(self._make_temp_dir(), 'linux.zip')
     storage.copy_file_from('gs://clusterfuzz-deployment/linux-3.zip',
                            zipfile_path)
-    archive = zipfile.ZipFile(zipfile_path)
-    for member in archive.infolist():
-      if os.path.basename(member.filename) == 'adb':
-        continue
-      mode = (member.external_attr >> 16) & 0o7777
-      filepath = os.path.join(self.temp_directory, member.filename)
-      if mode & 0o100:
-        self.assertTrue(os.access(filepath, os.X_OK))
+    with zipfile.ZipFile(zipfile_path) as archive:
+      for member in archive.infolist():
+        if os.path.basename(member.filename) == 'adb':
+          continue
+        mode = (member.external_attr >> 16) & 0o7777
+        filepath = os.path.join(self.temp_directory, member.filename)
+        if mode & 0o100:
+          self.assertTrue(os.access(filepath, os.X_OK))

src/clusterfuzz/_internal/tests/core/bot/tasks/utasks/minimize_task_test.py

@@ -399,7 +399,8 @@ def test_irreducible_input(self):
 
     def mock_test_function_unicode_dependent(data_file):
       """If crash -> False, otherwise -> True."""
-      data = open(data_file, 'rb').read()
+      with open(data_file, 'rb') as f:
+        data = f.read()
 
       if data == b'consol\\u0065.log(42);':
         return False
@@ -417,7 +418,8 @@ def test_reducible_input(self):
 
     def mock_test_function_unicode_independent(data_file):
       """If crash -> False, otherwise -> True."""
-      data = open(data_file, 'rb').read()
+      with open(data_file, 'rb') as f:
+        data = f.read()
 
       if data in (b'consol\\u0065.log(42);', b'console.log(42);'):
         return False

src/local/butler/common.py

@@ -227,13 +227,12 @@ def _install_chromedriver():
       constants.CHROMEDRIVER_DOWNLOAD_PATTERN.format(
           version=version, archive_name=archive_name))
   archive_io = io.BytesIO(archive_request.read())
-  chromedriver_archive = zipfile.ZipFile(archive_io)
+  with zipfile.ZipFile(archive_io) as chromedriver_archive:
+    chromedriver_path = get_chromedriver_path()
+    output_directory = os.path.dirname(chromedriver_path)
+    chromedriver_binary = os.path.basename(chromedriver_path)
 
-  chromedriver_path = get_chromedriver_path()
-  output_directory = os.path.dirname(chromedriver_path)
-  chromedriver_binary = os.path.basename(chromedriver_path)
-
-  chromedriver_archive.extract(chromedriver_binary, output_directory)
+    chromedriver_archive.extract(chromedriver_binary, output_directory)
   os.chmod(chromedriver_path, 0o750)
   print(f'Installed chromedriver at: {chromedriver_path}')