fixes #337

Author: jph00

fastcore/xtras.py

@@ -147,19 +147,26 @@ def loads_multi(s:str):
         yield obj
         s = s[pos:]
 
+# Cell
+def _unpack(fname, out):
+    shutil.unpack_archive(str(fname), str(out))
+    ls = out.ls()
+    return ls[0] if len(ls) == 1 else out
+
 # Cell
 def untar_dir(fname, dest, rename=False, overwrite=False):
     "untar `file` into `dest`, creating a directory if the root contains more than one item"
     with tempfile.TemporaryDirectory(dir='.') as d:
         out = Path(d)/remove_suffix(Path(fname).stem, '.tar')
         out.mkdir()
-        shutil.unpack_archive(str(fname), str(out))
-        ls = out.ls()
-        src = ls[0] if len(ls) == 1 else out
-        dest = dest/(out if rename else src).name
+        if rename: dest = dest/out.name
+        else:
+            src = _unpack(fname, out)
+            dest = dest/src.name
         if dest.exists():
             if overwrite: shutil.rmtree(dest)
             else: return dest
+        if rename: src = _unpack(fname, out)
         shutil.move(str(src), dest)
         return dest
 

nbs/03_xtras.ipynb

@@ -618,7 +618,7 @@
     {
      "data": {
       "text/plain": [
-       "['g', 'a', 'b', 'h', 'd', 'f', 'c', 'e']"
+       "['b', 'a', 'g', 'h', 'f', 'e', 'c', 'd']"
       ]
      },
      "execution_count": null,
@@ -1000,6 +1000,19 @@
     "test_eq(list(loads_multi(tst)), [{'a': 1}, {'b': 2}])"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def _unpack(fname, out):\n",
+    "    shutil.unpack_archive(str(fname), str(out))\n",
+    "    ls = out.ls()\n",
+    "    return ls[0] if len(ls) == 1 else out"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -1012,13 +1025,14 @@
     "    with tempfile.TemporaryDirectory(dir='.') as d:\n",
     "        out = Path(d)/remove_suffix(Path(fname).stem, '.tar')\n",
     "        out.mkdir()\n",
-    "        shutil.unpack_archive(str(fname), str(out))\n",
-    "        ls = out.ls()\n",
-    "        src = ls[0] if len(ls) == 1 else out\n",
-    "        dest = dest/(out if rename else src).name\n",
+    "        if rename: dest = dest/out.name\n",
+    "        else:\n",
+    "            src = _unpack(fname, out)\n",
+    "            dest = dest/src.name\n",
     "        if dest.exists():\n",
     "            if overwrite: shutil.rmtree(dest)\n",
     "            else: return dest\n",
+    "        if rename: src = _unpack(fname, out)\n",
     "        shutil.move(str(src), dest)\n",
     "        return dest"
    ]
@@ -1029,13 +1043,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def test_untar(foldername, **kwargs):\n",
+    "def test_untar(foldername, rename=False, **kwargs):\n",
     "    with tempfile.TemporaryDirectory() as d:\n",
     "        nm = os.path.join(d, 'a')\n",
     "        shutil.make_archive(nm, 'gztar', **kwargs)\n",
     "        with tempfile.TemporaryDirectory() as d2:\n",
     "            d2 = Path(d2)\n",
-    "            untar_dir(nm+'.tar.gz', d2)\n",
+    "            untar_dir(nm+'.tar.gz', d2, rename=rename)\n",
     "            test_eq(d2.ls(), [d2/foldername])"
    ]
   },
@@ -1056,6 +1070,22 @@
     "test_untar('images', base_dir='images')"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If `rename` then the directory created is named based on the archive, without extension:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_untar('a', base_dir='images', rename=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},