format with black instead of ruff

Author: GiacomoPope

src/kyber_py/drbg/aes256_ctr_drbg.py

@@ -5,7 +5,9 @@
 
 
 class AES256_CTR_DRBG:
-    def __init__(self, seed: Optional[bytes] = None, personalization: bytes = b""):
+    def __init__(
+        self, seed: Optional[bytes] = None, personalization: bytes = b""
+    ):
         """
         DRBG implementation based on AES-256 CTR following the document NIST SP
         800-90A Section 10.2.1
@@ -57,7 +59,9 @@ def __instantiate(self, personalization: bytes = b"") -> bytes:
                 f"{self.seed_length}. Input has length {len(personalization)}"
             )
         # Ensure personalization has exactly seed_length bytes
-        personalization += bytes([0]) * (self.seed_length - len(personalization))
+        personalization += bytes([0]) * (
+            self.seed_length - len(personalization)
+        )
         # debugging
         assert len(personalization) == self.seed_length
         return xor_bytes(self.entropy_input, personalization)
@@ -93,7 +97,9 @@ def __ctr_drbg_update(self, provided_data: bytes) -> None:
         self.key = tmp[:32]
         self.V = tmp[32:]
 
-    def random_bytes(self, num_bytes: int, additional: Optional[bytes] = None) -> bytes:
+    def random_bytes(
+        self, num_bytes: int, additional: Optional[bytes] = None
+    ) -> bytes:
         """
         Generate pseudorandom bytes without a generating function
 

src/kyber_py/kyber/kyber.py

@@ -135,7 +135,9 @@ def _generate_matrix_from_seed(self, rho, transpose=False):
 
         When `transpose` is set to True, the matrix A is built as the transpose.
         """
-        A_data = [[self.R.zero() for _ in range(self.k)] for _ in range(self.k)]
+        A_data = [
+            [self.R.zero() for _ in range(self.k)] for _ in range(self.k)
+        ]
         for i in range(self.k):
             for j in range(self.k):
                 input_bytes = self._xof(rho, bytes([j]), bytes([i]))

src/kyber_py/ml_kem/ml_kem.py

@@ -130,15 +130,19 @@ def _G(s: bytes) -> tuple[bytes, bytes]:
         h = sha3_512(s).digest()
         return h[:32], h[32:]
 
-    def _generate_matrix_from_seed(self, rho: bytes, transpose: bool = False) -> Matrix:
+    def _generate_matrix_from_seed(
+        self, rho: bytes, transpose: bool = False
+    ) -> Matrix:
         """
         Helper function which generates a element of size
         k x k from a seed `rho`.
 
         When `transpose` is set to True, the matrix A is
         built as the transpose.
         """
-        A_data = [[self.R.zero() for _ in range(self.k)] for _ in range(self.k)]
+        A_data = [
+            [self.R.zero() for _ in range(self.k)] for _ in range(self.k)
+        ]
         for i in range(self.k):
             for j in range(self.k):
                 xof_bytes = self._xof(rho, bytes([j]), bytes([i]))
@@ -236,7 +240,9 @@ def _k_pke_encrypt(self, ek_pke: bytes, m: bytes, r: bytes) -> bytes:
 
         # Next check that t_hat has been canonically encoded
         if t_hat.encode(12) != t_hat_bytes:
-            raise ValueError("Modulus check failed, t_hat does not encode correctly")
+            raise ValueError(
+                "Modulus check failed, t_hat does not encode correctly"
+            )
 
         # Generate A_hat^T from seed rho
         A_hat_T = self._generate_matrix_from_seed(rho, transpose=True)

src/kyber_py/modules/modules.py

@@ -89,7 +89,9 @@ def decode_vector(self, input_bytes, k, d, is_ntt=False):
         # Ensure the input bytes are the correct length to create k elements with
         # d bits used for each coefficient
         if self.ring.n * d * k != len(input_bytes) * 8:
-            raise ValueError("Byte length is the wrong length for given k, d values")
+            raise ValueError(
+                "Byte length is the wrong length for given k, d values"
+            )
 
         # Bytes needed to decode a polynomial
         n = 32 * d

src/kyber_py/modules/modules_generic.py

@@ -54,7 +54,9 @@ def __getitem__(self, idx):
         """
         matrix[i, j] returns the element on row i, column j
         """
-        assert isinstance(idx, tuple) and len(idx) == 2, "Can't access individual rows"
+        assert (
+            isinstance(idx, tuple) and len(idx) == 2
+        ), "Can't access individual rows"
         if not self._transpose:
             return self._data[idx[0]][idx[1]]
         else:
@@ -64,7 +66,9 @@ def __eq__(self, other):
         if self.dim() != other.dim():
             return False
         m, n = self.dim()
-        return all([self[i, j] == other[i, j] for i in range(m) for j in range(n)])
+        return all(
+            [self[i, j] == other[i, j] for i in range(m) for j in range(n)]
+        )
 
     def __neg__(self):
         """
