Add static type checking of tests using MyPy

.github/workflows/build.yaml

@@ -301,3 +301,18 @@ jobs:
           READTHEDOCS: 'True'
         run: |
             make html
+
+  check-types:
+    runs-on: ubuntu-24.04
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        submodules: true
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: Install Python dependencies
+      run: |
+        pip install -r requirements.txt -r requirements.dev.txt
+    - name: Run type checks
+      run: make check_types

Makefile

@@ -196,7 +196,7 @@ QEMU_ARCH ?= armv7emdp
 QEMU_FIRMWARE = $(QEMU_PORT_DIR)/build-$(QEMU_BOARD)/firmware.elf
 
 # Build firmware for QEMU
-.PHONY: qemu_build
+.PHONY: qemu_build check_types
 qemu_build:
 	$(MAKE) -C $(QEMU_PORT_DIR) BOARD=$(QEMU_BOARD) MICROPY_HEAP_SIZE=1024000
 
@@ -207,3 +207,13 @@ check_qemu: $(QEMU_FIRMWARE)
 	python3 $(abspath tools/run_qemu_tests.py) --board $(QEMU_BOARD) --arch $(QEMU_ARCH) --abi-version $(MPY_ABI_VERSION) --mount $(abspath .)
 
 
+# Type-check examples/ and tests/ against stubs/ with mypy.
+# Uses --follow-imports=skip to skip imports of third-party packages not
+# installed.
+# Excludes subdirs that have duplicate module names or 3rd-party deps.
+check_types:
+	MYPYPATH=./stubs python3 -m mypy \
+		--follow-imports=skip \
+		--exclude tests/tools/ \
+		tests/
+

docs/api_reference.rst

@@ -18,7 +18,7 @@ API reference
 
 
 .. _emlearn_trees:
-emlearn_trees - Decision tree ensembles
+emlearn_trees - Decision tree ensemble inference
 ----------------------------------------------------
 
 .. autoapimodule:: emlearn_trees
@@ -31,8 +31,27 @@ emlearn_linreg - Linear regression
 .. autoapimodule:: emlearn_linreg
    :members:
 
+.. _emlearn_linreg:
+emlearn_logreg - Logistic regression classification
+----------------------------------------------------
+
+.. autoapimodule:: emlearn_logreg
+   :members:
+
+emlearn_extratrees - Learning decision tree ensembles
+----------------------------------------------------
+
+.. autoapimodule:: emlearn_extratrees
+   :members:
+
+emlearn_plsr - Partial Least Squares Regression (PLSR)
+----------------------------------------------------
+
+.. autoapimodule:: emlearn_plsr
+   :members:
+
 .. _emlearn_cnn:
-emlearn_cnn - Convolutional Neural Networks
+emlearn_cnn - Convolutional Neural Networks inference
 ----------------------------------------------------------
 
 .. autoapimodule:: emlearn_cnn

requirements.dev.txt

@@ -5,3 +5,4 @@ sphinx_rtd_theme>=0.5.2
 sphinx-gallery>=0.10.1
 sphinx-autodoc-typehints>=3.0.1
 myst_parser>=4.0.0
+mypy>=2.0.0

stubs/emlearn_arrayutils.pyi

@@ -1,6 +1,4 @@
-
 # Stub file (PEP 484) with API definitions and documentation for native module
-# Is called .py because Sphinx autodoc currently does not support .pyi files
 
 """
 Array utility functions 

stubs/emlearn_cnn.pyi

@@ -1,5 +1,4 @@
 # Stub file (PEP 484) with API definitions and documentation for native module
-# Is called .py because Sphinx autodoc currently does not support .pyi files
 
 """
 Convolutional Neural Network module.

stubs/emlearn_cnn_fp32.pyi

@@ -0,0 +1 @@
+emlearn_cnn.pyi

stubs/emlearn_cnn_int8.pyi

@@ -0,0 +1 @@
+emlearn_cnn.pyi

stubs/emlearn_extratrees.pyi

