pytrilogy 0.3.148__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

trilogy/core/processing/discovery_node_factory.py

@@ -0,0 +1,451 @@
+from dataclasses import dataclass
+from typing import List, Optional, Protocol, Union
+
+from trilogy.constants import logger
+from trilogy.core.enums import Derivation, Granularity
+from trilogy.core.graph_models import ReferenceGraph
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildWhereClause,
+)
+from trilogy.core.models.build_environment import BuildEnvironment
+from trilogy.core.processing.discovery_utility import LOGGER_PREFIX, depth_to_prefix
+from trilogy.core.processing.node_generators import (
+    gen_basic_node,
+    gen_constant_node,
+    gen_filter_node,
+    gen_group_node,
+    gen_group_to_node,
+    gen_merge_node,
+    gen_multiselect_node,
+    gen_recursive_node,
+    gen_rowset_node,
+    gen_synonym_node,
+    gen_union_node,
+    gen_unnest_node,
+    gen_window_node,
+)
+from trilogy.core.processing.nodes import (
+    History,
+    StrategyNode,
+)
+
+
+class SearchConceptsType(Protocol):
+    def __call__(
+        self,
+        mandatory_list: List[BuildConcept],
+        history: History,
+        environment: BuildEnvironment,
+        depth: int,
+        g: ReferenceGraph,
+        accept_partial: bool = False,
+        conditions: Optional[BuildWhereClause] = None,
+    ) -> Union[StrategyNode, None]: ...
+
+
+@dataclass
+class NodeGenerationContext:
+    """Encapsulates common parameters for node generation."""
+
+    concept: BuildConcept
+    local_optional: List[BuildConcept]
+    environment: BuildEnvironment
+    g: ReferenceGraph
+    depth: int
+    source_concepts: SearchConceptsType
+    history: History
+    accept_partial: bool = False
+    conditions: Optional[BuildWhereClause] = None
+
+    @property
+    def next_depth(self) -> int:
+        return self.depth + 1
+
+    def log_generation(self, node_type: str, extra_info: str = "") -> None:
+        """Centralized logging for node generation."""
+        optional_addresses = [x.address for x in self.local_optional]
+        base_msg = f"for {self.concept.address}, generating {node_type} node with optional {optional_addresses}"
+
+        if extra_info:
+            base_msg += f" and {extra_info}"
+
+        logger.info(f"{depth_to_prefix(self.depth)}{LOGGER_PREFIX} {base_msg}")
+
+
+def restrict_node_outputs_targets(
+    node: StrategyNode, targets: list[BuildConcept], depth: int
+) -> list[BuildConcept]:
+    """Restricts node outputs to target concepts and returns extra concepts."""
+    ex_resolve = node.resolve()
+    target_addresses = {y.address for y in targets}
+
+    extra = [
+        x
+        for x in ex_resolve.output_concepts
+        if x.address not in target_addresses
+        and not any(c in targets for c in x.pseudonyms)
+    ]
+
+    base = [
+        x
+        for x in ex_resolve.output_concepts
+        if x.address not in {c.address for c in extra}
+    ]
+
+    logger.info(
+        f"{depth_to_prefix(depth)}{LOGGER_PREFIX} reducing final outputs, "
+        f"was {[c.address for c in ex_resolve.output_concepts]} "
+        f"with extra {[c.address for c in extra]}, remaining {base} (targets {target_addresses})"
+    )
+
+    # Add missing targets
+    for target in targets:
+        if target.address not in {c.address for c in base}:
+            base.append(target)
+
+    node.set_output_concepts(base)
+    return extra
+
+
+# Simple factory functions for basic derivation types
+def _generate_window_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("window")
+    return gen_window_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_filter_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("filter")
+    return gen_filter_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_unnest_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("unnest", f"condition {ctx.conditions}")
+    return gen_unnest_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_recursive_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("recursive", f"condition {ctx.conditions}")
+    return gen_recursive_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_union_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("union", f"condition {ctx.conditions}")
+    return gen_union_node(
+        ctx.concept,
+        ctx.local_optional,
+        ctx.environment,
+        ctx.g,
+        ctx.next_depth,
+        ctx.source_concepts,
+        ctx.history,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_aggregate_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    # Filter out constants to avoid multiplication issues
+    agg_optional = [
+        x
+        for x in ctx.local_optional
+        if not (
+            x.granularity == Granularity.SINGLE_ROW
+            and x.derivation != Derivation.AGGREGATE
+        )
+    ]
+
+    logger.info(
+        f"{depth_to_prefix(ctx.depth)}{LOGGER_PREFIX} "
+        f"for {ctx.concept.address}, generating aggregate node with optional {agg_optional}"
+    )
+
+    return gen_group_node(
+        ctx.concept,
+        agg_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_rowset_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("rowset")
+    return gen_rowset_node(
+        ctx.concept,
+        ctx.local_optional,
+        ctx.environment,
+        ctx.g,
+        ctx.next_depth,
+        ctx.source_concepts,
+        ctx.history,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_multiselect_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("multiselect")
+    return gen_multiselect_node(
+        ctx.concept,
+        ctx.local_optional,
+        ctx.environment,
+        ctx.g,
+        ctx.next_depth,
+        ctx.source_concepts,
+        ctx.history,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_group_to_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("group to grain")
+    return gen_group_to_node(
+        ctx.concept,
+        ctx.local_optional,
+        ctx.environment,
+        ctx.g,
+        ctx.next_depth,
+        ctx.source_concepts,
+        ctx.history,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_basic_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("basic")
+    return gen_basic_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+    )
+
+
+def _generate_constant_node(ctx: NodeGenerationContext) -> StrategyNode | None:
+    ctx.log_generation("constant")
+    return gen_constant_node(
+        ctx.concept,
+        ctx.local_optional,
+        history=ctx.history,
+        environment=ctx.environment,
+        g=ctx.g,
+        depth=ctx.next_depth,
+        source_concepts=ctx.source_concepts,
+        conditions=ctx.conditions,
+        accept_partial=ctx.accept_partial,
+    )
+
+
+class RootNodeHandler:
+    """Handles complex root node generation logic."""
+
+    def __init__(self, context: NodeGenerationContext):
+        self.ctx = context
+
+    def generate(self) -> Optional[StrategyNode]:
+        self.ctx.log_generation("select", "including condition inputs")
+
+        root_targets = [self.ctx.concept] + self.ctx.local_optional
+
+        return self._resolve_root_concepts(root_targets)
+
+    def _resolve_root_concepts(
+        self, root_targets: List[BuildConcept]
+    ) -> Optional[StrategyNode]:
+        expanded_node = self._try_merge_expansion(root_targets)
+        if expanded_node:
+            return expanded_node
+        if self.ctx.accept_partial:
+            synonym_node = self._try_synonym_resolution(root_targets)
+            if synonym_node:
+                logger.info(
+                    f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+                    f"resolved root concepts through synonyms"
+                )
+                return synonym_node
+
+        return None
+
+    def _try_merge_expansion(
+        self, root_targets: List[BuildConcept]
+    ) -> Optional[StrategyNode]:
+        for accept_partial in [False, True]:
+            expanded = gen_merge_node(
+                all_concepts=root_targets,
+                environment=self.ctx.environment,
+                g=self.ctx.g,
+                depth=self.ctx.next_depth,
+                source_concepts=self.ctx.source_concepts,
+                history=self.ctx.history,
+                search_conditions=self.ctx.conditions,
+                accept_partial=accept_partial,
+            )
+
+            if expanded:
+                self._handle_expanded_node(expanded, root_targets)
+                return expanded
+
+        logger.info(
+            f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+            f"could not find additional concept(s) to inject"
+        )
+        return None
+
+    def _handle_expanded_node(
+        self, expanded: StrategyNode, root_targets: List[BuildConcept]
+    ) -> None:
+        extra = restrict_node_outputs_targets(expanded, root_targets, self.ctx.depth)
+
+        logger.info(
+            f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+            f"Found connections for {[c.address for c in root_targets]} "
+            f"via concept addition; removing extra {[c.address for c in extra]}"
+        )
+
+    def _try_synonym_resolution(
+        self, root_targets: List[BuildConcept]
+    ) -> Optional[StrategyNode]:
+        logger.info(
+            f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+            f"Could not resolve root concepts, checking for synonyms for {root_targets}"
+        )
+
+        if not self.ctx.history.check_started(
+            root_targets,
+            accept_partial=self.ctx.accept_partial,
+            conditions=self.ctx.conditions,
+        ):
+            self.ctx.history.log_start(
+                root_targets,
+                accept_partial=self.ctx.accept_partial,
+                conditions=self.ctx.conditions,
+            )
+
+            resolved = gen_synonym_node(
+                all_concepts=root_targets,
+                environment=self.ctx.environment,
+                g=self.ctx.g,
+                depth=self.ctx.next_depth,
+                source_concepts=self.ctx.source_concepts,
+                history=self.ctx.history,
+                conditions=self.ctx.conditions,
+                accept_partial=self.ctx.accept_partial,
+            )
+
+            if resolved:
+                logger.info(
+                    f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+                    f"resolved concepts through synonyms"
+                )
+                return resolved
+        else:
+            logger.info(
+                f"{depth_to_prefix(self.ctx.depth)}{LOGGER_PREFIX} "
+                f"skipping synonym search, already in a recursion for these concepts"
+            )
+
+        return None
+
+
+def generate_node(
+    concept: BuildConcept,
+    local_optional: List[BuildConcept],
+    environment: BuildEnvironment,
+    g: ReferenceGraph,
+    depth: int,
+    source_concepts: SearchConceptsType,
+    history: History,
+    accept_partial: bool,
+    conditions: BuildWhereClause | None = None,
+) -> StrategyNode | None:
+
+    context = NodeGenerationContext(
+        concept=concept,
+        local_optional=local_optional,
+        environment=environment,
+        g=g,
+        depth=depth,
+        source_concepts=source_concepts,
+        history=history,
+        accept_partial=accept_partial,
+        conditions=conditions,
+    )
+
+    # Try materialized concept first
+    # this is worth checking every loop iteration
+    candidate = history.gen_select_node(
+        [concept] + local_optional,
+        environment,
+        g,
+        depth + 1,
+        fail_if_not_found=False,
+        accept_partial=accept_partial,
+        conditions=conditions,
+    )
+
+    if candidate:
+        return candidate
+
+    # Delegate to appropriate handler based on derivation
+    derivation_handlers = {
+        Derivation.WINDOW: lambda: _generate_window_node(context),
+        Derivation.FILTER: lambda: _generate_filter_node(context),
+        Derivation.UNNEST: lambda: _generate_unnest_node(context),
+        Derivation.RECURSIVE: lambda: _generate_recursive_node(context),
+        Derivation.UNION: lambda: _generate_union_node(context),
+        Derivation.AGGREGATE: lambda: _generate_aggregate_node(context),
+        Derivation.ROWSET: lambda: _generate_rowset_node(context),
+        Derivation.MULTISELECT: lambda: _generate_multiselect_node(context),
+        Derivation.GROUP_TO: lambda: _generate_group_to_node(context),
+        Derivation.BASIC: lambda: _generate_basic_node(context),
+        Derivation.ROOT: lambda: RootNodeHandler(context).generate(),
+        Derivation.CONSTANT: lambda: _generate_constant_node(context),
+    }
+
+    handler = derivation_handlers.get(concept.derivation)
+    if not handler:
+        raise ValueError(f"Unknown derivation {concept.derivation} on {concept}")
+
+    return handler()