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

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

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