repaired KNNClassifier cleaning behaviour, tested with a simple script

docs/releases/unreleased.md

@@ -33,3 +33,9 @@ The `dummy` module is now fully type-annotated.
 ## tree
 
 - Added handling for division by zero in `tree.hoeffding_tree` for leaf size estimation.
+
+## neighbors
+
+- Added function in nearest-neighbor engines to gather relevant classes/targets from the window.
+- Added a virtual function to the base engine class; New NN engines need to override `refresh_targets` function
+  - Classifier KNN now calls this engine-specific function under `clean_up_classes()`

river/neighbors/ann/swinn.py

@@ -491,3 +491,17 @@ def connectivity(self) -> list[int]:
                 forest.add(frozenset(u))
 
         return [len(tree) for tree in forest]
+
+    def refresh_targets(self) -> set:
+        """Refresh the set of classes in the window. Used by classifiers where labels are added as [1] in the vertex tuple.
+
+        This is used to clean up classes that are no longer in the window, and
+        ensure we do not consider "None" a class. It is called every `cleanup_every`
+        step, or can be called manually.
+
+        """
+        return {
+            vertex.item[1]
+            for vertex in self._data
+            if vertex is not None and vertex.item[1] is not None
+        }

river/neighbors/base.py

@@ -43,3 +43,7 @@ def append(self, item: typing.Any, **kwargs) -> None:
     @abc.abstractmethod
     def search(self, item: typing.Any, n_neighbors: int, **kwargs) -> tuple[list, list]:
         pass
+
+    @abc.abstractmethod
+    def refresh_targets(self) -> set:
+        pass

river/neighbors/knn_classifier.py

@@ -122,7 +122,7 @@ def clean_up_classes(self):
         `cleanup_every` step, or can be called manually.
 
         """
-        self.classes = {x for x in self.window if x[0][1] is not None}
+        self.classes = self._nn.refresh_targets()
 
     def learn_one(self, x, y):
         # Update the data buffer

river/neighbors/lazy.py

@@ -123,3 +123,15 @@ def search(self, item: typing.Any, n_neighbors: int, **kwargs):
 
         # Return the k closest points
         return tuple(map(list, zip(*sorted(points, key=operator.itemgetter(-1))[:n_neighbors])))
+
+    def refresh_targets(self) -> set:
+        """Refresh the set of classes in the window. Used by classifiers where labels are added as [1] in the vertex tuple.
+
+        This is used to clean up classes that are no longer in the window, and
+        ensure we do not consider "None" a class. It is called every `cleanup_every`
+        step, or can be called manually.
+
+        """
+        return {  # self.window is a deque of items, where each item is a tuple (x, y)
+            pos.item[1] for pos in self.window if pos is not None and pos.item[1] is not None
+        }