Merge pull request #395 from hkchekc/enum

Change base class of LinkType and DimensionType to Enum

LGTM

nixio/data_array.py

@@ -186,11 +186,11 @@ def _dimension_count(self):
     def _get_dimension_by_pos(self, index):
         h5dim = self._h5group.open_group("dimensions").open_group(str(index))
         dimtype = h5dim.get_attr("dimension_type")
-        if dimtype == DimensionType.Sample:
+        if DimensionType(dimtype) == DimensionType.Sample:
             return SampledDimension(h5dim, index)
-        elif dimtype == DimensionType.Range:
+        elif DimensionType(dimtype) == DimensionType.Range:
             return RangeDimension(h5dim, index)
-        elif dimtype == DimensionType.Set:
+        elif DimensionType(dimtype) == DimensionType.Set:
             return SetDimension(h5dim, index)
         else:
             raise TypeError("Invalid Dimension object in file.")
@@ -281,7 +281,7 @@ def unit(self, u):
             u = None
         util.check_attr_type(u, str)
         if (self._dimension_count() == 1 and
-                self.dimensions[0].dimension_type == DimensionType.Range and
+            self.dimensions[0].dimension_type == DimensionType.Range and
                 self.dimensions[0].is_alias and u is not None):
             if not (util.units.is_si(u) or util.units.is_compound(u)):
                 raise InvalidUnit(

nixio/dimension_type.py

@@ -7,8 +7,11 @@
 # modification, are permitted under the terms of the BSD License. See
 # LICENSE file in the root of the Project.
 
+from enum import Enum
 
-class DimensionType(object):
+
+class DimensionType(Enum):
+    # Values of enum are also part of NIX format
     Sample = "sample"
     Range = "range"
     Set = "set"

nixio/dimensions.py

@@ -14,6 +14,7 @@
 from .dimension_type import DimensionType
 from . import util
 from .container import Container
+from six import string_types
 
 
 class DimensionContainer(Container):
@@ -27,7 +28,7 @@ def _inst_item(self, item):
             DimensionType.Range: RangeDimension,
             DimensionType.Sample: SampledDimension,
             DimensionType.Set: SetDimension,
-        }[item.get_attr("dimension_type")]
+        }[DimensionType(item.get_attr("dimension_type"))]
         idx = item.name
         return cls(item, idx)
 
@@ -46,14 +47,14 @@ def _create_new(cls, parent, index):
 
     @property
     def dimension_type(self):
-        return self._h5group.get_attr("dimension_type")
+        return DimensionType(self._h5group.get_attr("dimension_type"))
 
     @dimension_type.setter
     def dimension_type(self, dimtype):
-        if dimtype not in (DimensionType.Sample, DimensionType.Range,
-                           DimensionType.Set):
-            raise ValueError("Invalid dimension type.")
-        self._h5group.set_attr("dimension_type", dimtype)
+        dimtype = DimensionType(dimtype)
+        if dimtype not in DimensionType:
+            raise TypeError("Invalid dimension type.")
+        self._h5group.set_attr("dimension_type", dimtype.value)
 
     @property
     def index(self):

nixio/feature.py

@@ -8,7 +8,8 @@
 # LICENSE file in the root of the Project.
 from .entity import Entity
 from .data_array import DataArray
-from .util import util
+from .link_type import LinkType
+from six import string_types
 
 
 class Feature(Entity):
@@ -29,11 +30,12 @@ def id(self):
 
     @property
     def link_type(self):
-        return util.link_type_from_string(self._h5group.get_attr("link_type"))
+        return LinkType(self._h5group.get_attr("link_type"))
 
     @link_type.setter
     def link_type(self, lt):
-        self._h5group.set_attr("link_type", util.link_type_to_string(lt))
+        lt = LinkType(lt)
+        self._h5group.set_attr("link_type", lt.value)
         if self._parent._parent._parent.time_auto_update:
             self.force_updated_at()
 

nixio/link_type.py

@@ -7,8 +7,11 @@
 # modification, are permitted under the terms of the BSD License. See
 # LICENSE file in the root of the Project.
 
+from enum import Enum
 
-class LinkType(object):
-    Tagged = "Tagged"
-    Untagged = "Untagged"
-    Indexed = "Indexed"
+
+class LinkType(Enum):
+    # Values of enum are also part of NIX format
+    Tagged = "tagged"
+    Untagged = "untagged"
+    Indexed = "indexed"

nixio/tag.py

