pytrilogy 0.3.149__cp313-cp313-win_amd64.whl

trilogy/core/graph_models.py

@@ -0,0 +1,142 @@
+from logging import Logger
+from typing import Union
+
+import networkx as nx
+
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildDatasource,
+    BuildUnionDatasource,
+    BuildWhereClause,
+)
+
+
+def get_graph_exact_match(
+    g: Union[nx.DiGraph, "ReferenceGraph"],
+    accept_partial: bool,
+    conditions: BuildWhereClause | None,
+) -> set[str]:
+    exact: set[str] = set()
+    for node, ds in g.datasources.items():
+        if isinstance(ds, BuildUnionDatasource):
+            exact.add(node)
+            continue
+
+        if not conditions and not ds.non_partial_for:
+            exact.add(node)
+            continue
+        elif not conditions and accept_partial and ds.non_partial_for:
+            exact.add(node)
+            continue
+        elif conditions:
+            if not ds.non_partial_for:
+                continue
+            if ds.non_partial_for and conditions == ds.non_partial_for:
+                exact.add(node)
+                continue
+        else:
+            continue
+
+    return exact
+
+
+def prune_sources_for_conditions(
+    g: "ReferenceGraph",
+    accept_partial: bool,
+    conditions: BuildWhereClause | None,
+):
+    complete = get_graph_exact_match(g, accept_partial, conditions)
+    to_remove = []
+    for node in g.datasources:
+        if node not in complete:
+            to_remove.append(node)
+
+    for node in to_remove:
+        g.remove_node(node)
+
+
+def prune_sources_for_aggregates(
+    g: "ReferenceGraph",
+    all_concepts: list[BuildConcept],
+    logger: Logger,
+) -> None:
+    required_grains = []
+    for x in all_concepts:
+        logger.debug(
+            f"Checking concept {x.address} at grain {x.grain}, is_aggregate={x.is_aggregate}"
+        )
+        if x.is_aggregate:
+            logger.debug(f"Aggregate found: {x.address} at grain {x.grain}")
+            required_grains.append(x.grain)
+    # if no aggregates, exit
+    logger.debug(f"Required grains for aggregates: {required_grains}")
+    if not required_grains:
+
+        return
+    # if we have distinct grains required, exit
+    if len(required_grains) > 1:
+        logger.debug("Multiple required grains found, cannot prune datasources.")
+        return
+    to_remove = []
+    for node, ds in g.datasources.items():
+        if ds.grain != required_grains[0]:
+            logger.debug(f"Removing datasource {node} at grain {ds.grain}")
+            to_remove.append(node)
+    for node in to_remove:
+        g.remove_node(node)
+    return
+
+
+def concept_to_node(input: BuildConcept, stash: dict[str, str] | None = None) -> str:
+    lookup = input.canonical_address_grain
+    if stash and lookup in stash:
+        return stash[lookup]
+    # if input.purpose == Purpose.METRIC:
+    #     return f"c~{input.namespace}.{input.name}@{input.grain}"
+
+    r = f"c~{input.canonical_address}@{input.grain.str_no_condition}"
+    if stash is not None:
+        stash[lookup] = r
+    return r
+
+
+def datasource_to_node(input: BuildDatasource) -> str:
+    # if isinstance(input, JoinedDataSource):
+    #     return "ds~join~" + ",".join(
+    #         [datasource_to_node(sub) for sub in input.datasources]
+    #     )
+    return f"ds~{input.identifier}"
+
+
+class ReferenceGraph(nx.DiGraph):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.concepts: dict[str, BuildConcept] = {}
+        self.datasources: dict[str, BuildDatasource | BuildUnionDatasource] = {}
+        self.pseudonyms: set[tuple[str, str]] = set()
+
+    def copy(self) -> "ReferenceGraph":
+        g = ReferenceGraph()
+        g.concepts = self.concepts.copy()
+        g.datasources = self.datasources.copy()
+        g.pseudonyms = self.pseudonyms.copy()
+        g._node.update(self._node)
+        g._adj.update({k: dict(v) for k, v in self._adj.items()})
+        g._pred.update({k: dict(v) for k, v in self._pred.items()})
+        return g
+
+    def remove_node(self, n) -> None:
+        if n in self.concepts:
+            del self.concepts[n]
+        if n in self.datasources:
+            del self.datasources[n]
+        super().remove_node(n)
+
+    def add_node(self, node_for_adding, **attr):
+        return super().add_node(node_for_adding, **attr)
+
+    def add_datasource_node(self, node_name, datasource) -> None:
+        super().add_node(node_name, datasource=datasource)
+
+    def add_edge(self, u_of_edge, v_of_edge, **attr):
+        return super().add_edge(u_of_edge, v_of_edge, **attr)

trilogy/core/internal.py

@@ -0,0 +1,85 @@
+from trilogy.core.constants import ALL_ROWS_CONCEPT, INTERNAL_NAMESPACE
+from trilogy.core.enums import Derivation, FunctionType, Granularity, Purpose
+from trilogy.core.models.author import Concept, Function, Grain
+from trilogy.core.models.core import DataType
+
+DEFAULT_CONCEPTS = {
+    ALL_ROWS_CONCEPT: Concept(
+        name=ALL_ROWS_CONCEPT,
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.INTEGER,
+        purpose=Purpose.CONSTANT,
+        grain=Grain(),
+        lineage=Function(
+            operator=FunctionType.CONSTANT,
+            arguments=[1],
+            output_datatype=DataType.INTEGER,
+            output_purpose=Purpose.CONSTANT,
+        ),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "concept_address": Concept(
+        name="concept_address",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "concept_datatype": Concept(
+        name="concept_datatype",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "concept_description": Concept(
+        name="concept_description",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "datasource": Concept(
+        name="datasource",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "query_text": Concept(
+        name="query_text",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "label": Concept(
+        name="label",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "expected": Concept(
+        name="expected_value",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+}