60class open():
 61    """
 62    GRIB2 File Object.
 63
 64    A physical file can contain one or more GRIB2 messages.  When instantiated, class `grib2io.open`, 
 65    the file named `filename` is opened for reading (`mode = 'r'`) and is automatically indexed.  
 66    The indexing procedure reads some of the GRIB2 metadata for all GRIB2 Messages.  A GRIB2 Message 
 67    may contain submessages whereby Section 2-7 can be repeated.  grib2io accommodates for this by 
 68    flattening any GRIB2 submessages into multiple individual messages.
 69
 70    Attributes
 71    ----------
 72    **`mode : str`**
 73        File IO mode of opening the file.
 74
 75    **`name : str`**
 76        Full path name of the GRIB2 file.
 77
 78    **`messages : int`**
 79        Count of GRIB2 Messages contained in the file.
 80
 81    **`current_message : int`**
 82        Current position of the file in units of GRIB2 Messages.
 83
 84    **`size : int`**
 85        Size of the file in units of bytes.
 86
 87    **`closed : bool`**
 88        `True` is file handle is close; `False` otherwise.
 89
 90    **`variables : tuple`**
 91        Tuple containing a unique list of variable short names (i.e. GRIB2 abbreviation names).
 92
 93    **`levels : tuple`**
 94        Tuple containing a unique list of wgrib2-formatted level/layer strings.
 95    """
 96    __slots__ = ('_filehandle','_hasindex','_index','mode','name','messages',
 97                 'current_message','size','closed','variables','levels','_pos')
 98    def __init__(self, filename, mode='r', **kwargs):
 99        """
100        Initialize GRIB2 File object instance.
101
102        Parameters
103        ----------
104
105        **`filename : str`**
106            File name containing GRIB2 messages.
107
108        **`mode : str, optional`**
109            File access mode where `r` opens the files for reading only; `w` opens the file for writing.
110        """
111        if mode in {'a','r','w'}:
112            mode = mode+'b'
113            if 'w' in mode: mode += '+'
114            if 'a' in mode: mode += '+'
115        self._filehandle = builtins.open(filename,mode=mode,buffering=ONE_MB)
116        self._hasindex = False
117        self._index = {}
118        self.mode = mode
119        self.name = os.path.abspath(filename)
120        self.messages = 0
121        self.current_message = 0
122        self.size = os.path.getsize(self.name)
123        self.closed = self._filehandle.closed
124        self.levels = None
125        self.variables = None
126        if 'r' in self.mode:
127            try:
128                self._build_index(no_data=kwargs['_xarray_backend'])
129            except(KeyError):
130                self._build_index()
131        # FIX: Cannot perform reads on mode='a'
132        #if 'a' in self.mode and self.size > 0: self._build_index()
133
134
135    def __delete__(self, instance):
136        self.close()
137        del self._index
138
139
140    def __enter__(self):
141        return self
142
143
144    def __exit__(self, atype, value, traceback):
145        self.close()
146
147
148    def __iter__(self):
149        yield from self._index['msg']
150
151
152    def __len__(self):
153        return self.messages
154
155
156    def __repr__(self):
157        strings = []
158        for k in self.__slots__:
159            if k.startswith('_'): continue
160            strings.append('%s = %s\n'%(k,eval('self.'+k)))
161        return ''.join(strings)
162
163
164    def __getitem__(self, key):
165        if isinstance(key,int):
166            if abs(key) >= len(self._index['msg']):
167                raise IndexError("index out of range")
168            else:
169                return self._index['msg'][key]
170        elif isinstance(key,str):
171            return self.select(shortName=key)
172        elif isinstance(key,slice):
173            return self._index['msg'][key]
174        else:
175            raise KeyError('Key must be an integer, slice, or GRIB2 variable shortName.')
176
177
178    def _build_index(self, no_data=False):
179        """
180        Perform indexing of GRIB2 Messages.
181        """
182        # Initialize index dictionary
183        if not self._hasindex:
184            self._index['offset'] = []
185            self._index['bitmap_offset'] = []
186            self._index['data_offset'] = []
187            self._index['size'] = []
188            self._index['data_size'] = []
189            self._index['submessageOffset'] = []
190            self._index['submessageBeginSection'] = []
191            self._index['isSubmessage'] = []
192            self._index['messageNumber'] = []
193            self._index['msg'] = []
194            self._hasindex = True
195
196        # Iterate
197        while True:
198            try:
199                # Read first 4 bytes and decode...looking for "GRIB"
200                pos = self._filehandle.tell()
201
202                # Ignore headers (usually text) that are not part of the GRIB2
203                # file.  For example, NAVGEM files have a http header at the
204                # beginning that needs to be ignored.
205
206                # Read a byte at a time until "GRIB" is found.  Using
207                # "wgrib2" on a NAVGEM file, the header was 421 bytes and
208                # decided to go to 2048 bytes to be safe. For normal GRIB2
209                # files this should be quick and break out of the first
210                # loop when "test_offset" is 0.
211                for test_offset in range(2048):
212                    self._filehandle.seek(pos + test_offset)
213                    header = struct.unpack(">i", self._filehandle.read(4))[0]
214                    if header.to_bytes(4, "big") == b"GRIB":
215                        pos = pos + test_offset
216                        break
217
218                # Test header. Then get information from GRIB2 Section 0: the discipline
219                # number, edition number (should always be 2), and GRIB2 message size.
220                # Then iterate to check for submessages.
221                if header.to_bytes(4,'big') == b'GRIB':
222
223                    _issubmessage = False
224                    _submsgoffset = 0
225                    _submsgbegin = 0
226                    _bmapflag = None
227
228                    # Read the rest of Section 0 using struct.
229                    section0 = np.concatenate(([header],list(struct.unpack('>HBBQ',self._filehandle.read(12)))),dtype=np.int64)
230                    assert section0[3] == 2
231
232                    # Read and unpack Section 1
233                    secsize = struct.unpack('>i',self._filehandle.read(4))[0]
234                    secnum = struct.unpack('>B',self._filehandle.read(1))[0]
235                    assert secnum == 1
236                    self._filehandle.seek(self._filehandle.tell()-5)
237                    _grbmsg = self._filehandle.read(secsize)
238                    _grbpos = 0
239                    section1,_grbpos = g2clib.unpack1(_grbmsg,_grbpos,np.empty)
240                    secrange = range(2,8)
241                    while 1:
242                        section2 = b''
243                        for num in secrange:
244                            secsize = struct.unpack('>i',self._filehandle.read(4))[0]
245                            secnum = struct.unpack('>B',self._filehandle.read(1))[0]
246                            if secnum == num:
247                                if secnum == 2:
248                                    if secsize > 0:
249                                        section2 = self._filehandle.read(secsize-5)
250                                elif secnum == 3:
251                                    self._filehandle.seek(self._filehandle.tell()-5)
252                                    _grbmsg = self._filehandle.read(secsize)
253                                    _grbpos = 0
254                                    # Unpack Section 3
255                                    _gds,_gdt,_deflist,_grbpos = g2clib.unpack3(_grbmsg,_grbpos,np.empty)
256                                    _gds = _gds.tolist()
257                                    _gdt = _gdt.tolist()
258                                    section3 = np.concatenate((_gds,_gdt))
259                                    section3 = np.where(section3==4294967295,-1,section3)
260                                elif secnum == 4:
261                                    self._filehandle.seek(self._filehandle.tell()-5)
262                                    _grbmsg = self._filehandle.read(secsize)
263                                    _grbpos = 0
264                                    # Unpack Section 4
265                                    _numcoord,_pdt,_pdtnum,_coordlist,_grbpos = g2clib.unpack4(_grbmsg,_grbpos,np.empty)
266                                    _pdt = _pdt.tolist()
267                                    section4 = np.concatenate((np.array((_numcoord,_pdtnum)),_pdt))
268                                elif secnum == 5:
269                                    self._filehandle.seek(self._filehandle.tell()-5)
270                                    _grbmsg = self._filehandle.read(secsize)
271                                    _grbpos = 0
272                                    # Unpack Section 5
273                                    _drt,_drtn,_npts,self._pos = g2clib.unpack5(_grbmsg,_grbpos,np.empty)
274                                    section5 = np.concatenate((np.array((_npts,_drtn)),_drt))
275                                    section5 = np.where(section5==4294967295,-1,section5)
276                                elif secnum == 6:
277                                    # Unpack Section 6. Not really...just get the flag value.
278                                    _bmapflag = struct.unpack('>B',self._filehandle.read(1))[0]
279                                    if _bmapflag == 0:
280                                        _bmappos = self._filehandle.tell()-6
281                                    elif _bmapflag == 254:
282                                        pass # Do this to keep the previous position value
283                                    else:
284                                        _bmappos = None
285                                    self._filehandle.seek(self._filehandle.tell()+secsize-6)
286                                elif secnum == 7:
287                                    # Unpack Section 7. No need to read it, just index the position in file.
288                                    _datapos = self._filehandle.tell()-5
289                                    _datasize = secsize
290                                    self._filehandle.seek(self._filehandle.tell()+secsize-5)
291                                else:
292                                    self._filehandle.seek(self._filehandle.tell()+secsize-5)
293                            else:
294                                if num == 2 and secnum == 3:
295                                    pass # Allow this.  Just means no Local Use Section.
296                                else:
297                                    _issubmessage = True
298                                    _submsgoffset = (self._filehandle.tell()-5)-(self._index['offset'][-1])
299                                    _submsgbegin = secnum
300                                self._filehandle.seek(self._filehandle.tell()-5)
301                                continue
302                        trailer = struct.unpack('>4s',self._filehandle.read(4))[0]
303                        if trailer == b'7777':
304                            self.messages += 1
305                            self._index['offset'].append(pos)
306                            self._index['bitmap_offset'].append(_bmappos)
307                            self._index['data_offset'].append(_datapos)
308                            self._index['size'].append(section0[-1])
309                            self._index['data_size'].append(_datasize)
310                            self._index['messageNumber'].append(self.messages)
311                            self._index['isSubmessage'].append(_issubmessage)
312                            if _issubmessage:
313                                self._index['submessageOffset'].append(_submsgoffset)
314                                self._index['submessageBeginSection'].append(_submsgbegin)
315                            else:
316                                self._index['submessageOffset'].append(0)
317                                self._index['submessageBeginSection'].append(_submsgbegin)
318
319                            # Create Grib2Message with data.
320                            msg = Grib2Message(section0,section1,section2,section3,section4,section5,_bmapflag)
321                            msg._msgnum = self.messages-1
322                            msg._deflist = _deflist
323                            msg._coordlist = _coordlist
324                            if not no_data:
325                                msg._data = Grib2MessageOnDiskArray((msg.ny,msg.nx), 2,
326                                                                    TYPE_OF_VALUES_DTYPE[msg.typeOfValues],
327                                                                    self._filehandle,
328                                                                    msg, pos, _bmappos, _datapos)
329                            self._index['msg'].append(msg)
330
331                            break
332                        else:
333                            self._filehandle.seek(self._filehandle.tell()-4)
334                            self.messages += 1
335                            self._index['offset'].append(pos)
336                            self._index['bitmap_offset'].append(_bmappos)
337                            self._index['data_offset'].append(_datapos)
338                            self._index['size'].append(section0[-1])
339                            self._index['data_size'].append(_datasize)
340                            self._index['messageNumber'].append(self.messages)
341                            self._index['isSubmessage'].append(_issubmessage)
342                            self._index['submessageOffset'].append(_submsgoffset)
343                            self._index['submessageBeginSection'].append(_submsgbegin)
344
345                            # Create Grib2Message with data.
346                            msg = Grib2Message(section0,section1,section2,section3,section4,section5,_bmapflag)
347                            msg._msgnum = self.messages-1
348                            msg._deflist = _deflist
349                            msg._coordlist = _coordlist
350                            if not no_data:
351                                msg._data = Grib2MessageOnDiskArray((msg.ny,msg.nx), 2,
352                                                                    TYPE_OF_VALUES_DTYPE[msg.typeOfValues],
353                                                                    self._filehandle,
354                                                                    msg, pos, _bmappos, _datapos)
355                            self._index['msg'].append(msg)
356
357                            continue
358
359            except(struct.error):
360                if 'r' in self.mode:
361                    self._filehandle.seek(0)
362                break
363
364        # Index at end of _build_index()
365        if self._hasindex and not no_data:
366             self.variables = tuple(sorted(set([msg.shortName for msg in self._index['msg']])))
367             self.levels = tuple(sorted(set([msg.level for msg in self._index['msg']])))
368
369
370    def close(self):
371        """
372        Close the file handle
373        """
374        if not self._filehandle.closed:
375            self.messages = 0
376            self.current_message = 0
377            self._filehandle.close()
378            self.closed = self._filehandle.closed
379
380
381    def read(self, size=None):
382        """
383        Read size amount of GRIB2 messages from the current position.
384
385        If no argument is given, then size is None and all messages are returned from 
386        the current position in the file. This read method follows the behavior of 
387        Python's builtin open() function, but whereas that operates on units of bytes, 
388        we operate on units of GRIB2 messages.
389
390        Parameters
391        ----------
392        **`size : int, optional`**
393            The number of GRIB2 messages to read from the current position. If no argument is
394            give, the default value is `None` and remainder of the file is read.
395
396        Returns
397        -------
398        `Grib2Message` object when size = 1 or a `list` of Grib2Messages when size > 1.
399        """
400        if size is not None and size < 0:
401            size = None
402        if size is None or size > 1:
403            start = self.tell()
404            stop = self.messages if size is None else start+size
405            if size is None:
406                self.current_message = self.messages-1
407            else:
408                self.current_message += size
409            return self._index['msg'][slice(start,stop,1)]
410        elif size == 1:
411            self.current_message += 1
412            return self._index['msg'][self.current_message]
413        else:
414            None
415
416
417    def seek(self, pos):
418        """
419        Set the position within the file in units of GRIB2 messages.
420
421        Parameters
422        ----------
423        **`pos : int`**
424            The GRIB2 Message number to set the file pointer to.
425        """
426        if self._hasindex:
427            self._filehandle.seek(self._index['offset'][pos])
428            self.current_message = pos
429
430
431    def tell(self):
432        """
433        Returns the position of the file in units of GRIB2 Messages.
434        """
435        return self.current_message
436
437
438    def select(self, **kwargs):
439        """
440        Select GRIB2 messages by `Grib2Message` attributes.
441        """
442        # TODO: Added ability to process multiple values for each keyword (attribute)
443        idxs = []
444        nkeys = len(kwargs.keys())
445        for k,v in kwargs.items():
446            for m in self._index['msg']:
447                if hasattr(m,k) and getattr(m,k) == v: idxs.append(m._msgnum)
448        idxs = np.array(idxs,dtype=np.int32)
449        return [self._index['msg'][i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]
450
451
452    def write(self, msg):
453        """
454        Writes GRIB2 message object to file.
455
456        Parameters
457        ----------
458        **`msg : Grib2Message or sequence of Grib2Messages`**
459            GRIB2 message objects to write to file.
460        """
461        if isinstance(msg,list):
462            for m in msg:
463                self.write(m)
464            return
465
466        if issubclass(msg.__class__,_Grib2Message):
467            if hasattr(msg,'_msg'):
468                self._filehandle.write(msg._msg)
469            else:
470                if msg._signature != msg._generate_signature():
471                    msg.pack()
472                    self._filehandle.write(msg._msg)
473                else:
474                    if hasattr(msg._data,'filehandle'):
475                        msg._data.filehandle.seek(msg._data.offset)
476                        self._filehandle.write(msg._data.filehandle.read(msg.section0[-1]))
477                    else:
478                        msg.pack()
479                        self._filehandle.write(msg._msg)
480            self.flush()
481            self.size = os.path.getsize(self.name)
482            self._filehandle.seek(self.size-msg.section0[-1])
483            self._build_index()
484        else:
485            raise TypeError("msg must be a Grib2Message object.")
486        return
487
488
489    def flush(self):
490        """
491        Flush the file object buffer.
492        """
493        self._filehandle.flush()
494
495
496    def levels_by_var(self, name):
497        """
498        Return a list of level strings given a variable shortName.
499
500        Parameters
501        ----------
502        **`name : str`**
503            Grib2Message variable shortName
504
505        Returns
506        -------
507        A list of strings of unique level strings.
508        """
509        return list(sorted(set([msg.level for msg in self.select(shortName=name)])))
510
511
512    def vars_by_level(self, level):
513        """
514        Return a list of variable shortName strings given a level.
515
516        Parameters
517        ----------
518        **`level : str`**
519            Grib2Message variable level
520
521        Returns
522        -------
523        A list of strings of variable shortName strings.
524        """
525        return list(sorted(set([msg.shortName for msg in self.select(level=level)])))

528class Grib2Message:
529    """
530    Creation class for a GRIB2 message.
531    """
532    def __new__(self, section0: np.array = np.array([struct.unpack('>I',b'GRIB')[0],0,0,2,0]),
533                      section1: np.array = np.zeros((13),dtype=np.int64),
534                      section2: bytes = None,
535                      section3: np.array = None,
536                      section4: np.array = None,
537                      section5: np.array = None, *args, **kwargs):
538
539        if np.all(section1==0):
540            section1[5:11] = datetime.datetime.utcfromtimestamp(0).timetuple()[:6]
541
542        bases = list()
543        if section3 is None:
544            if 'gdtn' in kwargs.keys():
545                gdtn = kwargs['gdtn']
546                Gdt = templates.gdt_class_by_gdtn(gdtn)
547                bases.append(Gdt)
548                section3 = np.zeros((Gdt._len+5),dtype=np.int64)
549                section3[4] = gdtn
550            else:
551                raise ValueError("Must provide GRIB2 Grid Definition Template Number or section 3 array")
552        else:
553            gdtn = section3[4]
554            Gdt = templates.gdt_class_by_gdtn(gdtn)
555            bases.append(Gdt)
556
557        if section4 is None:
558            if 'pdtn' in kwargs.keys():
559                pdtn = kwargs['pdtn']
560                Pdt = templates.pdt_class_by_pdtn(pdtn)
561                bases.append(Pdt)
562                section4 = np.zeros((Pdt._len+2),dtype=np.int64)
563                section4[1] = pdtn
564            else:
565                raise ValueError("Must provide GRIB2 Production Definition Template Number or section 4 array")
566        else:
567            pdtn = section4[1]
568            Pdt = templates.pdt_class_by_pdtn(pdtn)
569            bases.append(Pdt)
570
571        if section5 is None:
572            if 'drtn' in kwargs.keys():
573                drtn = kwargs['drtn']
574                Drt = templates.drt_class_by_drtn(drtn)
575                bases.append(Drt)
576                section5 = np.zeros((Drt._len+2),dtype=np.int64)
577                section5[1] = drtn
578            else:
579                raise ValueError("Must provide GRIB2 Data Representation Template Number or section 5 array")
580        else:
581            drtn = section5[1]
582            Drt = templates.drt_class_by_drtn(drtn)
583            bases.append(Drt)
584
585        # attempt to use existing Msg class if it has already been made with gdtn,pdtn,drtn combo
586        try:
587            Msg = _msg_class_store[f"{gdtn}:{pdtn}:{drtn}"]
588        except KeyError:
589            @dataclass(init=False, repr=False)
590            class Msg(_Grib2Message, *bases):
591                pass
592            _msg_class_store[f"{gdtn}:{pdtn}:{drtn}"] = Msg
593
594        return Msg(section0, section1, section2, section3, section4, section5, *args)

 597@dataclass
 598class _Grib2Message:
 599    """GRIB2 Message base class"""
 600    # GRIB2 Sections
 601    section0: np.array = field(init=True,repr=False)
 602    section1: np.array = field(init=True,repr=False)
 603    section2: bytes = field(init=True,repr=False)
 604    section3: np.array = field(init=True,repr=False)
 605    section4: np.array = field(init=True,repr=False)
 606    section5: np.array = field(init=True,repr=False)
 607    bitMapFlag: templates.Grib2Metadata = field(init=True,repr=False,default=255)
 608
 609    # Section 0 looked up attributes
 610    indicatorSection: np.array = field(init=False,repr=False,default=templates.IndicatorSection())
 611    discipline: templates.Grib2Metadata = field(init=False,repr=False,default=templates.Discipline())
 612
 613    # Section 1 looked up attributes
 614    identificationSection: np.array = field(init=False,repr=False,default=templates.IdentificationSection())
 615    originatingCenter: templates.Grib2Metadata = field(init=False,repr=False,default=templates.OriginatingCenter())
 616    originatingSubCenter: templates.Grib2Metadata = field(init=False,repr=False,default=templates.OriginatingSubCenter())
 617    masterTableInfo: templates.Grib2Metadata = field(init=False,repr=False,default=templates.MasterTableInfo())
 618    localTableInfo: templates.Grib2Metadata = field(init=False,repr=False,default=templates.LocalTableInfo())
 619    significanceOfReferenceTime: templates.Grib2Metadata = field(init=False,repr=False,default=templates.SignificanceOfReferenceTime())
 620    year: int = field(init=False,repr=False,default=templates.Year())
 621    month: int = field(init=False,repr=False,default=templates.Month())
 622    day: int = field(init=False,repr=False,default=templates.Day())
 623    hour: int = field(init=False,repr=False,default=templates.Hour())
 624    minute: int = field(init=False,repr=False,default=templates.Minute())
 625    second: int = field(init=False,repr=False,default=templates.Second())
 626    refDate: datetime.datetime = field(init=False,repr=False,default=templates.RefDate())
 627    productionStatus: templates.Grib2Metadata = field(init=False,repr=False,default=templates.ProductionStatus())
 628    typeOfData: templates.Grib2Metadata = field(init=False,repr=False,default=templates.TypeOfData())
 629
 630    @property
 631    def _isNDFD(self):
 632        """Check if GRIB2 message is from NWS NDFD"""
 633        return np.all(self.section1[0:2]==[8,65535])
 634
 635    # Section 3 looked up common attributes.  Other looked up attributes are available according
 636    # to the Grid Definition Template.
 637    gridDefinitionSection: np.array = field(init=False,repr=False,default=templates.GridDefinitionSection())
 638    sourceOfGridDefinition: int = field(init=False,repr=False,default=templates.SourceOfGridDefinition())
 639    numberOfDataPoints: int = field(init=False,repr=False,default=templates.NumberOfDataPoints())
 640    interpretationOfListOfNumbers: templates.Grib2Metadata = field(init=False,repr=False,default=templates.InterpretationOfListOfNumbers())
 641    gridDefinitionTemplateNumber: templates.Grib2Metadata = field(init=False,repr=False,default=templates.GridDefinitionTemplateNumber())
 642    gridDefinitionTemplate: list = field(init=False,repr=False,default=templates.GridDefinitionTemplate())
 643    _earthparams: dict = field(init=False,repr=False,default=templates.EarthParams())
 644    _dxsign: float = field(init=False,repr=False,default=templates.DxSign())
 645    _dysign: float = field(init=False,repr=False,default=templates.DySign())
 646    _llscalefactor: float = field(init=False,repr=False,default=templates.LLScaleFactor())
 647    _lldivisor: float = field(init=False,repr=False,default=templates.LLDivisor())
 648    _xydivisor: float = field(init=False,repr=False,default=templates.XYDivisor())
 649    shapeOfEarth: templates.Grib2Metadata = field(init=False,repr=False,default=templates.ShapeOfEarth())
 650    earthShape: str = field(init=False,repr=False,default=templates.EarthShape())
 651    earthRadius: float = field(init=False,repr=False,default=templates.EarthRadius())
 652    earthMajorAxis: float = field(init=False,repr=False,default=templates.EarthMajorAxis())
 653    earthMinorAxis: float = field(init=False,repr=False,default=templates.EarthMinorAxis())
 654    resolutionAndComponentFlags: list = field(init=False,repr=False,default=templates.ResolutionAndComponentFlags())
 655    ny: int = field(init=False,repr=False,default=templates.Ny())
 656    nx: int = field(init=False,repr=False,default=templates.Nx())
 657    scanModeFlags: list = field(init=False,repr=False,default=templates.ScanModeFlags())
 658    projParameters: dict = field(init=False,repr=False,default=templates.ProjParameters())
 659
 660    # Section 4 attributes. Listed here are "extra" or "helper" attrs that use metadata from
 661    # the given PDT, but not a formal part of the PDT.
 662    productDefinitionTemplateNumber: templates.Grib2Metadata = field(init=False,repr=False,default=templates.ProductDefinitionTemplateNumber())
 663    productDefinitionTemplate: np.array = field(init=False,repr=False,default=templates.ProductDefinitionTemplate())
 664    _varinfo: list = field(init=False, repr=False, default=templates.VarInfo())
 665    _fixedsfc1info: list = field(init=False, repr=False, default=templates.FixedSfc1Info())
 666    _fixedsfc2info: list = field(init=False, repr=False, default=templates.FixedSfc2Info())
 667    fullName: str = field(init=False, repr=False, default=templates.FullName())
 668    units: str = field(init=False, repr=False, default=templates.Units())
 669    shortName: str = field(init=False, repr=False, default=templates.ShortName())
 670    leadTime: datetime.timedelta = field(init=False,repr=False,default=templates.LeadTime())
 671    unitOfFirstFixedSurface: str = field(init=False,repr=False,default=templates.UnitOfFirstFixedSurface())
 672    valueOfFirstFixedSurface: int = field(init=False,repr=False,default=templates.ValueOfFirstFixedSurface())
 673    unitOfSecondFixedSurface: str = field(init=False,repr=False,default=templates.UnitOfSecondFixedSurface())
 674    valueOfSecondFixedSurface: int = field(init=False,repr=False,default=templates.ValueOfSecondFixedSurface())
 675    level: str = field(init=False, repr=False, default=templates.Level())
 676    duration: datetime.timedelta = field(init=False,repr=False,default=templates.Duration())
 677    validDate: datetime.datetime = field(init=False,repr=False,default=templates.ValidDate())
 678
 679    # Section 5 looked up common attributes.  Other looked up attributes are available according
 680    # to the Data Representation Template.
 681    numberOfPackedValues: int = field(init=False,repr=False,default=templates.NumberOfPackedValues())
 682    dataRepresentationTemplateNumber: templates.Grib2Metadata = field(init=False,repr=False,default=templates.DataRepresentationTemplateNumber())
 683    dataRepresentationTemplate: list = field(init=False,repr=False,default=templates.DataRepresentationTemplate())
 684    typeOfValues: templates.Grib2Metadata = field(init=False,repr=False,default=templates.TypeOfValues())
 685
 686
 687    def __post_init__(self):
 688        """Set some attributes after init"""
 689        self._msgnum = -1
 690        self._deflist = None
 691        self._coordlist = None
 692        self._signature = self._generate_signature()
 693        try:
 694            self._sha1_section3 = hashlib.sha1(self.section3).hexdigest()
 695        except(TypeError):
 696            pass
 697        self.bitMapFlag = templates.Grib2Metadata(self.bitMapFlag,table='6.0')
 698
 699
 700    @property
 701    def gdtn(self):
 702        """Return Grid Definition Template Number"""
 703        return self.section3[4]
 704
 705
 706    @property
 707    def gdt(self):
 708        """Return Grid Definition Template"""
 709        return self.gridDefinitionTemplate
 710
 711
 712    @property
 713    def pdtn(self):
 714        """Return Product Definition Template Number"""
 715        return self.section4[1]
 716
 717
 718    @property
 719    def pdt(self):
 720        """Return Product Definition Template"""
 721        return self.productDefinitionTemplate
 722
 723
 724    @property
 725    def drtn(self):
 726        """Return Data Representation Template Number"""
 727        return self.section5[1]
 728
 729
 730    @property
 731    def drt(self):
 732        """Return Data Representation Template"""
 733        return self.dataRepresentationTemplate
 734
 735
 736    @property
 737    def pdy(self):
 738        """Return the PDY ('YYYYMMDD')"""
 739        return ''.join([str(i) for i in self.section1[5:8]])
 740
 741
 742    @property
 743    def griddef(self):
 744        """Return a Grib2GridDef instance for a GRIB2 message"""
 745        return Grib2GridDef.from_section3(self.section3)
 746
 747
 748    def __repr__(self):
 749        info = ''
 750        for sect in [0,1,3,4,5,6]:
 751            for k,v in self.attrs_by_section(sect,values=True).items():
 752                info += f'Section {sect}: {k} = {v}\n'
 753        return info
 754
 755
 756    def __str__(self):
 757        return (f'{self._msgnum}:d={self.refDate}:{self.shortName}:'
 758                f'{self.fullName} ({self.units}):{self.level}:'
 759                f'{self.leadTime}')
 760
 761
 762    def _generate_signature(self):
 763        """Generature SHA-1 hash string from GRIB2 integer sections"""
 764        return hashlib.sha1(np.concatenate((self.section0,self.section1,
 765                                            self.section3,self.section4,
 766                                            self.section5))).hexdigest()
 767
 768
 769    def attrs_by_section(self, sect, values=False):
 770        """
 771        Provide a tuple of attribute names for the given GRIB2 section.
 772
 773        Parameters
 774        ----------
 775        **`sect : int`**
 776            The GRIB2 section number.
 777
 778        **`values : bool, optional`**
 779            Optional (default is `False`) arugment to return attributes values.
 780
 781        Returns
 782        -------
 783        A List attribute names or Dict if `values = True`.
 784        """
 785        if sect in {0,1,6}:
 786            attrs = templates._section_attrs[sect]
 787        elif sect in {3,4,5}:
 788            def _find_class_index(n):
 789                _key = {3:'Grid', 4:'Product', 5:'Data'}
 790                for i,c in enumerate(self.__class__.__mro__):
 791                    if _key[n] in c.__name__:
 792                        return i
 793                else:
 794                    return []
 795            if sys.version_info.minor <= 8:
 796                attrs = templates._section_attrs[sect]+\
 797                        [a for a in dir(self.__class__.__mro__[_find_class_index(sect)]) if not a.startswith('_')]
 798            else:
 799                attrs = templates._section_attrs[sect]+\
 800                        self.__class__.__mro__[_find_class_index(sect)]._attrs
 801        else:
 802            attrs = []
 803        if values:
 804            return {k:getattr(self,k) for k in attrs}
 805        else:
 806            return attrs
 807
 808
 809    def pack(self):
 810        """
 811        Packs GRIB2 section data into a binary message.  It is the user's responsibility
 812        to populate the GRIB2 section information with appropriate metadata.
 813        """
 814        # Create beginning of packed binary message with section 0 and 1 data.
 815        self._sections = []
 816        self._msg,self._pos = g2clib.grib2_create(self.indicatorSection[2:4],self.identificationSection)
 817        self._sections += [0,1]
 818
 819        # Add section 2 if present.
 820        if isinstance(self.section2,bytes) and len(self.section2) > 0:
 821            self._msg,self._pos = g2clib.grib2_addlocal(self._msg,self.section2)
 822            self._sections.append(2)
 823
 824        # Add section 3.
 825        self.section3[1] = self.nx * self.ny
 826        self._msg,self._pos = g2clib.grib2_addgrid(self._msg,self.gridDefinitionSection,
 827                                                   self.gridDefinitionTemplate,
 828                                                   self._deflist)
 829        self._sections.append(3)
 830
 831        # Prepare data.
 832        field = np.copy(self.data)
 833        if self.scanModeFlags is not None:
 834            if self.scanModeFlags[3]:
 835                fieldsave = field.astype('f') # Casting makes a copy
 836                field[1::2,:] = fieldsave[1::2,::-1]
 837        fld = field.astype('f')
 838
 839        # Prepare bitmap, if necessary
 840        bitmapflag = self.bitMapFlag.value
 841        if bitmapflag == 0:
 842            bmap = np.ravel(np.where(np.isnan(fld),0,1)).astype(DEFAULT_NUMPY_INT)
 843        else:
 844            bmap = None
 845
 846        # Prepare optional coordinate list
 847        if self._coordlist is not None:
 848            crdlist = np.array(self._coordlist,'f')
 849        else:
 850            crdlist = None
 851
 852        # Prepare data for packing if nans are present
 853        fld = np.ravel(fld)
 854        if np.isnan(fld).any() and hasattr(self,'_missvalmap'):
 855            fld = np.where(self._missvalmap==1,self.priMissingValue,fld)
 856            fld = np.where(self._missvalmap==2,self.secMissingValue,fld)
 857
 858        # Add sections 4, 5, 6 (if present), and 7.
 859        self._msg,self._pos = g2clib.grib2_addfield(self._msg,self.pdtn,
 860                                                    self.productDefinitionTemplate,
 861                                                    crdlist,
 862                                                    self.drtn,
 863                                                    self.dataRepresentationTemplate,
 864                                                    fld,
 865                                                    bitmapflag,
 866                                                    bmap)
 867        self._sections.append(4)
 868        self._sections.append(5)
 869        if bmap is not None: self._sections.append(6)
 870        self._sections.append(7)
 871
 872        # Finalize GRIB2 message with section 8.
 873        self._msg, self._pos = g2clib.grib2_end(self._msg)
 874        self._sections.append(8)
 875        self.section0[-1] = len(self._msg)
 876
 877
 878    @property
 879    def data(self) -> np.array:
 880        """
 881        Accessing the data attribute loads data into memmory
 882        """
 883        if not hasattr(self,'_auto_nans'): self._auto_nans = _AUTO_NANS
 884        if hasattr(self,'_data'):
 885            if self._auto_nans != _AUTO_NANS:
 886                self._data = self._ondiskarray
 887            if isinstance(self._data, Grib2MessageOnDiskArray):
 888                self._ondiskarray = self._data
 889                self._data = np.asarray(self._data)
 890            return self._data
 891        raise ValueError
 892
 893    @data.setter
 894    def data(self, data):
 895        if not isinstance(data, np.ndarray):
 896            raise ValueError('Grib2Message data only supports numpy arrays')
 897        self._data = data
 898
 899
 900    def __getitem__(self, item):
 901        return self.data[item]
 902
 903
 904    def __setitem__(self, item):
 905        raise NotImplementedError('assignment of data not supported via setitem')
 906
 907
 908    def latlons(self, *args, **kwrgs):
 909        """Alias for `grib2io.Grib2Message.grid` method"""
 910        return self.grid(*args, **kwrgs)
 911
 912
 913    def grid(self, unrotate=True):
 914        """
 915        Return lats,lons (in degrees) of grid.
 916
 917        Currently can handle reg. lat/lon,cglobal Gaussian, mercator, stereographic, 
 918        lambert conformal, albers equal-area, space-view and azimuthal equidistant grids.
 919
 920        Parameters
 921        ----------
 922        **`unrotate : bool`**
 923            If `True` [DEFAULT], and grid is rotated lat/lon, then unrotate the grid,
 924            otherwise `False`, do not.
 925
 926        Returns
 927        -------
 928        **`lats, lons : numpy.ndarray`**
 929            Returns two numpy.ndarrays with dtype=numpy.float32 of grid latitudes and
 930            longitudes in units of degrees.
 931        """
 932        if self._sha1_section3 in _latlon_datastore.keys():
 933            return (_latlon_datastore[self._sha1_section3]['latitude'],
 934                    _latlon_datastore[self._sha1_section3]['longitude'])
 935        gdtn = self.gridDefinitionTemplateNumber.value
 936        gdtmpl = self.gridDefinitionTemplate
 937        reggrid = self.gridDefinitionSection[2] == 0 # This means regular 2-d grid
 938        if gdtn == 0:
 939            # Regular lat/lon grid
 940            lon1, lat1 = self.longitudeFirstGridpoint, self.latitudeFirstGridpoint
 941            lon2, lat2 = self.longitudeLastGridpoint, self.latitudeLastGridpoint
 942            dlon = self.gridlengthXDirection
 943            dlat = self.gridlengthYDirection
 944            if lon2 < lon1 and dlon < 0: lon1 = -lon1
 945            lats = np.linspace(lat1,lat2,self.ny)
 946            if reggrid:
 947                lons = np.linspace(lon1,lon2,self.nx)
 948            else:
 949                lons = np.linspace(lon1,lon2,self.ny*2)
 950            lons,lats = np.meshgrid(lons,lats) # Make 2-d arrays.
 951        elif gdtn == 1: # Rotated Lat/Lon grid
 952            pj = pyproj.Proj(self.projParameters)
 953            lat1,lon1 = self.latitudeFirstGridpoint,self.longitudeFirstGridpoint
 954            lat2,lon2 = self.latitudeLastGridpoint,self.longitudeLastGridpoint
 955            if lon1 > 180.0: lon1 -= 360.0
 956            if lon2 > 180.0: lon2 -= 360.0
 957            lats = np.linspace(lat1,lat2,self.ny)
 958            lons = np.linspace(lon1,lon2,self.nx)
 959            lons,lats = np.meshgrid(lons,lats) # Make 2-d arrays.
 960            if unrotate:
 961                from grib2io.utils import rotated_grid
 962                lats,lons = rotated_grid.unrotate(lats,lons,self.anglePoleRotation,
 963                                                  self.latitudeSouthernPole,
 964                                                  self.longitudeSouthernPole)
 965        elif gdtn == 40: # Gaussian grid (only works for global!)
 966            from grib2io.utils.gauss_grid import gaussian_latitudes
 967            lon1, lat1 = self.longitudeFirstGridpoint, self.latitudeFirstGridpoint
 968            lon2, lat2 = self.longitudeLastGridpoint, self.latitudeLastGridpoint
 969            nlats = self.ny
 970            if not reggrid: # Reduced Gaussian grid.
 971                nlons = 2*nlats
 972                dlon = 360./nlons
 973            else:
 974                nlons = self.nx
 975                dlon = self.gridlengthXDirection
 976            lons = np.linspace(lon1,lon2,nlons)
 977            # Compute Gaussian lats (north to south)
 978            lats = gaussian_latitudes(nlats)
 979            if lat1 < lat2:  # reverse them if necessary
 980                lats = lats[::-1]
 981            lons,lats = np.meshgrid(lons,lats)
 982        elif gdtn in {10,20,30,31,110}:
 983            # Mercator, Lambert Conformal, Stereographic, Albers Equal Area, Azimuthal Equidistant
 984            dx,dy = self.gridlengthXDirection, self.gridlengthYDirection
 985            lon1,lat1 = self.longitudeFirstGridpoint, self.latitudeFirstGridpoint
 986            pj = pyproj.Proj(self.projParameters)
 987            llcrnrx, llcrnry = pj(lon1,lat1)
 988            x = llcrnrx+dx*np.arange(self.nx)
 989            y = llcrnry+dy*np.arange(self.ny)
 990            x,y = np.meshgrid(x, y)
 991            lons,lats = pj(x, y, inverse=True)
 992        elif gdtn == 90:
 993            # Satellite Projection
 994            dx = self.gridlengthXDirection
 995            dy = self.gridlengthYDirection
 996            pj = pyproj.Proj(self.projParameters)
 997            x = dx*np.indices((self.ny,self.nx),'f')[1,:,:]
 998            x -= 0.5*x.max()
 999            y = dy*np.indices((self.ny,self.nx),'f')[0,:,:]
1000            y -= 0.5*y.max()
1001            lons,lats = pj(x,y,inverse=True)
1002            # Set lons,lats to 1.e30 where undefined
1003            abslons = np.fabs(lons)
1004            abslats = np.fabs(lats)
1005            lons = np.where(abslons < 1.e20, lons, 1.e30)
1006            lats = np.where(abslats < 1.e20, lats, 1.e30)
1007        elif gdtn == 32769:
1008            # Special NCEP Grid, Rotated Lat/Lon, Arakawa E-Grid (Non-Staggered)
1009            from grib2io.utils import arakawa_rotated_grid
1010            from grib2io.utils.rotated_grid import DEG2RAD
1011            di, dj = 0.0, 0.0
1012            do_180 = False
1013            idir = 1 if self.scanModeFlags[0] == 0 else -1
1014            jdir = -1 if self.scanModeFlags[1] == 0 else 1
1015            grid_oriented = 0 if self.resolutionAndComponentFlags[4] == 0 else 1
1016            do_rot = 1 if grid_oriented == 1 else 0
1017            la1 = self.latitudeFirstGridpoint
1018            lo1 = self.longitudeFirstGridpoint
1019            clon = self.longitudeCenterGridpoint
1020            clat = self.latitudeCenterGridpoint
1021            lasp = clat - 90.0
1022            losp = clon
1023            llat, llon = arakawa_rotated_grid.ll2rot(la1,lo1,lasp,losp)
1024            la2, lo2 = arakawa_rotated_grid.rot2ll(-llat,-llon,lasp,losp)
1025            rlat = -llat
1026            rlon = -llon
1027            if self.nx == 1:
1028                di = 0.0
1029            elif idir == 1:
1030                ti = rlon
1031                while ti < llon:
1032                    ti += 360.0
1033                di = (ti - llon)/float(self.nx-1)
1034            else:
1035                ti = llon
1036                while ti < rlon:
1037                    ti += 360.0
1038                di = (ti - rlon)/float(self.nx-1)
1039            if self.ny == 1:
1040               dj = 0.0
1041            else:
1042                dj = (rlat - llat)/float(self.ny-1)
1043                if dj < 0.0:
1044                    dj = -dj
1045            if idir == 1:
1046                if llon > rlon:
1047                    llon -= 360.0
1048                if llon < 0 and rlon > 0:
1049                    do_180 = True
1050            else:
1051                if rlon > llon:
1052                    rlon -= 360.0
1053                if rlon < 0 and llon > 0:
1054                    do_180 = True
1055            xlat1d = llat + (np.arange(self.ny)*jdir*dj)
1056            xlon1d = llon + (np.arange(self.nx)*idir*di)
1057            xlons, xlats = np.meshgrid(xlon1d,xlat1d)
1058            rot2ll_vectorized = np.vectorize(arakawa_rotated_grid.rot2ll)
1059            lats, lons = rot2ll_vectorized(xlats,xlons,lasp,losp)
1060            if do_180:
1061                lons = np.where(lons>180.0,lons-360.0,lons)
1062            vector_rotation_angles_vectorized = np.vectorize(arakawa_rotated_grid.vector_rotation_angles)
1063            rots = vector_rotation_angles_vectorized(lats, lons, clat, losp, xlats)
1064            del xlat1d, xlon1d, xlats, xlons
1065        else:
1066            raise ValueError('Unsupported grid')
1067
1068        _latlon_datastore[self._sha1_section3] = dict(latitude=lats,longitude=lons)
1069        try:
1070            _latlon_datastore[self._sha1_section3]['vector_rotation_angles'] = rots
1071        except(NameError):
1072            pass
1073
1074        return lats, lons
1075
1076
1077    def map_keys(self):
1078        """
1079        Returns an unpacked data grid where integer grid values are replaced with
1080        a string.in which the numeric value is a representation of.
1081
1082        These types of fields are cateogrical or classifications where data values 
1083        do not represent an observable or predictable physical quantity. An example 
1084        of such a field field would be [Dominant Precipitation Type -
1085        DPTYPE](https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-201.shtml)
1086
1087        Returns
1088        -------
1089        **`numpy.ndarray`** of string values per element.
1090        """
1091        hold_auto_nans = _AUTO_NANS
1092        set_auto_nans(False)
1093        if (np.all(self.section1[0:2]==[7,14]) and self.shortName == 'PWTHER') or \
1094        (np.all(self.section1[0:2]==[8,65535]) and self.shortName == 'WX'):
1095            keys = utils.decode_wx_strings(self.section2)
1096            if hasattr(self,'priMissingValue') and self.priMissingValue not in [None,0]:
1097                keys[int(self.priMissingValue)] = 'Missing'
1098            if hasattr(self,'secMissingValue') and self.secMissingValue not in [None,0]:
1099                keys[int(self.secMissingValue)] = 'Missing'
1100            u,inv = np.unique(self.data,return_inverse=True)
1101            fld = np.array([keys[x] for x in u])[inv].reshape(self.data.shape)
1102        else:
1103            # For data whose units are defined in a code table (i.e. classification or mask)
1104            tblname = re.findall(r'\d\.\d+',self.units,re.IGNORECASE)[0]
1105            fld = self.data.astype(np.int32).astype(str)
1106            tbl = tables.get_table(tblname,expand=True)
1107            for val in np.unique(fld):
1108                fld = np.where(fld==val,tbl[val],fld)
1109        set_auto_nans(hold_auto_nans)
1110        return fld
1111
1112
1113    def to_bytes(self, validate=True):
1114        """
1115        Return packed GRIB2 message in bytes format.
1116
1117        This will be Useful for exporting data in non-file formats. For example, 
1118        can be used to output grib data directly to S3 using the boto3 client 
1119        without the need to write a temporary file to upload first.
1120
1121        Parameters
1122        ----------
1123        **`validate : bool, optional`**
1124            If `True` (DEFAULT), validates first/last four bytes for proper
1125            formatting, else returns None. If `False`, message is output as is.
1126
1127        Returns
1128        -------
1129        Returns GRIB2 formatted message as bytes.
1130        """
1131        if validate:
1132            if self._msg[0:4]+self._msg[-4:] == b'GRIB7777':
1133                return self._msg
1134            else:
1135                return None
1136        else:
1137            return self._msg
1138
1139
1140    def interpolate(self, method, grid_def_out, method_options=None):
1141        """
1142        Perform grid spatial interpolation via the [NCEPLIBS-ip library](https://github.com/NOAA-EMC/NCEPLIBS-ip).
1143
1144        **IMPORTANT:**  This interpolate method only supports scalar interpolation. If you
1145        need to perform vector interpolation, use the module-level `grib2io.interpolate` function.
1146
1147        Parameters
1148        ----------
1149        **`method : int or str`**
1150            Interpolate method to use. This can either be an integer or string using
1151            the following mapping:
1152
1153        | Interpolate Scheme | Integer Value |
1154        | :---:              | :---:         |
1155        | 'bilinear'         | 0             |
1156        | 'bicubic'          | 1             |
1157        | 'neighbor'         | 2             |
1158        | 'budget'           | 3             |
1159        | 'spectral'         | 4             |
1160        | 'neighbor-budget'  | 6             |
1161
1162        **`grid_def_out : grib2io.Grib2GridDef`**
1163            Grib2GridDef object of the output grid.
1164
1165        **`method_options : list of ints, optional`**
1166            Interpolation options. See the NCEPLIBS-ip doucmentation for
1167            more information on how these are used.
1168
1169        Returns
1170        -------
1171        If interpolating to a grid, a new Grib2Message object is returned.  The GRIB2 metadata of
1172        the new Grib2Message object is indentical to the input except where required to be different
1173        because of the new grid specs.
1174
1175        If interpolating to station points, the interpolated data values are returned as a numpy.ndarray.
1176        """
1177        section0 = self.section0
1178        section0[-1] = 0
1179        gds = [0, grid_def_out.npoints, 0, 255, grid_def_out.gdtn]
1180        section3 = np.concatenate((gds,grid_def_out.gdt))
1181
1182        msg = Grib2Message(section0,self.section1,self.section2,section3,
1183                           self.section4,self.section5,self.bitMapFlag.value)
1184
1185        msg._msgnum = -1
1186        msg._deflist = self._deflist
1187        msg._coordlist = self._coordlist
1188        shape = (msg.ny,msg.nx)
1189        ndim = 2
1190        if msg.typeOfValues == 0:
1191            dtype = 'float32'
1192        elif msg.typeOfValues == 1:
1193            dtype = 'int32'
1194        msg._data = interpolate(self.data,method,Grib2GridDef.from_section3(self.section3),grid_def_out,
1195                                method_options=method_options).reshape(msg.ny,msg.nx)
1196        return msg