Merge pull request #190 from dstansby/attdata

Add AttData

cdflib/cdfread.py

@@ -12,7 +12,7 @@
 import numpy as np
 
 import cdflib.epochs as epoch
-from cdflib.dataclasses import AEDR, VDR, ADRInfo, CDRInfo, GDRInfo
+from cdflib.dataclasses import AEDR, VDR, ADRInfo, AttData, CDRInfo, GDRInfo
 
 __all__ = ["CDF"]
 
@@ -331,7 +331,7 @@ def print_attrs(self):
             print("NAME: " + name + ", NUMBER: " + str(x) + ", SCOPE: " + attrs[x][name])
         return attrs
 
-    def attget(self, attribute=None, entry=None):
+    def attget(self, attribute=None, entry=None) -> AttData:
         """
         Returns the value of the attribute at the entry number provided.
 
@@ -343,24 +343,10 @@ def attget(self, attribute=None, entry=None):
         attribute : str, int, optional
             Attribute name or number to get.
         entry : int, optional
-        tp_np : bool, optional
-            If True, return a numpy array.
 
         Returns
         -------
-        dict
-            A dictionary is returned with the following defined keys
-
-            +-----------------+--------------------------------------------------------------------------------+
-            | ['Item_Size']   | the number of bytes for each entry value                                       |
-            +-----------------+--------------------------------------------------------------------------------+
-            | ['Num_Items']   | total number of values extracted                                               |
-            +-----------------+--------------------------------------------------------------------------------+
-            | ['Data_Type']   | the CDF data type                                                              |
-            +-----------------+--------------------------------------------------------------------------------+
-            | ['Data']        | retrieved attribute data as a scalar value, a numpy array or a string          |
-            +-----------------+--------------------------------------------------------------------------------+
-
+        AttData
         """
         # Starting position
         position = self._first_adr
@@ -1663,7 +1649,7 @@ def _endian(self) -> str:
             return "little-endian"
 
     @staticmethod
-    def _type_size(data_type, num_elms):
+    def _type_size(data_type, num_elms: int) -> int:
         # DATA TYPES
         #
         # 1 - 1 byte signed int
@@ -1701,6 +1687,8 @@ def _type_size(data_type, num_elms):
                 return 16
             elif (data_type == 51) or (data_type == 52):
                 return num_elms
+            else:
+                raise TypeError("Unknown data type....")
         elif isinstance(data_type, str):
             data_typeU = data_type.upper()
             if (data_typeU == "CDF_INT1") or (data_typeU == "CDF_UINT1") or (data_typeU == "CDF_BYTE"):
@@ -1719,6 +1707,8 @@ def _type_size(data_type, num_elms):
                 return 16
             elif (data_typeU == "CDF_CHAR") or (data_typeU == "CDF_UCHAR"):
                 return num_elms
+            else:
+                raise TypeError("Unknown data type....")
         else:
             raise TypeError("Unknown data type....")
 
@@ -1846,25 +1836,24 @@ def _num_values(self, vdr: VDR) -> int:
                 values = values * vdr.dim_sizes[x]
         return values
 
-    def _get_attdata(self, adr_info, entry_num, num_entry, first_entry):
+    def _get_attdata(self, adr_info, entry_num, num_entry, first_entry) -> AttData:
         position = first_entry
         for _ in range(0, num_entry):
             got_entry_num, next_aedr = self._read_aedr_fast(position)
             if entry_num == got_entry_num:
                 aedr_info = self._read_aedr(position)
-                return_dict = {}
-                return_dict["Item_Size"] = self._type_size(aedr_info.data_type, aedr_info.num_elements)
-                return_dict["Data_Type"] = self._datatype_token(aedr_info.data_type)
+                item_size = self._type_size(aedr_info.data_type, aedr_info.num_elements)
+                data_type = self._datatype_token(aedr_info.data_type)
 
-                return_dict["Num_Items"] = aedr_info.num_elements
-                return_dict["Data"] = aedr_info.entry
+                num_items = aedr_info.num_elements
+                data = aedr_info.entry
                 if aedr_info.data_type == 51 or aedr_info.data_type == 52:
                     if aedr_info.num_strings is not None:
                         num_strings = aedr_info.num_strings
-                        return_dict["Num_Items"] = num_strings
+                        num_items = num_strings
                         if num_strings > 1:
-                            return_dict["Data"] = aedr_info.entry.split("\\N ")
-                return return_dict
+                            data = aedr_info.entry.split("\\N ")
+                return AttData(item_size, data_type, num_items, data)
             else:
                 position = next_aedr
 
@@ -1936,13 +1925,13 @@ def _findtimerecords(self, var_name, starttime, endtime, epoch=None):
                         "variable"
                     )
 
