Allow load_icgem_gdf to take open file objects (#155)

Previously just assumed `fname` was a string and called `open` on it.
Now, check if it's a file-like object (has a `read` method) and don't
open it if that's the case. Use `contextlib` to keep the `with`
statement in case we do have to open the file. This is useful if the
file is compressed and can be opened with `bz2.open`, `gzip.open`, etc
before passing on to the function.

harmonica/io.py

@@ -2,6 +2,8 @@
 Functions for interacting with standarized data files that contain gravity or
 magnetic geophysical data.
 """
+import contextlib
+
 import numpy as np
 import xarray as xr
 
@@ -19,8 +21,8 @@ def load_icgem_gdf(fname, **kwargs):
 
  Parameters
  ----------
- fname : string
- Name of the ICGEM .gdf file
+ fname : string or file-like object
+ Name of the ICGEM ``.gdf`` file or an open file object to read from.
  kwargs
  Extra keyword arguments to this function will be passed to
  :func:`numpy.loadtxt`.
@@ -82,7 +84,15 @@ def _read_gdf_file(fname, **kwargs):
  """
  Read ICGEM gdf file and returns metadata dict and data in cols as np.array
  """
- with open(fname) as gdf_file:
+ # If it's an open file, does nothing. Otherwise, open and add to the
+ # context lib stack to make it close when exiting the with block.
+ with contextlib.ExitStack() as stack:
+ if hasattr(fname, "read"):
+ # It's a file object
+ gdf_file = fname
+ else:
+ # It's a file path
+ gdf_file = stack.enter_context(open(fname))
  # Read the header and extract metadata
  metadata = {}
  metadata_line = True

harmonica/tests/test_icgem.py

@@ -33,6 +33,28 @@ def test_load_icgem_gdf():
  npt.assert_allclose(height, icgem_grd.height_over_ell.values)
 
 
+def test_load_icgem_gdf_open_file():
+ "Check if load_icgem_gdf works if given an open file instead of string"
+ fname = os.path.join(TEST_DATA_DIR, "icgem-sample.gdf")
+ with open(fname, "r") as open_file:
+ icgem_grd = load_icgem_gdf(open_file)
+
+ s, n, w, e = 16, 28, 150, 164
+ nlat, nlon = 7, 8
+ shape = (nlat, nlon)
+ lat = np.linspace(s, n, nlat, dtype="float64")
+ lon = np.linspace(w, e, nlon, dtype="float64")
+ true_data = np.array([np.arange(nlon)] * nlat, dtype="float64")
+ height = 1100 * np.ones(shape)
+
+ assert icgem_grd.dims["latitude"] == nlat
+ assert icgem_grd.dims["longitude"] == nlon
+ npt.assert_equal(icgem_grd.longitude.values, lon)
+ npt.assert_equal(icgem_grd.latitude.values, lat)
+ npt.assert_allclose(true_data, icgem_grd.sample_data.values)
+ npt.assert_allclose(height, icgem_grd.height_over_ell.values)
+
+
 def test_load_icgem_gdf_with_height():
  "Check if load_icgem_gdf reads an ICGEM file with height column"
  fname = os.path.join(TEST_DATA_DIR, "icgem-sample-with-height.gdf")