Feature/cleanup and guardrails (#34)

Add typing guardrails

Author: mathias-nillion

nada_algebra/array.py

@@ -5,15 +5,16 @@
 
 # pylint:disable=too-many-lines
 
-from types import NoneType
-from typing import Any, Callable, Optional, Sequence, Type, Union
+from typing import Any, Callable, Optional, Sequence, Union, get_args, overload
 
 import numpy as np
 from nada_dsl import (Input, Integer, Output, Party, PublicInteger,
                       PublicUnsignedInteger, SecretInteger,
                       SecretUnsignedInteger, UnsignedInteger)
 
 from nada_algebra.context import UnsafeArithmeticSession
+from nada_algebra.nada_typing import (NadaBoolean, NadaInteger, NadaRational,
+                                      NadaUnsignedInteger)
 from nada_algebra.types import (Rational, SecretRational, get_log_scale,
                                 public_rational, rational, secret_rational)
 from nada_algebra.utils import copy_metadata
@@ -22,9 +23,6 @@
 class NadaArray:  # pylint:disable=too-many-public-methods
     """
     Represents an array-like object with additional functionality.
-
-    Attributes:
-        inner (np.ndarray): The underlying NumPy array.
     """
 
     def __init__(self, inner: np.ndarray):
@@ -41,6 +39,7 @@ def __init__(self, inner: np.ndarray):
             raise ValueError(f"inner must be a numpy array and is: {type(inner)}")
         if isinstance(inner, NadaArray):
             inner = inner.inner
+        _check_type_conflicts(inner)
         self.inner = inner
 
     def __getitem__(self, item):
@@ -65,7 +64,11 @@ def __setitem__(self, key, value):
         Args:
             key: The key to set.
             value: The value to set.
+
+        Raises:
+            ValueError: Raised when value with incompatible type is passed.
         """
+        _check_type_compatibility(value, self.dtype)
         if isinstance(value, NadaArray):
             self.inner[key] = value.inner
         else:
@@ -462,7 +465,7 @@ def apply(self, func: Callable[[Any], Any]) -> "NadaArray":
     def mean(self, axis=None, dtype=None, out=None) -> Any:
         sum_arr = self.inner.sum(axis=axis, dtype=dtype)
 
-        if self.dtype in (Rational, SecretRational):
+        if self.is_rational:
             nada_type = rational
         else:
             nada_type = Integer
@@ -483,7 +486,7 @@ def mean(self, axis=None, dtype=None, out=None) -> Any:
         return mean_arr
 
     @staticmethod
-    def output_array(array: Any, party: Party, prefix: str) -> list:
+    def _output_array(array: Any, party: Party, prefix: str) -> list:
         """
         Generate a list of Output objects for each element in the input array.
 
@@ -512,7 +515,7 @@ def output_array(array: Any, party: Party, prefix: str) -> list:
 
         if len(array.shape) == 0:
             # For compatibility we're leaving this here.
-            return NadaArray.output_array(array.item(), party, prefix)
+            return NadaArray._output_array(array.item(), party, prefix)
         if len(array.shape) == 1:
             return [
                 (
@@ -525,7 +528,7 @@ def output_array(array: Any, party: Party, prefix: str) -> list:
         return [
             v
             for i in range(array.shape[0])
-            for v in NadaArray.output_array(array[i], party, f"{prefix}_{i}")
+            for v in NadaArray._output_array(array[i], party, f"{prefix}_{i}")
         ]
 
     def output(self, party: Party, prefix: str) -> list:
@@ -539,10 +542,10 @@ def output(self, party: Party, prefix: str) -> list:
         Returns:
             list: A list of Output objects.
         """
-        return NadaArray.output_array(self.inner, party, prefix)
+        return NadaArray._output_array(self.inner, party, prefix)
 
     @staticmethod
-    def create_list(
+    def _create_list(
         dims: Sequence[int],
         party: Optional[Party],
         prefix: Optional[str],
@@ -564,7 +567,7 @@ def create_list(
         if len(dims) == 1:
             return [generator(f"{prefix}_{i}", party) for i in range(dims[0])]
         return [
-            NadaArray.create_list(
+            NadaArray._create_list(
                 dims[1:],
                 party,
                 f"{prefix}_{i}",
@@ -621,7 +624,7 @@ def array(
             raise ValueError(f"Unsupported nada_type: {nada_type}")
 
         return NadaArray(
-            np.array(NadaArray.create_list(dims, party, prefix, generator))
+            np.array(NadaArray._create_list(dims, party, prefix, generator))
         )
 
     @staticmethod
@@ -658,7 +661,7 @@ def random(
         else:
             raise ValueError(f"Unsupported nada_type: {nada_type}")
 
-        return NadaArray(np.array(NadaArray.create_list(dims, None, None, generator)))
+        return NadaArray(np.array(NadaArray._create_list(dims, None, None, generator)))
 
     def __len__(self):
         """
@@ -680,27 +683,58 @@ def empty(self) -> bool:
         return len(self.inner) == 0
 
     @property
-    def dtype(self) -> Type:
+    def dtype(
+        self,
+    ) -> Optional[Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]]:
         """
-        Gets inner data type of NadaArray values.
+        Gets data type of array.
 
         Returns:
-            Type: Inner data type.
+            Optional[
+                Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]
+            ]: Array data type if applicable.
         """
-        # TODO: account for mixed typed NadaArrays due to e.g. padding
-        if self.empty:
-            return NoneType
-        return type(self.inner.item(0))
+        return get_dtype(self.inner)
 
     @property
     def is_rational(self) -> bool:
         """
-        Returns whether or not the Array's type is a rational.
+        Returns whether or not the Array type contains rationals.
+
+        Returns:
+            bool: Boolean output.
+        """
+        return self.dtype == NadaRational
+
+    @property
+    def is_integer(self) -> bool:
+        """
+        Returns whether or not the Array type contains signed integers.
+
+        Returns:
+            bool: Boolean output.
+        """
+        return self.dtype == NadaInteger
+
+    @property
+    def is_unsigned_integer(self) -> bool:
+        """
+        Returns whether or not the Array type contains unsigned integers.
+
+        Returns:
+            bool: Boolean output.
+        """
+        return self.dtype == NadaUnsignedInteger
+
+    @property
+    def is_boolean(self) -> bool:
+        """
+        Returns whether or not the Array type contains signed integers.
 
         Returns:
             bool: Boolean output.
         """
-        return self.dtype in (Rational, SecretRational)
+        return self.dtype == NadaBoolean
 
     def __str__(self) -> str:
         """
@@ -817,9 +851,23 @@ def item(self, *args, **kwargs):
             return NadaArray(result)
         return result
 
+    @overload
+    def itemset(self, value: Any): ...
+    @overload
+    def itemset(self, item: Any, value: Any): ...
+
     # pylint:disable=missing-function-docstring
     @copy_metadata(np.ndarray.itemset)
     def itemset(self, *args, **kwargs):
+        value = None
+        if len(args) == 1:
+            value = args[0]
+        elif len(args) == 2:
+            value = args[1]
+        else:
+            value = kwargs["value"]
+
+        _check_type_compatibility(value, self.dtype)
         result = self.inner.itemset(*args, **kwargs)
         if isinstance(result, np.ndarray):
             return NadaArray(result)
@@ -835,9 +883,12 @@ def prod(self, *args, **kwargs):
 
     # pylint:disable=missing-function-docstring
     @copy_metadata(np.ndarray.put)
-    def put(self, *args, **kwargs):
-        result = self.inner.put(*args, **kwargs)
-        return result
+    def put(self, ind: Any, v: Any, mode: Any = None) -> None:
+        _check_type_compatibility(v, self.dtype)
+        if isinstance(v, NadaArray):
+            self.inner.put(ind, v.inner, mode)
+        else:
+            self.inner.put(ind, v, mode)
 
     # pylint:disable=missing-function-docstring
     @copy_metadata(np.ndarray.ravel)
@@ -1007,3 +1058,76 @@ def T(self):  # pylint:disable=invalid-name
         if isinstance(result, np.ndarray):
             return NadaArray(result)
         return result
+
+
+def _check_type_compatibility(
+    value: Any,
+    check_type: Optional[
+        Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]
+    ],
+) -> None:
+    """
+    Checks type compatibility between a type and a Nada base type.
+
+    Args:
+        value (Any): Value to be type-checked.
+        check_type (Optional[
+            Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]
+        ]): Base Nada type to check against.
+
+    Raises:
+        TypeError: Raised when types are not compatible.
+    """
+    if isinstance(value, (NadaArray, np.ndarray)):
+        if isinstance(value, NadaArray):
+            value = value.inner
+        dtype = get_dtype(value)
+        if dtype is None or check_type is None:
+            raise TypeError(f"Type {dtype} is not compatible with {check_type}")
+        if dtype == check_type:
+            return
+    else:
+        dtype = type(value)
+        if dtype in get_args(check_type):
+            return
+    raise TypeError(f"Type {dtype} is not compatible with {check_type}")
+
+
+def _check_type_conflicts(array: np.ndarray) -> None:
+    """
+    Checks for type conflicts
+
+    Args:
+        array (np.ndarray): Array to be checked.
+
+    Raises:
+        TypeError: Raised when incompatible dtypes are detected.
+    """
+    _ = get_dtype(array)
+
+
+def get_dtype(
+    array: np.ndarray,
+) -> Optional[Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]]:
+    """
+    Gets all data types present in array.
+
+    Args:
+        array (np.ndarray): Array to be checked.
+
+    Raises:
+        TypeError: Raised when incompatible dtypes are detected.
+
+    Returns:
+        Optional[Union[NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]: Array dtype].
+    """
+    if array.size == 0:
+        return None
+
+    unique_types = set(type(element) for element in array.flat)
+
+    base_dtypes = [NadaRational, NadaInteger, NadaUnsignedInteger, NadaBoolean]
+    for base_dtype in base_dtypes:
+        if all(unique_type in get_args(base_dtype) for unique_type in unique_types):
+            return base_dtype
+    raise TypeError(f"Nada-incompatible dtypes detected in `{unique_types}`.")

nada_algebra/funcs.py

@@ -6,8 +6,9 @@
 from typing import Any, Callable, List, Sequence, Tuple, Union
 
 import numpy as np
-from nada_dsl import (Integer, Party, PublicInteger, PublicUnsignedInteger,
-                      SecretInteger, SecretUnsignedInteger, UnsignedInteger)
+from nada_dsl import (Boolean, Integer, Party, PublicInteger,
+                      PublicUnsignedInteger, SecretInteger,
+                      SecretUnsignedInteger, UnsignedInteger)
 
 from nada_algebra.array import NadaArray
 from nada_algebra.types import Rational, SecretRational, rational
@@ -229,7 +230,8 @@ def output(arr: NadaArray, party: Party, prefix: str):
     Returns:
         list: A list of Output objects.
     """
-    return NadaArray.output_array(arr, party, prefix)
+    # pylint:disable=protected-access
+    return NadaArray._output_array(arr, party, prefix)
 
 
 def vstack(arr_list: list) -> NadaArray:
@@ -344,19 +346,18 @@ def pad(
 
     # Override python defaults by NadaType defaults
     overriden_kwargs = {}
-    if mode == "constant":
-        dtype = arr.dtype
-        if dtype in (Rational, SecretRational):
-            nada_type = rational
-        elif dtype in (PublicInteger, SecretInteger):
-            nada_type = Integer
-        elif dtype == (PublicUnsignedInteger, SecretUnsignedInteger):
-            nada_type = UnsignedInteger
+    if mode == "constant" and "constant_values" not in kwargs:
+        if arr.is_rational:
+            default = rational(0)
+        elif arr.is_integer:
+            default = Integer(0)
+        elif arr.is_unsigned_integer:
+            default = UnsignedInteger(0)
         else:
-            nada_type = dtype
+            default = Boolean(False)
 
         overriden_kwargs["constant_values"] = kwargs.get(
-            "constant_values", nada_type(0)
+            "constant_values", default
         )
 
     padded_inner = np.pad(  # type: ignore

nada_algebra/nada_typing.py

@@ -0,0 +1,33 @@
+"""Contains custom typing traits"""
+
+from typing import Union
+
+import nada_dsl as dsl
+
+from nada_algebra.types import (PublicBoolean, Rational, SecretBoolean,
+                                SecretRational)
+
+NadaRational = Union[
+    Rational,
+    SecretRational,
+]
+
+NadaInteger = Union[
+    dsl.Integer,
+    dsl.PublicInteger,
+    dsl.SecretInteger,
+]
+
+NadaUnsignedInteger = Union[
+    dsl.UnsignedInteger,
+    dsl.PublicUnsignedInteger,
+    dsl.SecretUnsignedInteger,
+]
+
+NadaBoolean = Union[
+    dsl.Boolean,
+    dsl.PublicBoolean,
+    dsl.SecretBoolean,
+    PublicBoolean,
+    SecretBoolean,
+]