Query key pipeline simplification

data_prototype/conversion_node.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterable, Callable, Sequence
 from collections import Counter
 from dataclasses import dataclass
+from enum import Enum, auto
 import inspect
 from functools import cached_property
 
@@ -15,39 +16,84 @@ def evaluate_pipeline(nodes: Sequence[ConversionNode], input: dict[str, Any]):
     return input
 
 
+def preview_pipeline(nodes: Sequence[ConversionNode], input: Iterable[str]):
+    out = [tuple(input)]
+    for node in nodes:
+        out.append(node.preview_keys(out[-1]))
+    return out
+
+
+def simplify_pipeline(nodes: Sequence[ConversionNode], output_keys: Iterable[str]):
+    out_nodes = []
+    required_keys = set(output_keys)
+    for node in reversed(nodes):
+        print(node, required_keys)
+        next_required = set()
+        next_removed = set()
+        for key in required_keys:
+            effect, added_reqs = node.affects_key(key)
+            print(key, effect, added_reqs)
+            if effect == AffectsKey.CREATE:
+                next_removed.update(key)
+                next_required |= set(added_reqs)
+                out_nodes.insert(0, node)
+            elif effect == AffectsKey.MODIFY:
+                next_required |= set(added_reqs)
+                out_nodes.insert(0, node)
+            elif effect == AffectsKey.REMOVE:
+                raise ValueError(f"Required key {key!r} removed by node")
+        required_keys |= next_required
+        required_keys -= next_removed
+    return out_nodes
+
+
+class AffectsKey(Enum):
+    CREATE = auto()
+    MODIFY = auto()
+    REMOVE = auto()
+    IGNORE = auto()
+    USE = auto()
+
+
 @dataclass
 class ConversionNode:
     required_keys: tuple[str, ...]
     output_keys: tuple[str, ...]
-    trim_keys: bool
 
     def preview_keys(self, input_keys: Iterable[str]) -> tuple[str, ...]:
         if missing_keys := set(self.required_keys) - set(input_keys):
             raise ValueError(f"Missing keys: {missing_keys}")
-        if self.trim_keys:
-            return tuple(sorted(set(self.output_keys)))
         return tuple(sorted(set(input_keys) | set(self.output_keys)))
 
-    def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
-        if self.trim_keys:
-            return {k: input[k] for k in self.output_keys}
+    def affects_key(self, key: str) -> tuple[AffectsKey, tuple[str, ...]]:
+        inp = key in self.required_keys
+        outp = key in self.output_keys
+        if inp and outp:
+            return AffectsKey.MODIFY, self.required_keys
+        elif inp:
+            return AffectsKey.USE, ()
+        elif outp:
+            return AffectsKey.CREATE, self.required_keys
         else:
-            if missing_keys := set(self.output_keys) - set(input):
-                raise ValueError(f"Missing keys: {missing_keys}")
-            return input
+            return AffectsKey.IGNORE, ()
+
+    def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
+        if missing_keys := set(self.output_keys) - set(input):
+            raise ValueError(f"Missing keys: {missing_keys}")
+        return input
 
 
 @dataclass
 class UnionConversionNode(ConversionNode):
     nodes: tuple[ConversionNode, ...]
 
     @classmethod
-    def from_nodes(cls, *nodes: ConversionNode, trim_keys=False):
+    def from_nodes(cls, *nodes: ConversionNode):
         required = tuple(set(k for n in nodes for k in n.required_keys))
         output = Counter(k for n in nodes for k in n.output_keys)
         if duplicate := {k for k, v in output.items() if v > 1}:
             raise ValueError(f"Duplicate keys from multiple input nodes: {duplicate}")
-        return cls(required, tuple(output), trim_keys, nodes)
+        return cls(required, tuple(output), nodes)
 
     def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
         return super().evaluate(
@@ -59,13 +105,25 @@ def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
 class RenameConversionNode(ConversionNode):
     mapping: dict[str, str]
 
+    def affects_key(self, key: str) -> tuple[AffectsKey, tuple[str, ...]]:
+        inp = key in self.mapping
+        outp = key in self.mapping.values()
+        if inp and outp:
+            return AffectsKey.MODIFY, self.required_keys
+        elif inp:
+            return AffectsKey.REMOVE, self.required_keys
+        elif outp:
+            return AffectsKey.CREATE, self.required_keys
+        else:
+            return AffectsKey.IGNORE, ()
+
     @classmethod
-    def from_mapping(cls, mapping: dict[str, str], trim_keys=False):
+    def from_mapping(cls, mapping: dict[str, str]):
         required = tuple(mapping)
         output = Counter(mapping.values())
         if duplicate := {k for k, v in output.items() if v > 1}:
             raise ValueError(f"Duplicate output keys in mapping: {duplicate}")
-        return cls(required, tuple(output), trim_keys, mapping)
+        return cls(required, tuple(output), mapping)
 
     def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
         return super().evaluate(
@@ -82,14 +140,14 @@ def _sigs(self):
         return {k: (f, inspect.signature(f)) for k, f in self.funcs.items()}
 
     @classmethod
-    def from_funcs(cls, funcs: dict[str, Callable], trim_keys=False):
+    def from_funcs(cls, funcs: dict[str, Callable]):
         sigs = {k: inspect.signature(f) for k, f in funcs.items()}
         output = tuple(sigs)
-        input = []
+        input = set()
         for v in sigs.values():
-            input.extend(v.parameters.keys())
-        input = tuple(set(input))
-        return cls(input, output, trim_keys, funcs)
+            input |= set(v.parameters.keys())
+        input = tuple(input)
+        return cls(input, output, funcs)
 
     def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
         return super().evaluate(
@@ -107,9 +165,14 @@ def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
 class LimitKeysConversionNode(ConversionNode):
     keys: set[str]
 
+    def affects_key(self, key: str) -> tuple[AffectsKey, tuple[str, ...]]:
+        if key in self.keys:
+            return AffectsKey.IGNORE, ()
+        return AffectsKey.REMOVE, ()
+
     @classmethod
     def from_keys(cls, keys: Sequence[str]):
-        return cls((), tuple(keys), trim_keys=True, keys=set(keys))
+        return cls((), tuple(keys), keys=set(keys))
 
     def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
         return {k: v for k, v in input.items() if k in self.keys}