Repository: fatiando/harmonica
Pull request #155 (2020-03-08 19:09:57)

Allow load_icgem_gdf to take open file objects
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.

<!--
Please describe changes proposed and WHY you made them. If fixing an issue,
include the text "Fixes #XXX" (replace XXX by the issue number. GitHub will
automatically close it when this gets merged.
-->





**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.



Patch:
diff --git a/harmonica/io.py b/harmonica/io.py
index b6e3956..4d1a712 100644
--- a/harmonica/io.py
+++ b/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


Test patch:
diff --git a/harmonica/tests/test_icgem.py b/harmonica/tests/test_icgem.py
index 836dbf6..a3a3205 100644
--- a/harmonica/tests/test_icgem.py
+++ b/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")