fixes #128

Author: jph00

fastcore/_nbdev.py

@@ -116,6 +116,7 @@
          "Path.readlines": "02_utils.ipynb",
          "Path.mk_write": "02_utils.ipynb",
          "Path.ls": "02_utils.ipynb",
+         "open_file": "02_utils.ipynb",
          "save_pickle": "02_utils.ipynb",
          "load_pickle": "02_utils.ipynb",
          "Path.__repr__": "02_utils.ipynb",

fastcore/utils.py

@@ -7,8 +7,8 @@
            'Float', 'tuplify', 'detuplify', 'replicate', 'uniqueify', 'setify', 'merge', 'is_listy', 'range_of',
            'groupby', 'last_index', 'shufflish', 'IterLen', 'ReindexCollection', 'num_methods', 'rnum_methods',
            'inum_methods', 'fastuple', 'trace', 'compose', 'maps', 'partialler', 'mapped', 'instantiate', 'using_attr',
-           'Self', 'Self', 'save_pickle', 'load_pickle', 'bunzip', 'join_path_file', 'urlread', 'urljson', 'run',
-           'do_request', 'sort_by_run', 'PrettyString', 'round_multiple', 'even_mults', 'num_cpus', 'add_props',
+           'Self', 'Self', 'open_file', 'save_pickle', 'load_pickle', 'bunzip', 'join_path_file', 'urlread', 'urljson',
+           'run', 'do_request', 'sort_by_run', 'PrettyString', 'round_multiple', 'even_mults', 'num_cpus', 'add_props',
            'ContextManagers', 'typed', 'set_num_threads', 'ProcessPoolExecutor', 'ThreadPoolExecutor', 'parallel',
            'run_procs', 'parallel_gen', 'threaded']
 
@@ -17,7 +17,7 @@
 from .foundation import *
 from functools import wraps
 
-import mimetypes,bz2,pickle,random,json,urllib,subprocess,shlex
+import mimetypes,bz2,pickle,random,json,urllib,subprocess,shlex,bz2,gzip
 from contextlib import contextmanager
 from pdb import set_trace
 from urllib.request import Request,urlopen
@@ -528,21 +528,24 @@ def ls(self:Path, n_max=None, file_type=None, file_exts=None):
     if n_max is not None: res = itertools.islice(res, n_max)
     return L(res)
 
+# Cell
+def open_file(fn, mode='r'):
+    "Open a file, with optional compression if gz or bz2 suffix"
+    if isinstance(fn, io.IOBase): return fn
+    fn = Path(fn)
+    if   fn.suffix=='.bz2': return bz2.BZ2File(fn, mode)
+    elif fn.suffix=='.gz' : return gzip.GzipFile(fn, mode)
+    else: return open(fn,mode)
+
 # Cell
 def save_pickle(fn, o):
     "Save a pickle file, to a file name or opened file"
-    fn = Path(fn)
-    if not isinstance(fn, io.IOBase): fn = open(fn,'wb')
-    try: pickle.dump(o, fn)
-    finally: fn.close()
+    with open_file(fn, 'wb') as f: pickle.dump(o, f)
 
 # Cell
 def load_pickle(fn):
     "Load a pickle file from a file name or opened file"
-    fn = Path(fn)
-    if not isinstance(fn, io.IOBase): fn = open(fn,'rb')
-    try: return pickle.load(fn)
-    finally: fn.close()
+    with open_file(fn, 'rb') as f: return pickle.load(f)
 
 # Cell
 @patch

nbs/02_utils.ipynb

@@ -20,7 +20,7 @@
     "from fastcore.foundation import *\n",
     "from functools import wraps\n",
     "\n",
-    "import mimetypes,bz2,pickle,random,json,urllib,subprocess,shlex\n",
+    "import mimetypes,bz2,pickle,random,json,urllib,subprocess,shlex,bz2,gzip\n",
     "from contextlib import contextmanager\n",
     "from pdb import set_trace\n",
     "from urllib.request import Request,urlopen\n",
@@ -2981,19 +2981,33 @@
     "txt_files[0],ipy_files[0]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#hide\n",
+    "path = Path()\n",
+    "pkl = pickle.dumps(path)\n",
+    "p2 = pickle.loads(pkl)\n",
+    "test_eq(path.ls()[0], p2.ls()[0])"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "#export\n",
-    "def save_pickle(fn, o):\n",
-    "    \"Save a pickle file, to a file name or opened file\"\n",
+    "def open_file(fn, mode='r'):\n",
+    "    \"Open a file, with optional compression if gz or bz2 suffix\"\n",
+    "    if isinstance(fn, io.IOBase): return fn\n",
     "    fn = Path(fn)\n",
-    "    if not isinstance(fn, io.IOBase): fn = open(fn,'wb')\n",
-    "    try: pickle.dump(o, fn)\n",
-    "    finally: fn.close()"
+    "    if   fn.suffix=='.bz2': return bz2.BZ2File(fn, mode)\n",
+    "    elif fn.suffix=='.gz' : return gzip.GzipFile(fn, mode)\n",
+    "    else: return open(fn,mode)"
    ]
   },
   {
@@ -3003,12 +3017,9 @@
    "outputs": [],
    "source": [
     "#export\n",
-    "def load_pickle(fn):\n",
-    "    \"Load a pickle file from a file name or opened file\"\n",
-    "    fn = Path(fn)\n",
-    "    if not isinstance(fn, io.IOBase): fn = open(fn,'rb')\n",
-    "    try: return pickle.load(fn)\n",
-    "    finally: fn.close()"
+    "def save_pickle(fn, o):\n",
+    "    \"Save a pickle file, to a file name or opened file\"\n",
+    "    with open_file(fn, 'wb') as f: pickle.dump(o, f)"
    ]
   },
   {
@@ -3017,11 +3028,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with tempfile.NamedTemporaryFile() as f:\n",
-    "    fn = Path(f.name)\n",
-    "    save_pickle(fn, 't')\n",
-    "    t = load_pickle(fn)\n",
-    "test_eq(t,'t')"
+    "#export\n",
+    "def load_pickle(fn):\n",
+    "    \"Load a pickle file from a file name or opened file\"\n",
+    "    with open_file(fn, 'rb') as f: return pickle.load(f)"
    ]
   },
   {
@@ -3030,10 +3040,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#hide\n",
-    "pkl = pickle.dumps(path)\n",
-    "p2 = pickle.loads(pkl)\n",
-    "test_eq(path.ls()[0], p2.ls()[0])"
+    "for suf in '.pkl','.bz2','.gz':\n",
+    "    with tempfile.NamedTemporaryFile(suffix=suf) as f:\n",
+    "        fn = Path(f.name)\n",
+    "        save_pickle(fn, 't')\n",
+    "        t = load_pickle(fn)\n",
+    "    test_eq(t,'t')"
    ]
   },
   {