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

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

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