@@ -128,7 +132,10 @@ def __matmul__(self, other):
 
         return self.parent(
             [
-                [sum(self[i, k] * other[k, j] for k in range(n)) for j in range(l)]
+                [
+                    sum(self[i, k] * other[k, j] for k in range(n))
+                    for j in range(l)
+                ]
                 for i in range(m)
             ]
         )
@@ -149,10 +156,17 @@ def __repr__(self):
         if m == 1:
             return str(self._data[0])
 
-        max_col_width = [max(len(str(self[i, j])) for i in range(m)) for j in range(n)]
+        max_col_width = [
+            max(len(str(self[i, j])) for i in range(m)) for j in range(n)
+        ]
         info = "]\n[".join(
             [
-                ", ".join([f"{str(self[i, j]):>{max_col_width[j]}}" for j in range(n)])
+                ", ".join(
+                    [
+                        f"{str(self[i, j]):>{max_col_width[j]}}"
+                        for j in range(n)
+                    ]
+                )
                 for i in range(m)
             ]
         )
@@ -175,7 +189,9 @@ def random_element(self, m, n):
         :param int m: the number of columns in tge matrix
         :return: an element of the module with dimension `m times n`
         """
-        elements = [[self.ring.random_element() for _ in range(n)] for _ in range(m)]
+        elements = [
+            [self.ring.random_element() for _ in range(n)] for _ in range(m)
+        ]
         return self(elements)
 
     def __repr__(self):
@@ -192,14 +208,18 @@ def __call__(self, matrix_elements, transpose=False):
 
         if isinstance(matrix_elements[0], list):
             for element_list in matrix_elements:
-                if not all(isinstance(aij, self.ring.element) for aij in element_list):
+                if not all(
+                    isinstance(aij, self.ring.element) for aij in element_list
+                ):
                     raise TypeError(
                         f"All elements of the matrix must be elements of the ring: {self.ring}"
                     )
             return self.matrix(self, matrix_elements, transpose=transpose)
 
         elif isinstance(matrix_elements[0], self.ring.element):
-            if not all(isinstance(aij, self.ring.element) for aij in matrix_elements):
+            if not all(
+                isinstance(aij, self.ring.element) for aij in matrix_elements
+            ):
                 raise TypeError(
                     f"All elements of the matrix must be elements of the ring: {self.ring}"
                 )

src/kyber_py/polynomials/polynomials.py

@@ -16,7 +16,9 @@ def __init__(self):
         self.element_ntt = PolynomialNTT
 
         root_of_unity = 17
-        self.ntt_zetas = [pow(root_of_unity, self._br(i, 7), 3329) for i in range(128)]
+        self.ntt_zetas = [
+            pow(root_of_unity, self._br(i, 7), 3329) for i in range(128)
+        ]
         self.ntt_f = pow(128, -1, 3329)
 
     @staticmethod
@@ -212,7 +214,9 @@ def to_ntt(self):
         """
         Not supported, raises a ``TypeError``
         """
-        raise TypeError(f"Polynomial is already in the NTT domain: {type(self) = }")
+        raise TypeError(
+            f"Polynomial is already in the NTT domain: {type(self) = }"
+        )
 
     def from_ntt(self):
         """
@@ -281,7 +285,9 @@ def _ntt_multiplication(self, other):
         """
         Number Theoretic Transform multiplication.
         """
-        new_coeffs = self._ntt_coefficient_multiplication(self.coeffs, other.coeffs)
+        new_coeffs = self._ntt_coefficient_multiplication(
+            self.coeffs, other.coeffs
+        )
         return new_coeffs
 
     def __add__(self, other):

src/kyber_py/polynomials/polynomials_generic.py

@@ -124,13 +124,16 @@ def __neg__(self):
     def _add_(self, other):
         if isinstance(other, type(self)):
             new_coeffs = [
-                self._add_mod_q(x, y) for x, y in zip(self.coeffs, other.coeffs)
+                self._add_mod_q(x, y)
+                for x, y in zip(self.coeffs, other.coeffs)
             ]
         elif isinstance(other, int):
             new_coeffs = self.coeffs.copy()
             new_coeffs[0] = self._add_mod_q(new_coeffs[0], other)
         else:
