118class open:
119    """
120    GRIB2 File Object.
121
122    This class can accommodate a physical file with one or more GRIB2
123    messages or a bytes object containing a GRIB2 messages.
124
125    A physical file can contain one or more GRIB2 messages.  When instantiated,
126    class `grib2io.open`, the file named `filename` is opened for reading (`mode
127    = 'r'`) and is automatically indexed.  The indexing procedure reads some of
128    the GRIB2 metadata for all GRIB2 Messages.  A GRIB2 Message may contain
129    submessages whereby Section 2-7 can be repeated.  grib2io accommodates for
130    this by flattening any GRIB2 submessages into multiple individual messages.
131
132    It is important to note that GRIB2 files from some Meteorological agencies
133    contain other data than GRIB2 messages.  GRIB2 files from ECMWF can contain
134    GRIB1 and GRIB2 messages.  grib2io checks for these and safely ignores them.
135
136    Attributes
137    ----------
138    closed : bool
139        `True` is file handle is close; `False` otherwise.
140    current_message : int
141        Current position of the file in units of GRIB2 Messages. (read only)
142    indexfile : str
143        Index file for the GRIB2 file.
144    levels : tuple
145        Tuple containing a unique list of wgrib2-formatted level/layer strings.
146    messages : int
147        Count of GRIB2 Messages contained in the file.
148    mode : str
149        File IO mode of opening the file.
150    name : str
151        Full path name of the GRIB2 file.
152    save_index : bool
153        Whether to save a pickle-based index file for the GRIB2 file. Default is `True`.
154    size : int
155        Size of the file in units of bytes.
156    use_index
157        Whether to use an existing pickle-based index file for the GRIB2 file. Default is `True`.
158    variables : tuple
159        Tuple containing a unique list of variable short names (i.e. GRIB2
160        abbreviation names).
161    """
162
163    __slots__ = (
164        "_fileid",
165        "_filehandle",
166        "_hasindex",
167        "_index",
168        "_msgs",
169        "_pos",
170        "closed",
171        "current_message",
172        "indexfile",
173        "messages",
174        "mode",
175        "name",
176        "size",
177        "save_index",
178        "use_index",
179    )
180
181    def __init__(
182        self,
183        filename: Union[bytes, str, Path, IO],
184        mode: Literal["r", "w", "x"] = "r",
185        *,
186        save_index=True,
187        use_index=True,
188        _xarray_backend=False,
189        **kwargs,
190    ):
191        """
192        Initialize GRIB2 File object instance.
193
194        Parameters
195        ----------
196        filename
197            File path containing GRIB2 messages OR bytes OR file-like object.
198        mode
199            File access mode where "r" opens the files for reading only; "w"
200            opens the file for overwriting and "x" for writing to a new file.
201        save_index
202            Whether to save a pickle-based index file for the GRIB2 file. Default is True.
203
204            .. versionadded:: 2.6.0
205        use_index
206            Whether to use an existing pickle-based index file for the GRIB2 file. Default is True.
207
208            .. versionadded:: 2.6.0
209        """
210
211        # All write modes are read/write; all modes are binary.
212        if mode in ("a", "x", "w"):
213            mode += "+"
214        mode = mode + "b"
215
216        self._hasindex = False
217        self.indexfile = None
218        self.mode = mode
219        self.save_index = save_index
220        self.size = 0
221        self.use_index = use_index
222
223        if isinstance(filename, bytes):
224            if "r" not in self.mode:
225                raise ValueError(
226                    "Invalid mode for bytes input: GRIB2 data supplied as bytes is read-only. Use mode='r' or provide a filename instead."
227                )
228
229            self.current_message = 0
230            if filename[:2] == _GZIP_HEADER:
231                filename = gzip.decompress(filename)
232            if filename[:4] + filename[-4:] != b"GRIB7777":
233                raise ValueError("Invalid GRIB bytes")
234            self._filehandle = BytesIO(filename)
235            self.name = "<in-memory-file>"
236            self.size = len(filename)
237            self._fileid = hashlib.sha1((self.name + str(self.size)).encode("ASCII")).hexdigest()
238            self._index = build_index(self._filehandle)
239            self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
240            self.messages = len(self._msgs)
241
242        elif all(hasattr(filename, attr) for attr in ("read", "seek", "tell")):
243            # Handle file-like objects (including S3File, etc.)
244            self.current_message = 0
245            self._filehandle = filename
246            self.name = getattr(filename, "name", filename.__repr__())
247            # For file-like objects, we can't get file stats, so create a unique ID
248            self._fileid = hashlib.sha1((self.name + str(id(filename))).encode("ASCII")).hexdigest()
249            # Disable index file usage for file-like objects since we can't save/load them
250            self.use_index = False
251            self.save_index = False
252
253            if "r" in self.mode:
254                # Build index from the file-like object
255                self._index = build_index(self._filehandle)
256                self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
257                self.messages = len(self._msgs)
258
259            # Get size for s3fs.core.S3File object.
260            self.size = self._filehandle.info()["size"] or 0
261
262        else:
263            self.current_message = 0
264            if "r" in mode:
265                self._filehandle = builtins.open(filename, mode=mode)
266                # Some GRIB2 files are gzipped, so check for that here, but
267                # raise error when using xarray backend.
268                # Gzip files contain a 2-byte header b'\x1f\x8b'.
269                if self._filehandle.read(2) == _GZIP_HEADER:
270                    self._filehandle.close()
271                    if _xarray_backend:
272                        raise RuntimeError("Gzip GRIB2 files are not supported by the Xarray backend.")
273                    self._filehandle = gzip.open(filename, mode=mode)
274                else:
275                    self._filehandle = builtins.open(filename, mode=mode)
276            else:
277                self._filehandle = builtins.open(filename, mode=mode)
278                self.name = os.path.abspath(filename)
279            self.name = os.path.abspath(filename)
280            fstat = os.stat(self.name)
281            self.size = fstat.st_size
282            self._fileid = hashlib.sha1((self.name + str(fstat.st_ino) + str(self.size)).encode("ASCII")).hexdigest()
283
284            if "r" in self.mode:
285                self.indexfile = f"{self.name}_{self._fileid}.grib2ioidx"
286                if self.use_index and os.path.exists(self.indexfile):
287                    try:
288                        with builtins.open(self.indexfile, "rb") as file:
289                            index = pickle.load(file)
290                        self._index = index
291                    except Exception as e:
292                        warnings.warn(
293                            f"found indexfile: {self.indexfile}, but unable to load it: {e}\n"
294                            f"re-forming index from grib2file, but not writing indexfile"
295                        )
296                        self._index = build_index(self._filehandle)
297                    self._hasindex = True
298                else:
299                    self._index = build_index(self._filehandle)
300                    if self.save_index:
301                        try:
302                            serialize_index(self._index, self.indexfile)
303                        except Exception as e:
304                            warnings.warn(f"index was not serialized for future use: {e}")
305
306                self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
307
308                self.messages = len(self._msgs)
309            elif "w" or "x" in self.mode:
310                self.messages = 0
311                self.current_message = None
312
313        self.closed = self._filehandle.closed
314
315    def __delete__(self, instance):
316        self.close()
317        del self._index
318
319    def __enter__(self):
320        return self
321
322    def __exit__(self, atype, value, traceback):
323        self.close()
324
325    def __iter__(self):
326        yield from self._msgs
327
328    def __len__(self):
329        return self.messages
330
331    def __repr__(self):
332        strings = []
333        for k in self.__slots__:
334            if k.startswith("_"):
335                continue
336            strings.append("%s = %s\n" % (k, eval("self." + k)))
337        return "".join(strings)
338
339    def __getitem__(self, key):
340        if isinstance(key, int):
341            if abs(key) >= len(self._msgs):
342                raise IndexError("index out of range")
343            else:
344                return self._msgs[key]
345        elif isinstance(key, str):
346            return self.select(shortName=key)
347        elif isinstance(key, slice):
348            return self._msgs[key]
349        elif isinstance(key, (list, tuple, set)):
350            if len(key) == 0:
351                return iter(self._msgs)
352            indices = sorted(key) if isinstance(key, set) else key
353
354            def _iter_msgs():
355                for i in indices:
356                    if not isinstance(i, int):
357                        raise TypeError(f"indices must be integers; got {type(i).__name__}")
358                    if abs(i) >= len(self._msgs):
359                        raise IndexError(f"index out of range: {i}")
360                    yield self._msgs[i]
361
362            return _iter_msgs()
363        else:
364            raise KeyError("Key must be an integer, slice, GRIB2 variable shortName, or an iterable of integer indices.")
365
366    @property
367    def levels(self):
368        """Provides a unique tuple of level strings."""
369        if self._hasindex:
370            return tuple(sorted(set([msg.level for msg in self._msgs])))
371        else:
372            return None
373
374    @property
375    def variables(self):
376        """Provides a unique tuple of variable shortName strings."""
377        if self._hasindex:
378            return tuple(sorted(set([msg.shortName for msg in self._msgs])))
379        else:
380            return None
381
382    def close(self):
383        """Close the file handle."""
384        if not self._filehandle.closed:
385            self.messages = 0
386            self.current_message = 0
387            self._filehandle.close()
388            self.closed = self._filehandle.closed
389
390    def read(self, size: Optional[int] = None):
391        """
392        Read size amount of GRIB2 messages from the current position.
393
394        If no argument is given, then size is None and all messages are returned
395        from the current position in the file. This read method follows the
396        behavior of Python's builtin open() function, but whereas that operates
397        on units of bytes, we operate on units of GRIB2 messages.
398
399        Parameters
400        ----------
401        size: default=None
402            The number of GRIB2 messages to read from the current position. If
403            no argument is give, the default value is None and remainder of
404            the file is read.
405
406        Returns
407        -------
408        read
409            ``Grib2Message`` object when size = 1 or a list of Grib2Messages
410            when size > 1.
411        """
412        if size is not None and size < 0:
413            size = None
414        if size is None or size > 1:
415            start = self.tell()
416            stop = self.messages if size is None else start + size
417            if size is None:
418                self.current_message = self.messages - 1
419            else:
420                self.current_message += size
421            return self._msgs[slice(start, stop, 1)]
422        elif size == 1:
423            self.current_message += 1
424            return self._msgs[self.current_message]
425        else:
426            None
427
428    def seek(self, pos: int):
429        """
430        Set the position within the file in units of GRIB2 messages.
431
432        Parameters
433        ----------
434        pos
435            The GRIB2 Message number to set the file pointer to.
436        """
437        if self._hasindex:
438            self._filehandle.seek(self._index["sectionOffset"][0][pos])
439            self.current_message = pos
440
441    def tell(self):
442        """Returns the position of the file in units of GRIB2 Messages."""
443        return self.current_message
444
445    def select(self, **kwargs):
446        """Select GRIB2 messages by `Grib2Message` attributes."""
447        # TODO: Added ability to process multiple values for each keyword (attribute)
448        idxs = []
449        nkeys = len(kwargs.keys())
450        for k, v in kwargs.items():
451            for m in self._msgs:
452                if hasattr(m, k) and getattr(m, k) == v:
453                    idxs.append(m._msgnum)
454        idxs = np.array(idxs, dtype=np.int32)
455        return [self._msgs[i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
456
457    def write(self, msg):
458        """
459        Writes GRIB2 message object to file.
460
461        Parameters
462        ----------
463        msg
464            GRIB2 message objects to write to file.
465        """
466        if isinstance(msg, list):
467            for m in msg:
468                self.write(m)
469            return
470
471        if issubclass(msg.__class__, _Grib2Message):
472            # TODO: We can consider letting pack return packed bytes instead of associating with message object
473            if hasattr(msg, "_msg"):
474                # write already packed bytes
475                self._filehandle.write(msg._msg)
476            else:
477                if msg._signature == msg._generate_signature() and msg._data is None and hasattr(msg._ondiskarray, "filehandle"):
478                    # write unchanged message from input
479                    offset = msg._ondiskarray.filehandle.tell()
480                    msg._ondiskarray.filehandle.seek(msg._ondiskarray.offset)
481                    self._filehandle.write(msg._ondiskarray.filehandle.read(msg.section0[-1]))
482                    msg._ondiskarray.filehandle.seek(offset)
483                else:
484                    msg.pack()
485                    self._filehandle.write(msg._msg)
486            self.size = os.path.getsize(self.name)
487            self.messages += 1
488        else:
489            raise TypeError("msg must be a Grib2Message object.")
490        return
491
492    def flush(self):
493        """Flush the file object buffer."""
494        self._filehandle.flush()
495
496    def levels_by_var(self, name: str):
497        """
498        Return a list of level strings given a variable shortName.
499
500        Parameters
501        ----------
502        name
503            Grib2Message variable shortName
504
505        Returns
506        -------
507        levels_by_var
508            A list of unique level strings.
509        """
510        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
511
512    def vars_by_level(self, level: str):
513        """
514        Return a list of variable shortName strings given a level.
515
516        Parameters
517        ----------
518        level
519            Grib2Message variable level
520
521        Returns
522        -------
523        vars_by_level
524            A list of unique variable shortName strings.
525        """
526        return list(sorted(set([msg.shortName for msg in self.select(level=level)])))

79def show_config():
80    """Print grib2io build configuration information."""
81    print(f"grib2io version {__version__} Configuration:")
82    print("")
83    print(f"NCEPLIBS-g2c library version: {__g2clib_version__}")
84    print(f"\tStatic library: {g2c_static}")
85    print(f"\tJPEG compression support: {has_jpeg_support}")
86    print(f"\tPNG compression support: {has_png_support}")
87    print(f"\tAEC compression support: {has_aec_support}")
88    print("")
89    print(f"NCEPLIPS-ip support: {has_interpolation}")
90    print(f"\tStatic library: {ip_static}")
91    print(f"\tOpenMP support: {has_openmp_support}")
92    print("")
93    print("Static libs:")
94    for lib in extra_objects:
95        print(f"\t{lib}")
96    print("")
97    print(f"NCEP GRIB2 Table Version: {_ncep_grib2_table_version}")

2115def interpolate(
2116    a,
2117    method: Union[int, str],
2118    grid_def_in,
2119    grid_def_out,
2120    method_options=None,
2121    num_threads=1,
2122):
2123    """
2124    This is the module-level interpolation function.
2125
2126    This interfaces with the "d" version of[NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
2127    through grib2io's internal iplib Cython extension module. The "d" version
2128    defines 4-byte integers and 8-byte reals.
2129
2130    Parameters
2131    ----------
2132    a : numpy.ndarray or tuple
2133        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
2134        performed.  If `a` is a `tuple`, then vector interpolation will be
2135        performed with the assumption that u = a[0] and v = a[1] and are both
2136        `numpy.ndarray`.
2137
2138        These data are expected to be in 2-dimensional form with shape (ny, nx)
2139        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
2140        spatial, temporal, or classification (i.e. ensemble members) dimension.
2141        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
2142        acceptable for input into the interpolation subroutines. If needed, these
2143        data will be converted to `np.float32`.
2144    method
2145        Interpolate method to use. This can either be an integer or string using
2146        the following mapping:
2147
2148        | Interpolate Scheme | Integer Value |
2149        | :---:              | :---:         |
2150        | 'bilinear'         | 0             |
2151        | 'bicubic'          | 1             |
2152        | 'neighbor'         | 2             |
2153        | 'budget'           | 3             |
2154        | 'spectral'         | 4             |
2155        | 'neighbor-budget'  | 6             |
2156
2157    grid_def_in : grib2io.Grib2GridDef
2158        Grib2GridDef object for the input grid.
2159    grid_def_out : grib2io.Grib2GridDef
2160        Grib2GridDef object for the output grid or station points.
2161    method_options : list of ints, optional
2162        Interpolation options. See the NCEPLIBS-ip documentation for
2163        more information on how these are used.
2164    num_threads : int, optional
2165        Number of OpenMP threads to use for interpolation. The default
2166        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
2167        was not built with OpenMP, then this keyword argument and value
2168        will have no impact.
2169
2170    Returns
2171    -------
2172    interpolate
2173        Returns a `numpy.ndarray` of dtype `np.float32` when scalar interpolation
2174        is performed or a `tuple` of `numpy.ndarray`s when vector interpolation is
2175        performed with the assumptions that 0-index is the interpolated u and
2176        1-index is the interpolated v.
2177    """
2178
2179    try:
2180        from . import iplib
2181    except ImportError:
2182        raise ImportError("NCEPLIBS-ip library not found. Interpolation is not available.")
2183
2184    prev_num_threads = 1
2185    try:
2186        prev_num_threads = iplib.openmp_get_num_threads()
2187        iplib.openmp_set_num_threads(num_threads)
2188    except AttributeError:
2189        pass
2190
2191    print(f"grib2io.interpolate thread report: OpenMP num threads = {iplib.openmp_get_num_threads()}")
2192
2193    if isinstance(method, int) and method not in _interp_schemes.values():
2194        raise ValueError("Invalid interpolation method.")
2195    elif isinstance(method, str):
2196        if method in _interp_schemes.keys():
2197            method = _interp_schemes[method]
2198        else:
2199            raise ValueError("Invalid interpolation method.")
2200
2201    if method_options is None:
2202        method_options = np.zeros((20), dtype=np.int32)
2203        if method in {3, 6}:
2204            method_options[0:2] = -1
2205
2206    mi = grid_def_in.npoints
2207    mo = grid_def_out.npoints
2208
2209    # Adjust shape of input array(s)
2210    a, newshp = _adjust_array_shape_for_interp(a, grid_def_in, grid_def_out)
2211
2212    # Call interpolation subroutines according to type of a.
2213    if isinstance(a, np.ndarray):
2214        # Scalar
2215        km = a.shape[0]
2216        if np.any(np.isnan(a)):
2217            ibi = np.ones((km), dtype=np.int32)
2218            li = np.where(np.isnan(a), 0, 1).astype(np.uint8)
2219        else:
2220            ibi = np.zeros((km), dtype=np.int32)
2221            li = np.zeros(a.shape, dtype=np.uint8)
2222        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
2223            method,
2224            method_options,
2225            grid_def_in.gdtn,
2226            np.array(grid_def_in.gdt, dtype=np.int32),
2227            grid_def_out.gdtn,
2228            np.array(grid_def_out.gdt, dtype=np.int32),
2229            mi,
2230            mo,
2231            km,
2232            ibi,
2233            li,
2234            a.astype(np.float64),
2235        )
2236        out = np.where(lo == 0, np.nan, go).reshape(newshp)
2237    elif isinstance(a, tuple):
2238        # Vector
2239        km = a[0].shape[0]
2240        if np.any(np.isnan(a)):
2241            ibi = np.ones((km), dtype=np.int32)
2242            li = np.where(np.isnan(a), 0, 1).astype(np.uint8)
2243        else:
2244            ibi = np.zeros((km), dtype=np.int32)
2245            li = np.zeros(a[0].shape, dtype=np.uint8)
2246        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
2247            method,
2248            method_options,
2249            grid_def_in.gdtn,
2250            np.array(grid_def_in.gdt, dtype=np.int32),
2251            grid_def_out.gdtn,
2252            np.array(grid_def_out.gdt, dtype=np.int32),
2253            mi,
2254            mo,
2255            km,
2256            ibi,
2257            li,
2258            a[0].astype(np.float64),
2259            a[1].astype(np.float64),
2260        )
2261        uo = np.where(lo == 0, np.nan, uo).reshape(newshp)
2262        vo = np.where(lo == 0, np.nan, vo).reshape(newshp)
2263        out = (uo.astype(np.float32), vo.astype(np.float32))
2264
2265    try:
2266        iplib.openmp_set_num_threads(prev_num_threads)
2267    except AttributeError:
2268        pass
2269
2270    return out