Fix indexed enum typings

Author: Mauko Quiroga

openfisca_core/indexed_enums/enum.py

@@ -7,40 +7,47 @@
 
 from openfisca_core.indexed_enums import config, EnumArray
 
+if typing.TYPE_CHECKING:
+    IndexedEnumArray = numpy.object_
+
 
 class Enum(enum.Enum):
     """
-    Enum based on `enum34 <https://pypi.python.org/pypi/enum34/>`_, whose items have an
-    index.
+    Enum based on `enum34 <https://pypi.python.org/pypi/enum34/>`_, whose items
+    have an index.
     """
 
     # Tweak enums to add an index attribute to each enum item
     def __init__(self, name: str) -> None:
-        # When the enum item is initialized, self._member_names_ contains the names of
-        # the previously initialized items, so its length is the index of this item.
+        # When the enum item is initialized, self._member_names_ contains the
+        # names of the previously initialized items, so its length is the index
+        # of this item.
         self.index = len(self._member_names_)
 
     # Bypass the slow Enum.__eq__
     __eq__ = object.__eq__
 
-    # In Python 3, __hash__ must be defined if __eq__ is defined to stay hashable.
+    # In Python 3, __hash__ must be defined if __eq__ is defined to stay
+    # hashable.
     __hash__ = object.__hash__
 
     @classmethod
     def encode(
             cls,
             array: typing.Union[
                 EnumArray,
-                numpy.ndarray[int],
-                numpy.ndarray[str],
-                numpy.ndarray[Enum],
+                numpy.int_,
+                numpy.float_,
+                IndexedEnumArray,
                 ],
             ) -> EnumArray:
         """
-        Encode a string numpy array, an enum item numpy array, or an int numpy array
-        into an :any:`EnumArray`. See :any:`EnumArray.decode` for decoding.
+        Encode a string numpy array, an enum item numpy array, or an int numpy
+        array into an :any:`EnumArray`. See :any:`EnumArray.decode` for
+        decoding.
 
-        :param ndarray array: Array of string identifiers, or of enum items, to encode.
+        :param ndarray array: Array of string identifiers, or of enum items, to
+                              encode.
 
         :returns: An :any:`EnumArray` encoding the input array values.
         :rtype: :any:`EnumArray`
@@ -59,24 +66,31 @@ def encode(
         >>> encoded_array[0]
         2  # Encoded value
         """
-        if type(array) is EnumArray:
+        if isinstance(array, EnumArray):
             return array
 
-        if array.dtype.kind in {'U', 'S'}:  # String array
+        # String array
+        if isinstance(array, numpy.ndarray) and \
+                array.dtype.kind in {'U', 'S'}:
             array = numpy.select(
                 [array == item.name for item in cls],
                 [item.index for item in cls],
                 ).astype(config.ENUM_ARRAY_DTYPE)
 
-        elif array.dtype.kind == 'O':  # Enum items arrays
+        # Enum items arrays
+        elif isinstance(array, numpy.ndarray) and \
+                array.dtype.kind == 'O':
             # Ensure we are comparing the comparable. The problem this fixes:
             # On entering this method "cls" will generally come from
-            # variable.possible_values, while the array values may come from directly
-            # importing a module containing an Enum class. However, variables (and
-            # hence their possible_values) are loaded by a call to load_module, which
-            # gives them a different identity from the ones imported in the usual way.
-            # So, instead of relying on the "cls" passed in, we use only its name to
-            # check that the values in the array, if non-empty, are of the right type.
+            # variable.possible_values, while the array values may come from
+            # directly importing a module containing an Enum class. However,
+            # variables (and hence their possible_values) are loaded by a call
+            # to load_module, which gives them a different identity from the
+            # ones imported in the usual way.
+            #
+            # So, instead of relying on the "cls" passed in, we use only its
+            # name to check that the values in the array, if non-empty, are of
+            # the right type.
             if len(array) > 0 and cls.__name__ is array[0].__class__.__name__:
                 cls = array[0].__class__
 

openfisca_core/indexed_enums/enum_array.py

@@ -7,6 +7,8 @@
 if typing.TYPE_CHECKING:
     from openfisca_core.indexed_enums import Enum
 
+    IndexedEnumArray = numpy.object_
+
 
 class EnumArray(numpy.ndarray):
     """
@@ -19,24 +21,24 @@ class EnumArray(numpy.ndarray):
     # To read more about the two following methods, see:
     # https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html#slightly-more-realistic-example-attribute-added-to-existing-array.
     def __new__(
-            cls,
-            input_array: numpy.ndarray[int],
+            cls: typing.Type[EnumArray],
+            input_array: numpy.int_,
             possible_values: typing.Optional[typing.Type[Enum]] = None,
             ) -> EnumArray:
         obj = numpy.asarray(input_array).view(cls)
         obj.possible_values = possible_values
         return obj
 
     # See previous comment
-    def __array_finalize__(self, obj: typing.Optional[numpy.ndarray[int]]) -> None:
+    def __array_finalize__(self, obj: typing.Optional[numpy.int_]) -> None:
         if obj is None:
             return
 
         self.possible_values = getattr(obj, "possible_values", None)
 
     def __eq__(self, other: typing.Any) -> bool:
-        # When comparing to an item of self.possible_values, use the item index to
-        # speed up the comparison.
+        # When comparing to an item of self.possible_values, use the item index
+        # to speed up the comparison.
         if other.__class__.__name__ is self.possible_values.__name__:
             # Use view(ndarray) so that the result is a classic ndarray, not an
             # EnumArray.
@@ -49,8 +51,8 @@ def __ne__(self, other: typing.Any) -> bool:
 
     def _forbidden_operation(self, other: typing.Any) -> typing.NoReturn:
         raise TypeError(
-            "Forbidden operation. The only operations allowed on EnumArrays are "
-            "'==' and '!='.",
+            "Forbidden operation. The only operations allowed on EnumArrays "
+            "are '==' and '!='.",
             )
 
     __add__ = _forbidden_operation
@@ -62,7 +64,7 @@ def _forbidden_operation(self, other: typing.Any) -> typing.NoReturn:
     __and__ = _forbidden_operation
     __or__ = _forbidden_operation
 
-    def decode(self) -> numpy.ndarray[Enum]:
+    def decode(self) -> IndexedEnumArray:
         """
         Return the array of enum items corresponding to self.
 
@@ -72,14 +74,16 @@ def decode(self) -> numpy.ndarray[Enum]:
         >>> enum_array[0]
         >>> 2  # Encoded value
         >>> enum_array.decode()[0]
-        <HousingOccupancyStatus.free_lodger: 'Free lodger'>  # Decoded value : enum item
+        <HousingOccupancyStatus.free_lodger: 'Free lodger'>
+
+        Decoded value: enum item
         """
         return numpy.select(
             [self == item.index for item in self.possible_values],
             list(self.possible_values),
             )
 
-    def decode_to_str(self) -> numpy.ndarray[str]:
+    def decode_to_str(self) -> numpy.str_:
         """
         Return the array of string identifiers corresponding to self.