-                vdr_info = self.varinq(dependVar["Data"])
+                vdr_info = self.varinq(dependVar.Data)
                 if vdr_info["Data_Type"] != 31 and vdr_info["Data_Type"] != 32 and vdr_info["Data_Type"] != 33:
                     raise ValueError(
                         "Corresponding variable from 'DEPEND_0' attribute "
                         "for variable: {}".format(var_name) + " is not a CDF epoch type"
                     )
-                epochtimes = self.varget(dependVar["Data"])
+                epochtimes = self.varget(dependVar.Data)
 
         return self._findrangerecords(vdr_info["Data_Type"], epochtimes, starttime, endtime)
 

cdflib/dataclasses.py

@@ -73,3 +73,25 @@ class VDR:
     num_elements: int
     sparse: int
     pad: Optional[bool] = None
+
+
+@dataclass
+class AttData:
+    """
+    Attribute data.
+
+    Attributes
+    ----------
+    Item_size : int
+        Number of bytes for each entry value.
+    Data_Type : int
+        CDF data type.
+    Num_Items : int
+        Number of values extracted.
+    Data : numpy.ndarray
+        Data as a scalar value, a numpy array or a string.
+    """
+    Item_Size: int
+    Data_Type: int
+    Num_Items: int
+    Data: np.ndarray

doc/changelog.rst

@@ -37,6 +37,7 @@ Breaking changes
   Use ``CDF.attinq`` to get attribute information instead.
 - ``CDF.vdr_info`` now returns a dataclass instead of a dict.
 - ``CDF.attinq`` now returns a dataclass instead of a dict.
+- ``CDF.attget`` now returns a dataclass instead of a dict.
 
 Bugfixes
 --------

doc/modules/dataclasses.rst

@@ -0,0 +1,8 @@
+Dataclasses
+===========
+
+CDF Writer Class
+=================
+
+.. automodapi:: cdflib.dataclasses
+    :no-inheritance-diagram:

doc/modules/index.rst

@@ -7,4 +7,5 @@ API Reference
    cdfread
    cdfwrite
    cdfepoch
+   dataclasses
    xarray

tests/test_cdfwrite.py

@@ -2,7 +2,6 @@
 from typing import Any, Dict
 
 import numpy as np
-import pytest
 
 from cdflib import cdfread, cdfwrite
 
@@ -83,10 +82,10 @@ def test_checksum_compressed(tmp_path):
     assert (var == v).all()
 
     att = reader.attget("Attribute1", entry=0)
-    assert att["Data"] == [1]
+    assert att.Data == [1]
 
     att = reader.attget("Attribute2", entry=0)
-    assert att["Data"] == "500"
+    assert att.Data == "500"
 
 
 def test_file_compression(tmp_path):
@@ -150,10 +149,10 @@ def test_globalattrs(tmp_path):
     assert attrib.num_gr_entry == 4
 
     entry = reader.attget("Global6", 3)
-    assert entry["Data_Type"] == "CDF_INT8"
+    assert entry.Data_Type == "CDF_INT8"
 
     for x in [0, 1, 2]:
-        assert entry["Data"][x] == x
+        assert entry.Data[x] == x
 
 
 def test_create_zvariable(tmp_path):
@@ -267,10 +266,10 @@ def test_create_zvariables_with_attributes(tmp_path):
 
     # Test CDF info
     att = reader.attget("Attribute1", entry=0)
-    assert att["Data"] == [1]
+    assert att.Data == [1]
 
     att = reader.attget("Attribute2", entry=1)
-    assert att["Data"] == "1000"
+    assert att.Data == "1000"
 
 
 def test_create_zvariables_then_attributes(tmp_path):
@@ -302,10 +301,10 @@ def test_create_zvariables_then_attributes(tmp_path):
 
     # Test CDF info
     att = reader.attget("Attribute1", entry=0)
-    assert att["Data"] == [1]
+    assert att.Data == [1]
 
     att = reader.attget("Attribute2", entry=1)
-    att["Data"] == "1000"
+    att.Data == "1000"
 
 
 def test_nonsparse_zvariable_blocking(tmp_path):
@@ -569,11 +568,7 @@ def test_create_2d_r_and_z_variables(tmp_path):
     assert (var == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).all()
 
     att = reader.attget("Attribute1", entry="Variable2")
-    assert att["Data"] == [2]
+    assert att.Data == [2]
 
     att = reader.attget("Attribute2", entry="Variable2")
-    assert att["Data"] == "1000"
-
-
-if __name__ == "__main__":
-    pytest.main([__file__])
+    assert att.Data == "1000"