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

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