tree parser done

Author: amol-

src/mustela/translation/steps/trees/tree.py

@@ -1,4 +1,6 @@
+"""Prase tree definitions and return a graph of nodes."""
 import itertools
+import typing
 
 import ibis
 
@@ -10,16 +12,16 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
     
     The tree is built based on the node and attributes of the translator.
     """
-    nodes_treeids: list[int] = translator._attributes["nodes_treeids"]
-    nodes_nodeids: list[int] = translator._attributes["nodes_nodeids"]
-    nodes_modes: list[str] = translator._attributes["nodes_modes"]
-    nodes_truenodeids: list[int] = translator._attributes["nodes_truenodeids"]
-    nodes_falsenodeids: list[int] = translator._attributes["nodes_falsenodeids"]
-    nodes_thresholds: list[float] = translator._attributes["nodes_values"]
-    nodes_featureids: list[int] = translator._attributes["nodes_featureids"]
-    nodes_missing_value_tracks_true: list[int] = translator._attributes[
+    nodes_treeids = typing.cast(list[int], translator._attributes["nodes_treeids"])
+    nodes_nodeids = typing.cast(list[int], translator._attributes["nodes_nodeids"])
+    nodes_modes = typing.cast(list[str], translator._attributes["nodes_modes"])
+    nodes_truenodeids = typing.cast(list[int], translator._attributes["nodes_truenodeids"])
+    nodes_falsenodeids = typing.cast(list[int], translator._attributes["nodes_falsenodeids"])
+    nodes_thresholds = typing.cast(list[float], translator._attributes["nodes_values"])
+    nodes_featureids = typing.cast(list[int], translator._attributes["nodes_featureids"])
+    nodes_missing_value_tracks_true = typing.cast(list[int], translator._attributes[
         "nodes_missing_value_tracks_true"
-    ]
+    ])
     node = translator._node
 
     # Assert a few things to ensure we don't ed up genearting a tree with wrong data
@@ -37,19 +39,19 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
     weights = {}
     if node.op_type == "TreeEnsembleClassifier":
         # Weights for classifier, in this case the weights are per-class
-        class_nodeids: list[int] = translator._attributes["class_nodeids"]
-        class_treeids: list[int] = translator._attributes["class_treeids"]
-        class_weights: list[float] = translator._attributes["class_weights"]
-        weights_classid: list[int] = translator._attributes["class_ids"]
+        class_nodeids = typing.cast(list[int], translator._attributes["class_nodeids"])
+        class_treeids = typing.cast(list[int], translator._attributes["class_treeids"])
+        class_weights = typing.cast(list[float], translator._attributes["class_weights"])
+        weights_classid = typing.cast(list[int], translator._attributes["class_ids"])
         assert (
             len(class_treeids)
             == len(class_nodeids)
             == len(class_weights)
             == len(weights_classid)
         )
-        classlabels: None|list[str|int] = translator._attributes.get(
+        classlabels = typing.cast(None|list[str|int], translator._attributes.get(
             "classlabels_strings"
-        ) or translator._attributes.get("classlabels_int64s")
+        ) or translator._attributes.get("classlabels_int64s"))
         if not classlabels:
             raise ValueError("Missing class labels when building tree")
 
@@ -61,9 +63,9 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
             )
     elif node.op_type == "TreeEnsembleRegressor":
         # Weights for the regressor, in this case leaf nodes have only 1 weight
-        target_weights: list[float] = translator._attributes["target_weights"]
-        target_nodeids: list[int] = translator._attributes["target_nodeids"]
-        target_treeids: list[int] = translator._attributes["target_treeids"]
+        target_weights = typing.cast(list[float], translator._attributes["target_weights"])
+        target_nodeids = typing.cast(list[int], translator._attributes["target_nodeids"])
+        target_treeids = typing.cast(list[int], translator._attributes["target_treeids"])
         assert len(target_treeids) == len(target_nodeids) == len(target_weights)
         for tree_id, node_id, weight in zip(
             target_treeids, target_nodeids, target_weights
@@ -125,7 +127,13 @@ def build_tree(translator: Translator) -> dict[int, dict[int, dict]]:
     return {tree_id: trees[tree_id][0] for tree_id in trees}
 
 
-def mode_to_condition(node, feature_expr: ibis.Expr) -> ibis.Expr:
+def mode_to_condition(node: dict, feature_expr: ibis.Expr) -> ibis.Expr:
+    """Build a comparison expression for a branch node.
+
+    The comparison is based on the mode of the node and the threshold
+    for that noode. The feature will be compared to the threshold
+    using the operator defined by the mode.
+    """
     threshold = node["treshold"]
     if node["mode"] == "BRANCH_LEQ":
         condition = feature_expr <= threshold