Merge pull request #455 from vieramercado/feature/local-index-file-367

Using local index file #367

Author: blaylockbk

src/herbie/core.py

@@ -19,6 +19,7 @@
 from io import StringIO
 from shutil import which
 from typing import Literal, Optional, Union
+from urllib.parse import urlparse
 
 import cfgrib
 import pandas as pd
@@ -397,6 +398,9 @@ def _check_idx(self, url: str, verbose: bool = False) -> tuple[bool, Optional[st
 
         if verbose:
             print(f"🐜 {self.IDX_SUFFIX=}")
+        
+         # Initialize variable to avoid UnboundLocalError
+        idx_exists = False
 
         # Loop through IDX_SUFFIX options until we find one that exists
         for i in self.IDX_SUFFIX:
@@ -405,7 +409,12 @@ def _check_idx(self, url: str, verbose: bool = False) -> tuple[bool, Optional[st
             else:
                 idx_url = url + i
 
-            idx_exists = requests.head(idx_url).ok
+            try:
+                idx_exists = requests.head(idx_url).ok
+            except Exception as e:
+                if verbose:
+                    print(f"Unable to get index file from {idx_url} due to error {str(e)}")
+
             if verbose:
                 print(f"🐜 {idx_url=}")
                 print(f"🐜 {idx_exists=}")
@@ -431,8 +440,6 @@ def find_grib(self) -> tuple[Optional[Union[Path, str]], Optional[str]]:
         local_grib = self.get_localFilePath()
         if local_grib.exists() and not self.overwrite:
             return local_grib, "local"
-            # NOTE: We will still get the idx files from a remote
-            #       because they aren't stored locally, or are they?   # TODO: If the idx file is local, then use that
 
         # If priority list is set, we want to search SOURCES in that
         # priority order. If priority is None, then search all SOURCES
@@ -465,8 +472,37 @@ def find_grib(self) -> tuple[Optional[Union[Path, str]], Optional[str]]:
 
         return (None, None)
 
-    def find_idx(self) -> tuple[Optional[Union[Path, str]], Optional[str]]:
+    def find_idx(self, overwrite: bool = False) -> tuple[Optional[Union[Path, str]], Optional[str]]:
         """Find an index file for the GRIB file."""
+
+        # But first, if overwrite is False, then check if the GRIB2 inventory file exists locally.
+        if not overwrite:
+            local_grib = self.get_localFilePath()
+            for suffix in self.IDX_SUFFIX:
+                # There is no easy way to know if the index suffix needs to be appended or overwrite the grib suffix.
+                # Hence there is a need to do various checks to verify if the index file exists locally or not.
+
+                # This is the case of GFS, where the grib file ends with ".grb2" and the index file is ".grb2.inv"
+                # The index suffix needs to overwrite the grib file suffix, now ending in ".grb2.inv"
+                if suffix.startswith(local_grib.suffix):
+                    local_idx = local_grib.with_suffix(suffix)
+                
+                # This is the case of GFS, where the grib file ends with ".f000" and the index file is ".idx"
+                # The index suffix needs to be appended to the grib file suffix, now ending in ".f000.idx"
+                else:
+                    local_idx = local_grib.with_suffix(local_grib.suffix + suffix)
+
+                if local_idx.exists() and not self.overwrite:
+                    return (local_idx, "local")
+                
+                # If the index file does not exists locally, we will need to try another variation of the index file name.
+                # This is the case of IFS, where the grib file ends with ".grib2" and the index file is ".index"
+                # The index suffix needs to overwrite the grib file suffix, now ending in ".index"
+                else:
+                    local_idx = local_grib.with_suffix(suffix)
+                    if local_idx.exists() and not self.overwrite:
+                        return (local_idx, "local")    
+
         # If priority list is set, we want to search SOURCES in that
         # priority order. If priority is None, then search all SOURCES
         # in the order given by the model template file.
@@ -591,10 +627,27 @@ def get_localFilePath(
 
         return localFilePath
 
+    def get_localIndexFilePath(self) -> Path:
+        """Get full path to the local index file."""
+
+        # Get the local file path, which creates the directory structure
+        local_file_path = self.get_localFilePath()
+
+        # Get the directory in the local file path
+        dir_path = os.path.dirname(local_file_path)
+        
+        # Get the filename from the index URL path
+        index_filename = os.path.basename(urlparse(self.idx).path)
+
+        # Create new path with the index file name
+        index_file_path = os.path.join(dir_path, index_filename)
+
+        return index_file_path
+
     @functools.cached_property
     def index_as_dataframe(self) -> pd.DataFrame:
         """Read and cache the full index file."""
-        if self.grib_source == "local" and wgrib2:
+        if self.idx_source is None and self.grib_source == "local" and wgrib2:
             # Generate IDX inventory with wgrib2
             self.idx = StringIO(wgrib2_idx(self.get_localFilePath()))
             self.idx_source = "generated"
@@ -625,6 +678,7 @@ def index_as_dataframe(self) -> pd.DataFrame:
                 read_this_idx = self.idx
             else:
                 read_this_idx = None
+                print(f"Downloading inventory file from {self.idx=}")
                 response = requests.get(self.idx)
                 if response.status_code != 200:
                     response.raise_for_status()
@@ -635,9 +689,18 @@ def index_as_dataframe(self) -> pd.DataFrame:
                         f"You will need to remake the Herbie object (H = `Herbie()`)\n"
                         f"or delete this cached property: `del H.index_as_dataframe()`"
                     )
+                
                 read_this_idx = StringIO(response.text)
                 response.close()
 
+                index_filepath = self.get_localIndexFilePath()
+                os.makedirs(os.path.dirname(index_filepath), exist_ok=True)
+
+                with open(index_filepath, "w") as file:
+                    file.write(read_this_idx.read())
+                    # reset the cursor to the beggining of the StringIO
+                    read_this_idx.seek(0)
+
             df = pd.read_csv(
                 read_this_idx,
                 sep=":",
@@ -697,9 +760,34 @@ def index_as_dataframe(self) -> pd.DataFrame:
             # eccodes keywords explained here:
             # https://confluence.ecmwf.int/display/UDOC/Identification+keywords
 
-            r = requests.get(self.idx)
-            idxs = [json.loads(x) for x in r.text.split("\n") if x]
-            r.close()
+            if self.idx_source in ["local"]:
+                with open(self.idx, "r") as file:
+                    read_this_idx = StringIO(file.read())
+            else:
+                print(f"Downloading inventory file from {self.idx=}")
+                response = requests.get(self.idx)
+                if response.status_code != 200:
+                    response.raise_for_status()
+                    response.close()
+                    raise ValueError(
+                        f"\nCant open index file {self.idx}\n"
+                        f"Download the full file first (with `H.download()`).\n"
+                        f"You will need to remake the Herbie object (H = `Herbie()`)\n"
+                        f"or delete this cached property: `del H.index_as_dataframe()`"
+                    )
+                
+                read_this_idx = StringIO(response.text)
+                response.close()
+
+                index_filepath = self.get_localIndexFilePath()
+                os.makedirs(os.path.dirname(index_filepath), exist_ok=True)
+
+                with open(index_filepath, "w") as file:
+                    file.write(read_this_idx.read())
+                    # reset the cursor to the beggining of the StringIO
+                    read_this_idx.seek(0)
+
+            idxs = [json.loads(x) for x in read_this_idx.getvalue().split("\n") if x]
             df = pd.DataFrame(idxs)
 
             # Format the DataFrame
@@ -1005,7 +1093,7 @@ def subset(search, outFile):
         if self.overwrite and self.grib_source.startswith("local"):
             # Search for the grib files on the remote archives again
             self.grib, self.grib_source = self.find_grib(overwrite=True)
-            self.idx, self.idx_source = self.find_idx()
+            self.idx, self.idx_source = self.find_idx(overwrite=True)
             print(f"Overwrite local file with file from [{self.grib_source}]")
 
         # Check that data exists

tests/test_local.py

@@ -0,0 +1,58 @@
+"""Test using local directory for data files."""
+
+from herbie import Herbie, config
+import pytest
+
+DATA_DIRECTORY = config["default"]["save_dir"] / "Herbie-Tests-Data/"
+HERBIE_OPTIONS = {
+    "verbose": True,
+    "overwrite": False,
+    "save_dir": DATA_DIRECTORY,
+}
+
+@pytest.fixture
+def setup_local_data(request):
+    date, model, product = request.param    
+    H = Herbie(date, 
+        model=model, 
+        product=product,
+        **HERBIE_OPTIONS
+    )
+
+    print(f"{H.save_dir=}")
+
+    H.download(verbose=True)
+    print(f"{H.idx_source=}")
+
+    print(H.inventory())
+    return date, model, product
+
+
+@pytest.mark.parametrize(
+    "setup_local_data",
+    [
+        ("2025-08-08", "gfs", "pgrb2.0p25"),
+        ("2020-10-27", "gfs", "0.5-degree"),
+        ("2025-08-10", "hrrr", "sfc"),
+        ("2025-08-01", "rap", "awp200"),
+        ("2025-08-05", "ifs", "wave"),
+    ],
+    indirect=True,
+    ids=[
+        "gfs-2025-08-08", 
+        "gfs-2020-10-27",
+        "hrrr-2025-08-10",
+        "rap-2025-08-01",
+        "ifs-2025-08-05",
+    ],
+)
+def test_local_data(setup_local_data):
+    date, model, product = setup_local_data
+    H = Herbie(date, 
+        model=model, 
+        product=product,
+        **HERBIE_OPTIONS
+    )
+
+    assert "local" == H.idx_source
+    assert "local" == H.grib_source