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

56def show_config():
57    """Print grib2io build configuration information."""
58    print(f"grib2io version {__version__} Configuration:")
59    print("")
60    print(f"NCEPLIBS-g2c library version: {__g2clib_version__}")
61    print(f"\tStatic library: {g2c_static}")
62    print(f"\tJPEG compression support: {has_jpeg_support}")
63    print(f"\tPNG compression support: {has_png_support}")
64    print(f"\tAEC compression support: {has_aec_support}")
65    print("")
66    print(f"NCEPLIPS-ip support: {has_interpolation}")
67    print(f"\tStatic library: {ip_static}")
68    print(f"\tOpenMP support: {has_openmp_support}")
69    print("")
70    print("Static libs:")
71    for lib in extra_objects:
72        print(f"\t{lib}")
73    print("")
74    print(f"NCEP GRIB2 Table Version: {_ncep_grib2_table_version}")

2231def interpolate(
2232    a,
2233    method: Union[int, str],
2234    grid_def_in,
2235    grid_def_out,
2236    method_options=None,
2237    num_threads=1,
2238):
2239    """
2240    This is the module-level interpolation function.
2241
2242    This interfaces with the 4-byte (32-bit float) [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip)
2243    through grib2io's internal iplib Cython extension module.
2244
2245    Parameters
2246    ----------
2247    a : numpy.ndarray or tuple
2248        Input data.  If `a` is a `numpy.ndarray`, scalar interpolation will be
2249        performed.  If `a` is a `tuple`, then vector interpolation will be
2250        performed with the assumption that u = a[0] and v = a[1] and are both
2251        `numpy.ndarray`.
2252
2253        These data are expected to be in 2-dimensional form with shape (ny, nx)
2254        or 3-dimensional (:, ny, nx) where the 1st dimension represents another
2255        spatial, temporal, or classification (i.e. ensemble members) dimension.
2256        The function will properly flatten the (ny,nx) dimensions into (nx * ny)
2257        acceptable for input into the interpolation subroutines. If needed, these
2258        data will be converted to `np.float32`.
2259    method
2260        Interpolate method to use. This can either be an integer or string using
2261        the following mapping:
2262
2263        | Interpolate Scheme | Integer Value |
2264        | :---:              | :---:         |
2265        | 'bilinear'         | 0             |
2266        | 'bicubic'          | 1             |
2267        | 'neighbor'         | 2             |
2268        | 'budget'           | 3             |
2269        | 'spectral'         | 4             |
2270        | 'neighbor-budget'  | 6             |
2271
2272    grid_def_in : grib2io.Grib2GridDef
2273        Grib2GridDef object for the input grid.
2274    grid_def_out : grib2io.Grib2GridDef
2275        Grib2GridDef object for the output grid or station points.
2276    method_options : list of ints, optional
2277        Interpolation options. See the NCEPLIBS-ip documentation for
2278        more information on how these are used.
2279    num_threads : int, optional
2280        Number of OpenMP threads to use for interpolation. The default
2281        value is 1. If NCEPLIBS-ip and grib2io's iplib extension module
2282        was not built with OpenMP, then this keyword argument and value
2283        will have no impact.
2284
2285    Returns
2286    -------
2287    interpolate
2288        Returns a `numpy.ndarray` of dtype `np.float32` when scalar interpolation
2289        is performed or a `tuple` of `numpy.ndarray`s when vector interpolation is
2290        performed with the assumptions that 0-index is the interpolated u and
2291        1-index is the interpolated v.
2292    """
2293
2294    try:
2295        from . import iplib
2296    except ImportError:
2297        raise ImportError(
2298            "NCEPLIBS-ip library not found. Interpolation is not available."
2299        )
2300
2301    prev_num_threads = 1
2302    try:
2303        prev_num_threads = iplib.openmp_get_num_threads()
2304        iplib.openmp_set_num_threads(num_threads)
2305    except AttributeError:
2306        pass
2307
2308    print(
2309        f"grib2io.interpolate thread report: OpenMP num threads = {iplib.openmp_get_num_threads()}"
2310    )
2311
2312    if isinstance(method, int) and method not in _interp_schemes.values():
2313        raise ValueError("Invalid interpolation method.")
2314    elif isinstance(method, str):
2315        if method in _interp_schemes.keys():
2316            method = _interp_schemes[method]
2317        else:
2318            raise ValueError("Invalid interpolation method.")
2319
2320    if method_options is None:
2321        method_options = np.zeros((20), dtype=np.int32)
2322        if method in {3, 6}:
2323            method_options[0:2] = -1
2324
2325    mi = grid_def_in.npoints
2326    mo = grid_def_out.npoints
2327
2328    # Adjust shape of input array(s)
2329    a, newshp = _adjust_array_shape_for_interp(a, grid_def_in, grid_def_out)
2330
2331    # Call interpolation subroutines according to type of a.
2332    if isinstance(a, np.ndarray):
2333        # Scalar
2334        km = a.shape[0]
2335        if np.any(np.isnan(a)):
2336            ibi = np.ones((km), dtype=np.int32)
2337            li = np.where(np.isnan(a), 0, 1).astype(np.uint8)
2338        else:
2339            ibi = np.zeros((km), dtype=np.int32)
2340            li = np.zeros(a.shape, dtype=np.uint8)
2341        no, rlat, rlon, ibo, lo, go, iret = iplib.interpolate_scalar(
2342            method,
2343            method_options,
2344            grid_def_in.gdtn,
2345            np.array(grid_def_in.gdt, dtype=np.int32),
2346            grid_def_out.gdtn,
2347            np.array(grid_def_out.gdt, dtype=np.int32),
2348            mi,
2349            mo,
2350            km,
2351            ibi,
2352            li,
2353            a,
2354        )
2355        out = np.where(lo == 0, np.nan, go).reshape(newshp)
2356    elif isinstance(a, tuple):
2357        # Vector
2358        km = a[0].shape[0]
2359        if np.any(np.isnan(a)):
2360            ibi = np.ones((km), dtype=np.int32)
2361            li = np.where(np.isnan(a), 0, 1).astype(np.uint8)
2362        else:
2363            ibi = np.zeros((km), dtype=np.int32)
2364            li = np.zeros(a[0].shape, dtype=np.uint8)
2365        no, rlat, rlon, crot, srot, ibo, lo, uo, vo, iret = iplib.interpolate_vector(
2366            method,
2367            method_options,
2368            grid_def_in.gdtn,
2369            np.array(grid_def_in.gdt, dtype=np.int32),
2370            grid_def_out.gdtn,
2371            np.array(grid_def_out.gdt, dtype=np.int32),
2372            mi,
2373            mo,
2374            km,
2375            ibi,
2376            li,
2377            a[0],
2378            a[1],
2379        )
2380        uo = np.where(lo == 0, np.nan, uo).reshape(newshp)
2381        vo = np.where(lo == 0, np.nan, vo).reshape(newshp)
2382        out = (uo, vo)
2383
2384    try:
2385        iplib.openmp_set_num_threads(prev_num_threads)
2386    except AttributeError:
2387        pass
2388
2389    return out