From d9ae36728171f3c8bcb18a5c1a313c97516ab197 Mon Sep 17 00:00:00 2001
From: Ivan Kravets <me@ikravets.com>
Date: Sun, 24 Sep 2017 00:33:12 +0300
Subject: [PATCH] Improve archive unpacker

---
 platformio/managers/package.py |  4 ++--
 platformio/unpacker.py         | 35 ++++++++++++++++++++--------------
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/platformio/managers/package.py b/platformio/managers/package.py
index 7d1f9216..47e63539 100644
--- a/platformio/managers/package.py
+++ b/platformio/managers/package.py
@@ -187,8 +187,8 @@ class PkgInstallerMixin(object):
 
     @staticmethod
     def unpack(source_path, dest_dir):
-        fu = FileUnpacker(source_path, dest_dir)
-        return fu.start()
+        with FileUnpacker(source_path) as fu:
+            return fu.unpack(dest_dir)
 
     @staticmethod
     def get_install_dirname(manifest):
diff --git a/platformio/unpacker.py b/platformio/unpacker.py
index 036dcc4c..4e789d1c 100644
--- a/platformio/unpacker.py
+++ b/platformio/unpacker.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from os import chmod
-from os.path import join, splitext
+from os.path import join
 from tarfile import open as tarfile_open
 from time import mktime
 from zipfile import ZipFile
@@ -39,6 +39,9 @@ class ArchiveBase(object):
     def after_extract(self, item, dest_dir):
         pass
 
+    def close(self):
+        self._afo.close()
+
 
 class TARArchive(ArchiveBase):
 
@@ -76,28 +79,32 @@ class ZIPArchive(ArchiveBase):
 
 class FileUnpacker(object):
 
-    def __init__(self, archpath, dest_dir="."):
-        self._archpath = archpath
-        self._dest_dir = dest_dir
+    def __init__(self, archpath):
+        self.archpath = archpath
         self._unpacker = None
 
-        _, archext = splitext(archpath.lower())
-        if archext in (".gz", ".bz2"):
-            self._unpacker = TARArchive(archpath)
-        elif archext == ".zip":
-            self._unpacker = ZIPArchive(archpath)
-
+    def __enter__(self):
+        if self.archpath.lower().endswith((".gz", ".bz2")):
+            self._unpacker = TARArchive(self.archpath)
+        elif self.archpath.lower().endswith(".zip"):
+            self._unpacker = ZIPArchive(self.archpath)
         if not self._unpacker:
-            raise UnsupportedArchiveType(archpath)
+            raise UnsupportedArchiveType(self.archpath)
+        return self
 
-    def start(self):
+    def __exit__(self, *args):
+        if self._unpacker:
+            self._unpacker.close()
+
+    def unpack(self, dest_dir="."):
+        assert self._unpacker
         if app.is_disabled_progressbar():
             click.echo("Unpacking...")
             for item in self._unpacker.get_items():
-                self._unpacker.extract_item(item, self._dest_dir)
+                self._unpacker.extract_item(item, dest_dir)
         else:
             items = self._unpacker.get_items()
             with click.progressbar(items, label="Unpacking") as pb:
                 for item in pb:
-                    self._unpacker.extract_item(item, self._dest_dir)
+                    self._unpacker.extract_item(item, dest_dir)
         return True