@@ -98,6 +98,7 @@ def create_feature(self, data, link_type):
         :returns: The created feature object.
         :rtype: Feature
         """
+        link_type = LinkType(link_type)
         features = self._h5group.open_group("features")
         feat = Feature._create_new(self, features, data, link_type)
         return feat

nixio/test/test_feature.py

@@ -81,3 +81,6 @@ def test_feature_on_group(self):
 
         grouptag = self.group.tags[-1]
         grouptag.create_feature(self.movie1, nix.LinkType.Tagged)
+
+    def test_create_diff_link_type_style(self):
+        self.stimuli_tag.create_feature(self.movie1, nix.LinkType.Tagged)

nixio/test/test_nix_compatibility.py

@@ -179,7 +179,7 @@ def test_tags(tmpdir, bindir):
             grp.tags.append(tag)
 
         if idx == 5:
-            tag.create_feature(da1, "Tagged")
+            tag.create_feature(da1, nix.LinkType.Tagged)
     nix_file.close()
     # validate(nixfilepath)
     cmd = os.path.join(bindir, "readtags")
@@ -209,7 +209,7 @@ def test_multi_tags(tmpdir, bindir):
         if idx == 3:
             feada = blk.create_data_array("feature", "afea",
                                           data=np.random.random(200))
-            mt.create_feature(feada, "Tagged")
+            mt.create_feature(feada, nix.LinkType.Tagged)
 
         if idx != 1:
             mt.extents = extda

nixio/util/util.py

@@ -143,28 +143,3 @@ def apply_polynomial(coefficients, origin, data):
     data[:] = data[:] - origin
     if coefficients:
         data[:] = np.polynomial.polynomial.polyval(data, coefficients)
-
-
-# The following two functions currently behave as capitalised <-> lowercase
-# converters, but they are general solutions for alternate implementations of
-# LinkType (e.g., enum)
-def link_type_to_string(lt):
-    if lt == LinkType.Indexed:
-        return "indexed"
-    elif lt == LinkType.Tagged:
-        return "tagged"
-    elif lt == LinkType.Untagged:
-        return "untagged"
-    else:
-        raise RuntimeError("Invalid LinkType")
-
-
-def link_type_from_string(ltstr):
-    if ltstr == "indexed":
-        return LinkType.Indexed
-    elif ltstr == "tagged":
-        return LinkType.Tagged
-    elif ltstr == "untagged":
-        return LinkType.Untagged
-    else:
-        raise RuntimeError("Invalid string for LinkType")

nixio/validate.py

@@ -9,6 +9,7 @@
 from __future__ import (absolute_import, division, print_function)
 import numpy as np
 from .util import units
+from .dimension_type import DimensionType
 try:
     from collections.abc import OrderedDict
 except ImportError:
@@ -146,7 +147,7 @@ def check_data_arrays(self, da, da_idx, blk_idx):
 
         if da.dimensions:
             for i, dim in enumerate(da.dimensions):
-                if dim.dimension_type == 'range':
+                if dim.dimension_type == DimensionType.Range:
                     try:
                         if len(dim.ticks) != da.data_extent[i]:
                             # if data_extent is used instead of len()
@@ -159,7 +160,7 @@ def check_data_arrays(self, da, da_idx, blk_idx):
                         raise IndexError("Dimension of Dataarray and "
                                          "Number of Dimension object Mismatch")
 
-                if dim.dimension_type == 'set':
+                if dim.dimension_type == DimensionType.Set:
                     # same as above
                     try:
                         if len(dim.labels) != da.data_extent[i]:
@@ -221,7 +222,7 @@ def check_tag(self, tag, tag_idx, blk_idx):
                 unit_list = [un for un in unit_list if un]
                 dim_list = [dim for refer in tag.references
                             for dim in refer.dimensions
-                            if dim.dimension_type != 'set']
+                            if dim.dimension_type != DimensionType.Set]
                 if len(unit_list) != len(dim_list):
                     tag_err_list.append(
                         "Some dimensions of references have no units"
@@ -388,7 +389,7 @@ def check_range_dim(self, r_dim, dim_idx, da_idx, blk_idx):
         elif not all(r_dim.ticks[i] <= r_dim.ticks[i+1]
                      for i in range(len(r_dim.ticks)-1)):
             rdim_err_list.append("Ticks are not sorted!")
-        if r_dim.dimension_type != 'range':
+        if r_dim.dimension_type != DimensionType.Range:
             rdim_err_list.append("Dimension type is not correct!")
 
         # sorting is already covered in the dimensions.py file
@@ -412,7 +413,7 @@ def check_set_dim(self, set_dim, dim_idx, da_idx, blk_idx):
             da = self.errors['blocks'][blk_idx]['data_arrays'][da_idx]
             da['dimensions'][dim_idx]['errors'].append(self.check_dim(set_dim))
             self.error_count += 1
-        if set_dim.dimension_type != 'set':
+        if set_dim.dimension_type != DimensionType.Set:
             da = self.errors['blocks'][blk_idx]['data_arrays'][da_idx]
             dim = da['dimensions'][dim_idx]
             dim['errors'].append("Dimension type is not correct!")
@@ -435,7 +436,7 @@ def check_sampled_dim(self, sam_dim, dim_idx, da_idx, blk_idx):
             sdim_err_list.append("SamplingInterval is not set to valid value "
                                  "(> 0)!")
 
-        if sam_dim.dimension_type != 'sample':
+        if sam_dim.dimension_type != DimensionType.Sample:
             sdim_err_list.append("Dimension type is not correct!")
 
         if sam_dim.offset and not sam_dim.unit:
@@ -472,10 +473,7 @@ def get_dim_units(self, data_arrays):
         """
         unit_list = []
         for dim in data_arrays.dimensions:
-            if dim.dimension_type == 'range':
+            if dim.dimension_type == DimensionType.Range or\
+                    dim.dimension_type == DimensionType.Sample:
                 unit_list.append(dim.unit)
-            if dim.dimension_type == 'sample':
-                unit_list.append(dim.unit)
-            if dim.dimension_type == 'set':
-                pass
         return unit_list