Fix surprisal tree kwarg mismatch (#749)

* bug fix

* test: add tests for create_tree k-kwarg handling

src/dreamer/trees/__init__.py

@@ -29,6 +29,15 @@ def create_tree(tree_type: str, **kwargs: Any) -> SurprisalTree:
     Raises:
         ValueError: If tree_type is not recognized
     """
+    # `surprisal.py` calls this factory with a uniform `k=settings.DREAM.SURPRISAL.TREE_K` kwarg for every tree type,
+    # but `k` is only meaningful for the kNN-based trees (kdtree, balltree, graph).
+    # The other 4 use different tunables and raise TypeError if `k` is passed.
+    # Drop it here so the factory accepts a uniform kwargs dict.
+
+    trees_without_k = {"rptree", "covertree", "lsh", "prototype"}
+    if tree_type in trees_without_k:
+        kwargs.pop("k", None)
+
     if tree_type == "rptree":
         return RPTree(**kwargs)
     elif tree_type == "kdtree":

tests/dreamer/test_trees.py

@@ -0,0 +1,51 @@
+import pytest
+
+from src.dreamer.trees import (
+    CoverTree,
+    GraphSurprisal,
+    LSHSurprisal,
+    PrototypeSurprisal,
+    RPTree,
+    SklearnTreeWrapper,
+    SurprisalTree,
+    create_tree,
+)
+
+ALL_TREE_TYPES = [
+    "kdtree",
+    "balltree",
+    "rptree",
+    "covertree",
+    "lsh",
+    "graph",
+    "prototype",
+]
+
+EXPECTED_CLASS = {
+    "kdtree": SklearnTreeWrapper,
+    "balltree": SklearnTreeWrapper,
+    "rptree": RPTree,
+    "covertree": CoverTree,
+    "lsh": LSHSurprisal,
+    "graph": GraphSurprisal,
+    "prototype": PrototypeSurprisal,
+}
+
+
+@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES)
+def test_create_tree_accepts_uniform_k_kwarg(tree_type: str):
+    tree = create_tree(tree_type=tree_type, k=5)
+    assert isinstance(tree, SurprisalTree)
+    assert isinstance(tree, EXPECTED_CLASS[tree_type])
+
+
+@pytest.mark.parametrize("tree_type", ALL_TREE_TYPES)
+def test_create_tree_without_k(tree_type: str):
+    """The factory should also work when no ``k`` is supplied."""
+    tree = create_tree(tree_type=tree_type)
+    assert isinstance(tree, EXPECTED_CLASS[tree_type])
+
+
+def test_create_tree_unknown_type_raises():
+    with pytest.raises(ValueError, match="Unknown tree type"):
+        create_tree(tree_type="not_a_tree")