Add type hints for misc functions

galois/_lfsr.py

@@ -1,9 +1,12 @@
 """
 A module containing classes and functions for generating and analyzing linear feedback shift registers and their sequences.
 """
+from typing import Tuple, List, Union, overload
+from typing_extensions import Literal
+
 import numba
-from numba import int64
 import numpy as np
+from numba import int64
 
 from ._fields import FieldClass, FieldArray
 from ._overrides import set_module
@@ -58,7 +61,12 @@ class LFSR:
  * https://www.cs.uky.edu/~klapper/pdf/galois.pdf
  """
 
- def __init__(self, poly, state=1, config="fibonacci"):
+ def __init__(
+ self,
+ poly: Poly,
+ state: Union[int, Tuple[int], List[int], np.ndarray, FieldArray] = 1,
+ config: Literal["fibonacci", "galois"] = "fibonacci"
+ ):
  r"""
  Constructs a linear-feedback shift register.
 
@@ -80,7 +88,7 @@ def __init__(self, poly, state=1, config="fibonacci"):
  """
  if not isinstance(poly, Poly):
  raise TypeError(f"Argument `poly` must be a galois.Poly, not {type(poly)}.")
- if not isinstance(state, (type(None), int, np.integer, tuple, list, np.ndarray, FieldArray)):
+ if not isinstance(state, (int, np.integer, tuple, list, np.ndarray, FieldArray)):
  raise TypeError(f"Argument `state` must be an int or array-like, not {type(state)}.")
  if not isinstance(config, str):
  raise TypeError(f"Argument `config` must be a string, not {type(config)}.")
@@ -94,7 +102,7 @@ def __init__(self, poly, state=1, config="fibonacci"):
  self._field = poly.field
  self._poly = poly
  self._initial_state = self.field(state)
- self._state = None
+ self._state = self.field(state)
  self._config = config
  self.reset()
 
@@ -134,7 +142,7 @@ def reset(self):
  """
  self._state = self.initial_state.copy()
 
- def step(self, steps=1):
+ def step(self, steps: int = 1) -> FieldArray:
  """
  Steps the LFSR and produces `steps` output symbols.
 
@@ -191,7 +199,7 @@ def step(self, steps=1):
  return y
 
  @property
- def field(self):
+ def field(self) -> FieldClass:
  """
  galois.FieldClass: The Galois field that defines the LFSR arithmetic. The generator polynomial :math:`g(x)` is over this
  field and the state vector contains values in this field.
@@ -207,7 +215,7 @@ def field(self):
  return self._field
 
  @property
- def poly(self):
+ def poly(self) -> Poly:
  r"""
  galois.Poly: The generator polynomial :math:`g(x) = g_n x^n + \dots + g_1 x + g_0`.
 
@@ -221,7 +229,7 @@ def poly(self):
  return self._poly
 
  @property
- def initial_state(self):
+ def initial_state(self) -> FieldArray:
  r"""
  galois.FieldArray: The initial state vector :math:`s = [s_{n-1}, \dots, s_1, s_0]`.
 
@@ -235,7 +243,7 @@ def initial_state(self):
  return self._initial_state.copy()
 
  @property
- def state(self):
+ def state(self) -> FieldArray:
  r"""
  galois.FieldArray: The current state vector :math:`s = [s_{n-1}, \dots, s_1, s_0]`.
 
@@ -251,7 +259,7 @@ def state(self):
  return self._state.copy()
 
  @property
- def config(self):
+ def config(self) -> str:
  """
  str: The LFSR configuration, either `"fibonacci"` or `"galois"`. See the Notes section of :obj:`LFSR` for descriptions of
  the two configurations.
@@ -285,6 +293,12 @@ def config(self):
 # return Poly(coeffs, field=field)
 
 
+@overload
+def berlekamp_massey(sequence: FieldArray, config: Literal["fibonacci", "galois"] = "fibonacci", state: Literal[False] = False) -> Poly:
+ ...
+@overload
+def berlekamp_massey(sequence: FieldArray, config: Literal["fibonacci", "galois"] = "fibonacci", state: Literal[True] = True) -> Tuple[Poly, FieldArray]:
+ ...
 @set_module("galois")
 def berlekamp_massey(sequence, config="fibonacci", state=False):
  r"""

galois/_ntt.py

@@ -1,6 +1,8 @@
 """
 A module that contains functions to perform the forward and reverse Number-Theoretic Transform (NTT).
 """
+from typing import Tuple, List, Optional, Union
+
 import numpy as np
 
 from ._factory import GF
@@ -13,7 +15,11 @@
 
 
 @set_module("galois")
-def ntt(x, size=None, modulus=None):
+def ntt(
+ x: Union[Tuple[int], List[int], np.ndarray, FieldArray],
+ size: Optional[int] = None,
+ modulus: Optional[int] = None
+) -> FieldArray:
  r"""
  Computes the Number-Theoretic Transform (NTT) of :math:`x`.
 
@@ -109,7 +115,12 @@ def ntt(x, size=None, modulus=None):
 
 
 @set_module("galois")
-def intt(X, size=None, modulus=None, scaled=True):
+def intt(
+ X: Union[Tuple[int], List[int], np.ndarray, FieldArray],
+ size: Optional[int] = None,
+ modulus: Optional[int] = None,
+ scaled: bool = True
+) -> FieldArray:
  r"""
  Computes the Inverse Number-Theoretic Transform (INTT) of :math:`X`.
 
@@ -205,7 +216,7 @@ def intt(X, size=None, modulus=None, scaled=True):
 
  galois.ntt(galois.intt(x))
  """
- if not isinstance(X, (tuple, list, np.ndarray)):
+ if not isinstance(X, (tuple, list, np.ndarray, FieldArray)):
  raise TypeError(f"Argument `X` must be array-like, not {type(X)}.")
  if isinstance(X, FieldArray) and not type(X).is_prime_field:
  raise ValueError(f"If argument `X` is a Galois field array, it must be a prime field, not {type(X)}.")