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", getattr(filename, "path", filename.__repr__()))
247            try:
248                self.size = self._filehandle.info().get("size", 0)
249            except (AttributeError, TypeError):
250                self.size = 0
251            self._fileid = hashlib.sha1((self.name + str(self.size)).encode("ASCII")).hexdigest()
252
253            if "r" in self.mode:
254                self.indexfile = f"{self.name}.grib2ioidx"
255                idx_loaded = False
256                if self.use_index:
257                    try:
258                        import fsspec
259
260                        with fsspec.open(self.indexfile, "rb") as f:
261                            self._index = pickle.load(f)
262                        idx_loaded = True
263                        self._hasindex = True
264                    except Exception:
265                        pass
266                if not idx_loaded:
267                    self._index = build_index(self._filehandle)
268                self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
269                self.messages = len(self._msgs)
270
271        else:
272            self.current_message = 0
273            self.name = str(filename)
274            is_remote = isinstance(filename, str) and "://" in filename
275            if is_remote:
276                import fsspec
277
278                self._filehandle = fsspec.open(filename, mode=mode, **kwargs).open()
279                self.name = filename
280                try:
281                    self.size = self._filehandle.info().get("size", 0)
282                except (AttributeError, TypeError):
283                    self.size = 0
284                self._fileid = hashlib.sha1((self.name + str(self.size)).encode("ASCII")).hexdigest()
285            else:
286                self.name = os.path.abspath(filename)
287                if "r" in mode:
288                    self._filehandle = builtins.open(filename, mode=mode)
289                    # Some GRIB2 files are gzipped, so check for that here, but
290                    # raise error when using xarray backend.
291                    # Gzip files contain a 2-byte header b'\x1f\x8b'.
292                    if self._filehandle.read(2) == _GZIP_HEADER:
293                        self._filehandle.close()
294                        if _xarray_backend:
295                            raise RuntimeError("Gzip GRIB2 files are not supported by the Xarray backend.")
296                        self._filehandle = gzip.open(filename, mode=mode)
297                    else:
298                        self._filehandle.seek(0)
299                else:
300                    self._filehandle = builtins.open(filename, mode=mode)
301                fstat = os.stat(self.name)
302                self.size = fstat.st_size
303                self._fileid = hashlib.sha1((self.name + str(fstat.st_ino) + str(self.size)).encode("ASCII")).hexdigest()
304
305            if "r" in self.mode:
306                index_paths = [f"{self.name}_{self._fileid}.grib2ioidx", f"{self.name}.grib2ioidx"]
307                idx_paths = [f"{self.name}.idx"]
308                idx_loaded = False
309                remote_index_cache_path = None
310                if is_remote:
311                    cache_root = os.path.join(os.path.expanduser("~"), ".cache", "grib2io")
312                    cache_key = hashlib.sha1((self.name + str(self.size)).encode("ASCII")).hexdigest()
313                    remote_index_cache_path = os.path.join(cache_root, f"{cache_key}.grib2ioidx")
314                if self.use_index:
315                    if is_remote and remote_index_cache_path and os.path.exists(remote_index_cache_path):
316                        try:
317                            with builtins.open(remote_index_cache_path, "rb") as f:
318                                self._index = pickle.load(f)
319                            self.indexfile = remote_index_cache_path
320                            idx_loaded = True
321                            self._hasindex = True
322                        except Exception:
323                            idx_loaded = False
324                    for idx_path in index_paths:
325                        if idx_loaded:
326                            break
327                        try:
328                            if is_remote:
329                                import fsspec
330
331                                with fsspec.open(idx_path, "rb", **kwargs) as f:
332                                    self._index = pickle.load(f)
333                            else:
334                                if os.path.exists(idx_path):
335                                    with builtins.open(idx_path, "rb") as f:
336                                        self._index = pickle.load(f)
337                                else:
338                                    continue
339                            self.indexfile = idx_path
340                            idx_loaded = True
341                            self._hasindex = True
342                            break
343                        except Exception:
344                            continue
345                    if not idx_loaded:
346                        for idx_path in idx_paths:
347                            try:
348                                if is_remote:
349                                    import fsspec
350
351                                    with fsspec.open(idx_path, "r", **kwargs) as f:
352                                        offsets = _parse_wgrib2_idx(f)
353                                else:
354                                    if os.path.exists(idx_path):
355                                        with builtins.open(idx_path, "r") as f:
356                                            offsets = _parse_wgrib2_idx(f)
357                                    else:
358                                        continue
359                                if offsets:
360                                    self._index = build_index(self._filehandle, offsets=offsets)
361                                    self.indexfile = idx_path
362                                    idx_loaded = True
363                                    self._hasindex = True
364                                    break
365                            except Exception:
366                                continue
367                    if is_remote and idx_loaded and self.save_index and remote_index_cache_path:
368                        if self.indexfile != remote_index_cache_path:
369                            try:
370                                os.makedirs(os.path.dirname(remote_index_cache_path), exist_ok=True)
371                                serialize_index(self._index, remote_index_cache_path)
372                            except Exception:
373                                pass
374                if not idx_loaded:
375                    self._index = build_index(self._filehandle)
376                    self.indexfile = index_paths[0]
377                    if self.save_index:
378                        try:
379                            if is_remote and remote_index_cache_path:
380                                os.makedirs(os.path.dirname(remote_index_cache_path), exist_ok=True)
381                                serialize_index(self._index, remote_index_cache_path)
382                                self.indexfile = remote_index_cache_path
383                            else:
384                                serialize_index(self._index, self.indexfile)
385                        except Exception as e:
386                            warnings.warn(f"index was not serialized for future use: {e}")
387                self._msgs = msgs_from_index(self._index, filehandle=self._filehandle)
388                self.messages = len(self._msgs)
389            elif "w" or "x" in self.mode:
390                self.messages = 0
391                self.current_message = None
392
393        self.closed = self._filehandle.closed
394
395    def __delete__(self, instance):
396        self.close()
397        del self._index
398
399    def __enter__(self):
400        return self
401
402    def __exit__(self, atype, value, traceback):
403        self.close()
404
405    def __iter__(self):
406        yield from self._msgs
407
408    def __len__(self):
409        return self.messages
410
411    def __repr__(self):
412        strings = []
413        for k in self.__slots__:
414            if k.startswith("_"):
415                continue
416            strings.append("%s = %s\n" % (k, eval("self." + k)))
417        return "".join(strings)
418
419    def __getitem__(self, key):
420        if isinstance(key, int):
421            if abs(key) >= len(self._msgs):
422                raise IndexError("index out of range")
423            else:
424                return self._msgs[key]
425        elif isinstance(key, str):
426            return self.select(shortName=key)
427        elif isinstance(key, slice):
428            return self._msgs[key]
429        elif isinstance(key, (list, tuple, set)):
430            if len(key) == 0:
431                return iter(self._msgs)
432            indices = sorted(key) if isinstance(key, set) else key
433
434            def _iter_msgs():
435                for i in indices:
436                    if not isinstance(i, int):
437                        raise TypeError(f"indices must be integers; got {type(i).__name__}")
438                    if abs(i) >= len(self._msgs):
439                        raise IndexError(f"index out of range: {i}")
440                    yield self._msgs[i]
441
442            return _iter_msgs()
443        else:
444            raise KeyError("Key must be an integer, slice, GRIB2 variable shortName, or an iterable of integer indices.")
445
446    @property
447    def levels(self):
448        """Provides a unique tuple of level strings."""
449        if self._hasindex:
450            return tuple(sorted(set([msg.level for msg in self._msgs])))
451        else:
452            return None
453
454    @property
455    def variables(self):
456        """Provides a unique tuple of variable shortName strings."""
457        if self._hasindex:
458            return tuple(sorted(set([msg.shortName for msg in self._msgs])))
459        else:
460            return None
461
462    def close(self):
463        """Close the file handle."""
464        if not self._filehandle.closed:
465            self.messages = 0
466            self.current_message = 0
467            self._filehandle.close()
468            self.closed = self._filehandle.closed
469
470    def read(self, size: Optional[int] = None):
471        """
472        Read size amount of GRIB2 messages from the current position.
473
474        If no argument is given, then size is None and all messages are returned
475        from the current position in the file. This read method follows the
476        behavior of Python's builtin open() function, but whereas that operates
477        on units of bytes, we operate on units of GRIB2 messages.
478
479        Parameters
480        ----------
481        size: default=None
482            The number of GRIB2 messages to read from the current position. If
483            no argument is give, the default value is None and remainder of
484            the file is read.
485
486        Returns
487        -------
488        read
489            ``Grib2Message`` object when size = 1 or a list of Grib2Messages
490            when size > 1.
491        """
492        if size is not None and size < 0:
493            size = None
494        if size is None or size > 1:
495            start = self.tell()
496            stop = self.messages if size is None else start + size
497            if size is None:
498                self.current_message = self.messages - 1
499            else:
500                self.current_message += size
501            return self._msgs[slice(start, stop, 1)]
502        elif size == 1:
503            self.current_message += 1
504            return self._msgs[self.current_message]
505        else:
506            None
507
508    def seek(self, pos: int):
509        """
510        Set the position within the file in units of GRIB2 messages.
511
512        Parameters
513        ----------
514        pos
515            The GRIB2 Message number to set the file pointer to.
516        """
517        if self._hasindex:
518            self._filehandle.seek(self._index["sectionOffset"][0][pos])
519            self.current_message = pos
520
521    def tell(self):
522        """Returns the position of the file in units of GRIB2 Messages."""
523        return self.current_message
524
525    def select(self, **kwargs):
526        """Select GRIB2 messages by `Grib2Message` attributes."""
527        # TODO: Added ability to process multiple values for each keyword (attribute)
528        idxs = []
529        nkeys = len(kwargs.keys())
530        for k, v in kwargs.items():
531            for m in self._msgs:
532                if hasattr(m, k) and getattr(m, k) == v:
533                    idxs.append(m._msgnum)
534        idxs = np.array(idxs, dtype=np.int32)
535        return [self._msgs[i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
536
537    def write(self, msg):
538        """
539        Writes GRIB2 message object to file.
540
541        Parameters
542        ----------
543        msg
544            GRIB2 message objects to write to file.
545        """
546        if isinstance(msg, list):
547            for m in msg:
548                self.write(m)
549            return
550
551        if issubclass(msg.__class__, _Grib2Message):
552            # TODO: We can consider letting pack return packed bytes instead of associating with message object
553            if hasattr(msg, "_msg"):
554                # write already packed bytes
555                self._filehandle.write(msg._msg)
556            else:
557                if msg._signature == msg._generate_signature() and msg._data is None and hasattr(msg._ondiskarray, "filehandle"):
558                    # write unchanged message from input
559                    offset = msg._ondiskarray.filehandle.tell()
560                    msg._ondiskarray.filehandle.seek(msg._ondiskarray.offset)
561                    self._filehandle.write(msg._ondiskarray.filehandle.read(msg.section0[-1]))
562                    msg._ondiskarray.filehandle.seek(offset)
563                else:
564                    msg.pack()
565                    self._filehandle.write(msg._msg)
566            self.size = os.path.getsize(self.name)
567            self.messages += 1
568        else:
569            raise TypeError("msg must be a Grib2Message object.")
570        return
571
572    def flush(self):
573        """Flush the file object buffer."""
574        self._filehandle.flush()
575
576    def levels_by_var(self, name: str):
577        """
578        Return a list of level strings given a variable shortName.
579
580        Parameters
581        ----------
582        name
583            Grib2Message variable shortName
584
585        Returns
586        -------
587        levels_by_var
588            A list of unique level strings.
589        """
590        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
591
592    def vars_by_level(self, level: str):
593        """
594        Return a list of variable shortName strings given a level.
595
596        Parameters
597        ----------
598        level
599            Grib2Message variable level
600
601        Returns
602        -------
603        vars_by_level
604            A list of unique variable shortName strings.
605        """
606        return list(sorted(set([msg.shortName for msg in self.select(level=level)])))