Fix type hinting (#40)

Enhance type hinting

Author: mathias-nillion

examples/dot_product/tests/dot_product.yaml

@@ -16,5 +16,5 @@ inputs:
       SecretInteger: "3"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "27"

nada_numpy/funcs.py

@@ -6,12 +6,12 @@
 from typing import Any, Callable, List, Sequence, Tuple, Union
 
 import numpy as np
-from nada_dsl import (Boolean, Integer, Party, PublicInteger,
+from nada_dsl import (Boolean, Integer, Output, Party, PublicInteger,
                       PublicUnsignedInteger, SecretInteger,
                       SecretUnsignedInteger, UnsignedInteger)
 
 from nada_numpy.array import NadaArray
-from nada_numpy.nada_typing import NadaCleartextNumber
+from nada_numpy.nada_typing import AnyNadaType, NadaCleartextNumber
 from nada_numpy.types import Rational, SecretRational, rational
 from nada_numpy.utils import copy_metadata
 
@@ -60,7 +60,7 @@
 ]
 
 
-def parties(num: int, prefix: str = "Party") -> list:
+def parties(num: int, prefix: str = "Party") -> List[Party]:
     """
     Create a list of Party objects with specified names.
 
@@ -69,12 +69,12 @@ def parties(num: int, prefix: str = "Party") -> list:
         prefix (str, optional): The prefix to use for party names. Defaults to "Party".
 
     Returns:
-        list: A list of Party objects with names in the format "{prefix}{i}".
+        List[Party]: A list of Party objects with names in the format "{prefix}{i}".
     """
     return [Party(name=f"{prefix}{i}") for i in range(num)]
 
 
-def __from_numpy(arr: np.ndarray, nada_type: NadaCleartextNumber) -> list:
+def __from_numpy(arr: np.ndarray, nada_type: NadaCleartextNumber) -> List:
     """
     Recursively convert a n-dimensional NumPy array to a nested list of NadaInteger objects.
 
@@ -83,7 +83,7 @@ def __from_numpy(arr: np.ndarray, nada_type: NadaCleartextNumber) -> list:
         nada_type (type): The type of NadaInteger objects to create.
 
     Returns:
-        list: A nested list of NadaInteger objects.
+        List: A nested list of NadaInteger objects.
     """
     if len(arr.shape) == 1:
         if isinstance(nada_type, Rational):
@@ -261,20 +261,24 @@ def random(
     return NadaArray.random(dims, nada_type)
 
 
-def output(arr: NadaArray, party: Party, prefix: str):
+def output(value: Union[NadaArray, AnyNadaType], party: Party, prefix: str) -> List[Output]:
     """
-    Generate a list of Output objects for each element in the input NadaArray.
+    Generate a list of Output objects for some provided value.
 
     Args:
-        arr (NadaArray): The input NadaArray.
+        value (Union[NadaArray, AnyNadaType]): The input NadaArray.
         party (Party): The party object.
         prefix (str): The prefix for naming the Output objects.
 
     Returns:
-        list: A list of Output objects.
-    """
-    # pylint:disable=protected-access
-    return NadaArray._output_array(arr, party, prefix)
+        List[Output]: A list of Output objects.
+    """
+    if isinstance(value, NadaArray):
+        # pylint:disable=protected-access
+        return NadaArray._output_array(value, party, prefix)
+    if isinstance(value, (Rational, SecretRational)):
+        value = value.value
+    return [Output(value, prefix, party)]
 
 
 def vstack(arr_list: list) -> NadaArray:

nada_numpy/nada_typing.py

@@ -44,3 +44,9 @@
     dsl.Boolean,
     Rational,
 ]
+
+AnyNadaType = Union[
+    dsl.NadaType,
+    Rational,
+    SecretRational,
+]

tests/nada-tests/tests/array_statistics.yaml

@@ -32,23 +32,23 @@ expected_outputs:
     SecretInteger: "4"
   b_mean_arr_0:
     SecretInteger: "3"
-  b_sum_0:
+  b_sum:
     SecretInteger: "21"
-  b_mean_0:
+  b_mean:
     SecretInteger: "3"
   a_mean_arr_0:
     SecretInteger: "3"
   b_mean_arr_1:
     SecretInteger: "4"
   b_sum_arr_0:
     SecretInteger: "9"
-  a_sum_0:
+  a_sum:
     SecretInteger: "21"
   a_sum_arr_0:
     SecretInteger: "9"
   b_sum_arr_1:
     SecretInteger: "12"
-  a_mean_0:
+  a_mean:
     SecretInteger: "3"
   a_sum_arr_1:
     SecretInteger: "12"

tests/nada-tests/tests/dot_product.yaml

@@ -16,5 +16,5 @@ inputs:
       SecretInteger: "3"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "27"

tests/nada-tests/tests/dot_product_rational.yaml

@@ -16,11 +16,11 @@ inputs:
       SecretInteger: "196608"
   public_variables: {}
 expected_outputs:
-  my_output_a_0:
+  my_output_a:
     SecretInteger: "917504"
-  my_output_b_0:
+  my_output_b:
     SecretInteger: "917504"
-  my_output_c_0:
+  my_output_c:
     SecretInteger: "393216"
-  my_output_d_0:
+  my_output_d:
     SecretInteger: "393216"

tests/nada-tests/tests/get_attr.yaml

@@ -10,5 +10,5 @@ inputs:
       SecretInteger: "3"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "9"

tests/nada-tests/tests/get_item.yaml

@@ -10,5 +10,5 @@ inputs:
       SecretInteger: "3"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "9"

tests/nada-tests/tests/logistic_regression.yaml

@@ -18,5 +18,5 @@ inputs:
       SecretInteger: "65536"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "12884967424"

tests/nada-tests/tests/rational_if_else.yaml

@@ -14,29 +14,29 @@ inputs:
       SecretUnsignedInteger: "1"
   public_variables: {}
 expected_outputs:
-  out_0_0:
+  out_0:
     SecretInteger: "294912"
-  out_1_0:
+  out_1:
     SecretInteger: "294912"
-  out_2_0:
+  out_2:
     SecretInteger: "78643"
-  out_3_0:
+  out_3:
     SecretInteger: "78643"
-  out_4_0:
+  out_4:
     SecretInteger: "1"
-  out_5_0:
+  out_5:
     SecretInteger: "1"
-  out_6_0:
+  out_6:
     SecretInteger: "1"
-  out_7_0:
+  out_7:
     SecretInteger: "1"
-  out_8_0:
+  out_8:
     Integer: "131072"
-  out_9_0:
+  out_9:
     Integer: "65536"
-  out_10_0:
+  out_10:
     Integer: "1"
-  out_11_0:
+  out_11:
     SecretInteger: "1"
-  out_12_0:
+  out_12:
     SecretUnsignedInteger: "1"

tests/nada-tests/tests/sum.yaml

@@ -10,5 +10,5 @@ inputs:
       SecretInteger: "3"
   public_variables: {}
 expected_outputs:
-  my_output_0:
+  my_output:
     SecretInteger: "9"

tests/nada-tests/tests/supported_operations.yaml

@@ -80,15 +80,15 @@ inputs:
       SecretInteger: "1"
   public_variables: {}
 expected_outputs:
-  out_0_0:
+  out_0:
     SecretInteger: 1
-  out_1_0:
+  out_1:
     SecretInteger: 362880
-  out_2_0:
+  out_2:
     SecretInteger: 31
-  out_3_0:
+  out_3:
     SecretInteger: 13
-  out_4_0:
+  out_4:
     SecretInteger: 21
-  out_5_0:
+  out_5:
     Integer: 42

tests/nada-tests/tests/supported_operations_return_types.yaml

@@ -80,15 +80,15 @@ inputs:
       SecretInteger: "1"
   public_variables: {}
 expected_outputs:
-  out_0_0:
+  out_0:
     SecretInteger: 1
-  out_1_0:
+  out_1:
     SecretInteger: 362880
-  out_2_0:
+  out_2:
     SecretInteger: 31
-  out_3_0:
+  out_3:
     SecretInteger: 13
-  out_4_0:
+  out_4:
     SecretInteger: 21
-  out_5_0:
+  out_5:
     Integer: 42