feat: add input validation in lossless source coding classes

Author: rwnobrega

src/komm/_lossless_coding/LempelZiv77Code.py

@@ -4,6 +4,7 @@
 import numpy.typing as npt
 from tqdm import tqdm
 
+from .._util.validators import validate_integer_range
 from .util import integer_to_symbols, symbols_to_integer
 
 Token = tuple[int, int, int]
@@ -100,7 +101,7 @@ def source_to_tokens(self, source: npt.ArrayLike) -> list[Token]:
             >>> lz77.source_to_tokens([0, 0, 1, 0, 1, 0, 2, 1, 0, 2, 1, 0, 2, 1, 2])
             [(8, 2, 1), (7, 3, 2), (6, 7, 2)]
         """
-        source = np.asarray(source, dtype=int)
+        source = validate_integer_range(source, high=self.source_cardinality)
         ss, ls = self.search_size, self.lookahead_size
         buffer = bytes(self.search_buffer + source.tolist())
         tokens: list[Token] = []

src/komm/_lossless_coding/LempelZiv78Code.py

@@ -5,6 +5,8 @@
 import numpy.typing as npt
 from tqdm import tqdm
 
+from komm._util.validators import validate_integer_range
+
 from .util import Word, integer_to_symbols, symbols_to_integer
 
 Token = tuple[int, int]
@@ -46,7 +48,7 @@ def source_to_tokens(self, source: npt.ArrayLike) -> list[Token]:
             >>> lz78.source_to_tokens([1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0])
             [(0, 1), (0, 0), (1, 1), (2, 1), (4, 0), (2, 0)]
         """
-        source = np.asarray(source)
+        source = validate_integer_range(source, high=self.source_cardinality)
         dictionary: dict[Word, int] = {(): 0}
         tokens: list[Token] = []
         word: Word = ()

src/komm/_lossless_coding/LempelZivWelchCode.py

@@ -5,6 +5,8 @@
 import numpy.typing as npt
 from tqdm import tqdm
 
+from komm._util.validators import validate_integer_range
+
 from .util import Word, integer_to_symbols, symbols_to_integer
 
 
@@ -54,7 +56,7 @@ def encode(self, input: npt.ArrayLike) -> npt.NDArray[np.integer]:
             array([0, 2, 3, 4, 5])
         """
         calX, calY = self.source_cardinality, self.target_cardinality
-        input = np.asarray(input)
+        input = validate_integer_range(input, high=calX)
         dictionary: dict[Word, int] = {(s,): s for s in range(calX)}
         output: list[int] = []
 

src/komm/_lossless_coding/VariableToFixedCode.py

@@ -360,6 +360,11 @@ def encode(self, input: npt.ArrayLike) -> Array1D[np.integer]:
             >>> code.encode([1, 0, 0])  # Incomplete input, completed as 1|000
             array([1, 1, 0, 0])
 
+            >>> code.encode([0, 7, 0, 0])  # 07 is not a valid source word
+            Traceback (most recent call last):
+            ...
+            ValueError: input contains invalid word
+
             >>> code = komm.VariableToFixedCode.from_sourcewords(2, [(0, 0), (0, 1)])
             >>> code.encode([1, 0, 0, 0])  # Code is not fully covering
             Traceback (most recent call last):

src/komm/_util/validators.py

@@ -57,3 +57,15 @@ def validate_transition_matrix(
             f"({value.shape[0]}, {value.shape[1]}))"
         )
     return value
+
+
+def validate_integer_range(
+    value: npt.ArrayLike,
+    *,
+    low: int = 0,
+    high: int = 2,
+) -> npt.NDArray[np.integer]:
+    value = np.asarray(value, dtype=int)
+    if not (np.all(value >= low) and np.all(value < high)):
+        raise ValueError(f"input contains invalid entries (expected in [{low}:{high}))")
+    return value

tests/lossless_coding/test_lz77.py

@@ -140,3 +140,12 @@ def test_lz77_zeros(k):
     code = komm.LempelZiv77Code(window_size=32, lookahead_size=8, source_cardinality=2)
     source = np.zeros(k, dtype=int)
     np.testing.assert_equal(code.decode(code.encode(source)), source)
+
+
+def test_lz77_invalid_input():
+    code = komm.LempelZiv77Code(window_size=32, lookahead_size=8, source_cardinality=27)
+    code.encode([0, 10, 26])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([0, 10, 27])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([-1, 10, 26])

tests/lossless_coding/test_lz78.py

@@ -149,3 +149,12 @@ def test_lz78_worst_case(k):
     assert len(message) == (k - 1) * 2 ** (k + 1) + 2
     assert len(compressed) == len_compressed(2 ** (k + 1) - 2, 2)
     np.testing.assert_equal(code.decode(compressed), message)
+
+
+def test_lz78_invalid_input():
+    code = komm.LempelZiv78Code(source_cardinality=27)
+    code.encode([0, 10, 26])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([0, 10, 27])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([-1, 10, 26])

tests/lossless_coding/test_lzw.py

@@ -67,3 +67,12 @@ def test_lzw_encode_decode(source_cardinality, target_cardinality):
     # Random message
     message = np.random.randint(0, source_cardinality, 1000)
     np.testing.assert_equal(code.decode(code.encode(message)), message)
+
+
+def test_lzw_invalid_input():
+    code = komm.LempelZivWelchCode(source_cardinality=27)
+    code.encode([0, 10, 26])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([0, 10, 27])
+    with pytest.raises(ValueError, match="invalid entries"):
+        code.encode([-1, 10, 26])