Add close() and context management to ubi_file

Author: AT0myks

ubireader/scripts/ubireader_display_blocks.py

@@ -184,6 +184,8 @@ def main():
         if match:                
             blocks.append(ubi_obj.blocks[block])
 
+    ufile_obj.close()
+
     print('\nBlock matches: %s' % len(blocks))
 
     for block in blocks:

ubireader/scripts/ubireader_display_info.py

@@ -185,5 +185,8 @@ def main():
     else:
         print('Something went wrong to get here.')
 
+    ufile_obj.close()
+
+
 if __name__=='__main__':
     main()

ubireader/scripts/ubireader_extract_files.py

@@ -198,6 +198,8 @@ def main():
     else:
         print('Something went wrong to get here.')
 
+    ufile_obj.close()
+
 
 if __name__=='__main__':
     main()

ubireader/scripts/ubireader_extract_images.py

@@ -167,5 +167,8 @@ def main():
                 for block in image.volumes[volume].reader(ubi_obj):
                     f.write(block)
 
+    ufile_obj.close()
+
+
 if __name__=='__main__':
     main()

ubireader/scripts/ubireader_list_files.py

@@ -173,5 +173,8 @@ def main():
     else:
         print('Something went wrong to get here.')
 
+    ufile_obj.close()
+
+
 if __name__=='__main__':
     main()

ubireader/scripts/ubireader_utils_info.py

@@ -343,5 +343,8 @@ def main():
         # Create build scripts.
         make_files(ubi_obj, outpath)
 
+    ufile_obj.close()
+
+
 if __name__=='__main__':
     main()

ubireader/ubi_io.py

@@ -88,6 +88,11 @@ def __init__(self, path, block_size, start_offset=0, end_offset=None):
         self._last_read_addr = self._fhandle.tell()
         self.is_valid = True
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
 
     def _set_start(self, i):
         self._start_offset = i
@@ -105,6 +110,8 @@ def _get_block_size(self):
         return self._block_size
     block_size = property(_get_block_size)
 
+    def close(self):
+        self._fhandle.close()
 
     def seek(self, offset):
         self._fhandle.seek(offset)