@@ -0,0 +1,126 @@
+# Stub file (PEP 484) with API definitions and documentation for native module
+
+"""
+Extra Trees (Extremely Randomized Trees) classification.
+
+Tree-based ensemble classifier with randomized splits.
+"""
+
+import array
+import typing
+from typing import Iterator
+
+
+class Model():
+    """An ExtraTrees ensemble model
+
+    Note: Use emlearn_extratrees.new to construct an instance
+    """
+    def train(self, X : array.array, y : array.array) -> None:
+        """
+        Train the model on the given data.
+
+        :param X: Training features, int16 array of n_samples * n_features
+        :param y: Training labels, int16 array of n_samples
+        """
+        pass
+
+    def train_init(self, X : array.array, y : array.array) -> None:
+        """
+        Initialize step-by-step training.
+
+        :param X: Training features, int16 array of n_samples * n_features
+        :param y: Training labels, int16 array of n_samples
+        """
+        pass
+
+    def train_step(self) -> int:
+        """
+        Process one node of step-by-step training.
+
+        :return: 1 if training is complete, 0 if more steps needed
+        """
+        pass
+
+    def predict(self, features : array.array, probabilities : array.array) -> int:
+        """
+        Make a prediction and fill the probabilities buffer.
+
+        :param features: Input features, int16 array of n_features
+        :param probabilities: Output buffer, float32 array of n_classes (modified in-place)
+        :return: Predicted class index
+        """
+        pass
+
+    def get_n_features(self) -> int:
+        """
+        Get the number of features.
+        """
+        pass
+
+    def get_n_classes(self) -> int:
+        """
+        Get the number of classes.
+        """
+        pass
+
+    def get_n_trees(self) -> int:
+        """
+        Get the number of trees in the ensemble.
+        """
+        pass
+
+    def get_n_nodes_used(self) -> int:
+        """
+        Get the number of nodes currently used in the model.
+        """
+        pass
+
+    def get_n_trees_trained(self) -> int:
+        """
+        Get the number of trees trained so far.
+        """
+        pass
+
+    def __del__(self) -> None:
+        pass
+
+
+def new(n_features : int, n_classes : int,
+        *, n_trees : int = ..., max_depth : int = ...,
+        min_samples_leaf : int = ..., n_thresholds : int = ...,
+        subsample_ratio : float = ..., feature_subsample_ratio : float = ...,
+        max_nodes : int = ..., max_samples : int = ...,
+        rng_seed : int = ..., use_global_feature_range : bool = ...) -> Model:
+    """
+    Construct a new ExtraTrees model.
+
+    :param n_features: Number of input features
+    :param n_classes: Number of output classes
+    :param n_trees: Number of trees in the ensemble
+    :param max_depth: Maximum tree depth
+    :param min_samples_leaf: Minimum samples at a leaf node
+    :param n_thresholds: Random thresholds drawn per feature split
+    :param subsample_ratio: Fraction of samples used per tree (0.0-1.0)
+    :param feature_subsample_ratio: Fraction of features considered per split
+    :param max_nodes: Maximum pre-allocated nodes
+    :param max_samples: Maximum pre-allocated samples
+    :param rng_seed: Random number generator seed
+    :param use_global_feature_range: Use global feature range
+    """
+    pass
+
+
+def train_steps(model : Model, X : array.array, y : array.array) -> Iterator[int]:
+    """
+    Generator for step-by-step training.
+
+    Yields the number of trees trained so far after each step.
+    Returns when training is complete.
+
+    :param model: ExtraTrees model instance
+    :param X: Training features, int16 array of n_samples * n_features
+    :param y: Training labels, int16 array of n_samples
+    :return: Iterator yielding trees_trained count
+    """
+    pass

stubs/emlearn_fft.pyi

@@ -1,6 +1,5 @@
 
 # Stub file (PEP 484) with API definitions and documentation for native module
-# Is called .py because Sphinx autodoc currently does not support .pyi files
 
 """
 Fast Fourier Transform (FFT)

stubs/emlearn_iir.pyi

@@ -1,6 +1,5 @@
 
 # Stub file (PEP 484) with API definitions and documentation for native module
-# Is called .py because Sphinx autodoc currently does not support .pyi files
 
 """
 Infinite Impulse Response (IIR) filters 

stubs/emlearn_iir_q15.pyi

@@ -0,0 +1 @@
+emlearn_iir.pyi

stubs/emlearn_kmeans.pyi

@@ -0,0 +1,51 @@
+# Stub file (PEP 484) with API definitions and documentation for native module
+
+"""
+K-Means clustering with C-accelerated distance computation.
+"""
+
+import array
+import typing
+from typing import Iterator
+
+
+def euclidean_argmin(vectors : array.array, point : array.array) -> typing.Tuple[int, int]:
+    """
+    Find the closest centroid/vector to a given point.
+
+    :param vectors: All vectors concatenated, uint8 array of n_vectors * n_channels
+    :param point: Query point, uint8 array of n_channels
+    :return: Tuple of (closest_vector_index, squared_distance)
+    """
+    pass
+
+
+def cluster_iter(values : array.array, centroids : array.array,
+                 assignments : array.array, features : int,
+                 max_iter : int = ..., stop_changes : int = ...) -> Iterator[int]:
+    """
+    Perform K-Means clustering with iteration yielding.
+
+    :param values: Input data, uint8 array of n_samples * features
+    :param centroids: Initial centroids, uint8 array of n_clusters * features (modified in-place)
+    :param assignments: Output assignments, uint8 array of n_samples (modified in-place)
+    :param features: Number of features
+    :param max_iter: Maximum number of iterations
+    :param stop_changes: Stop if changes in assignments drop below this count
+    :return: Iterator yielding number of assignment changes per iteration
+    """
+    pass
+
+
+def cluster(values : array.array, centroids : array.array,
+            features : int, **kwargs) -> array.array:
+    """
+    Run K-Means clustering and return sample assignments.
+
+    :param values: Input data, uint8 array of n_samples * features
+    :param centroids: Initial centroids, uint8 array of n_clusters * features
+    :param features: Number of features
+    :param kwargs: Additional arguments (max_iter, stop_changes)
+    :return: Array of cluster assignments for each sample
+    """
+    pass

stubs/emlearn_linreg.pyi

@@ -1,5 +1,4 @@
 # Stub file (PEP 484) with API definitions and documentation for native module
-# Is called .py because Sphinx autodoc currently does not support .pyi files
 
 """
 Linear Regression with support for training/learning/fitting as well as inference/predictions.