-            raise NotImplementedError("Polynomials can only be added to each other")
+            raise NotImplementedError(
+                "Polynomials can only be added to each other"
+            )
         return new_coeffs
 
     def __add__(self, other):
@@ -147,7 +150,8 @@ def __iadd__(self, other):
     def _sub_(self, other):
         if isinstance(other, type(self)):
             new_coeffs = [
-                self._sub_mod_q(x, y) for x, y in zip(self.coeffs, other.coeffs)
+                self._sub_mod_q(x, y)
+                for x, y in zip(self.coeffs, other.coeffs)
             ]
         elif isinstance(other, int):
             new_coeffs = self.coeffs.copy()
@@ -211,7 +215,10 @@ def __eq__(self, other):
         if isinstance(other, type(self)):
             return self.coeffs == other.coeffs
         elif isinstance(other, int):
-            if self.is_constant() and (other % self.parent.q) == self.coeffs[0]:
+            if (
+                self.is_constant()
+                and (other % self.parent.q) == self.coeffs[0]
+            ):
                 return True
         return False
 

tests/test_drbg.py

@@ -35,7 +35,9 @@ def test_bad_personalization(self):
         # if the personalization is longer than 48 bytes, fail
         seed = os.urandom(48)
         personalization = os.urandom(49)
-        self.assertRaises(ValueError, lambda: AES256_CTR_DRBG(seed, personalization))
+        self.assertRaises(
+            ValueError, lambda: AES256_CTR_DRBG(seed, personalization)
+        )
 
     def test_additional(self):
         drbg = AES256_CTR_DRBG()
@@ -47,4 +49,6 @@ def test_bad_additional(self):
         # if the additional data is longer than 48 bytes, fail
         drbg = AES256_CTR_DRBG()
         additional = os.urandom(49)
-        self.assertRaises(ValueError, lambda: drbg.random_bytes(32, additional))
+        self.assertRaises(
+            ValueError, lambda: drbg.random_bytes(32, additional)
+        )

tests/test_kyber.py

@@ -10,7 +10,9 @@ def parse_kat_data(data):
     count_blocks = data.split("\n\n")
     for block in count_blocks[:-1]:
         block_data = block.split("\n")
-        count, seed, pk, sk, ct, ss = [line.split(" = ")[-1] for line in block_data]
+        count, seed, pk, sk, ct, ss = [
+            line.split(" = ")[-1] for line in block_data
+        ]
         parsed_data[int(count)] = {
             "seed": bytes.fromhex(seed),
             "pk": bytes.fromhex(pk),

tests/test_ml_kem.py

@@ -51,7 +51,9 @@ def test_encaps_modulus_check_failure(self):
         self.assertRaises(ValueError, lambda: ML_KEM_512.encaps(bad_ek))
 
     def test_xof_failure(self):
-        self.assertRaises(ValueError, lambda: ML_KEM_512._xof(b"1", b"2", b"3"))
+        self.assertRaises(
+            ValueError, lambda: ML_KEM_512._xof(b"1", b"2", b"3")
+        )
 
     def test_prf_failure(self):
         self.assertRaises(ValueError, lambda: ML_KEM_512._prf(2, b"1", b"2"))
@@ -256,7 +258,9 @@ def generic_keygen_kat(self, ML_KEM, index):
             self.assertEqual(dk, dk_kat)
 
     def generic_encap_kat(self, ML_KEM, index):
-        with open("assets/ML-KEM-encapDecap-FIPS203/internalProjection.json") as f:
+        with open(
+            "assets/ML-KEM-encapDecap-FIPS203/internalProjection.json"
+        ) as f:
             data = json.load(f)
         kat_data = data["testGroups"][index]["tests"]
 
@@ -275,7 +279,9 @@ def generic_encap_kat(self, ML_KEM, index):
             self.assertEqual(K_prime, k_kat)
 
     def generic_decap_kat(self, ML_KEM, index):
-        with open("assets/ML-KEM-encapDecap-FIPS203/internalProjection.json") as f:
+        with open(
+            "assets/ML-KEM-encapDecap-FIPS203/internalProjection.json"
+        ) as f:
             data = json.load(f)
         kat_data = data["testGroups"][3 + index]["tests"]