pytrilogy 0.0.2.58__py3-none-any.whl → 0.0.3.0__py3-none-any.whl

trilogy/core/processing/node_generators/select_helpers/datasource_injection.py

@@ -2,24 +2,23 @@ from collections import defaultdict
 from datetime import date, datetime, timedelta
 from typing import List, Tuple, TypeVar
 
-from trilogy.core.enums import ComparisonOperator
-from trilogy.core.models import (
-    Comparison,
-    Concept,
-    Conditional,
-    Datasource,
-    DataType,
-    Function,
-    FunctionType,
-    Parenthetical,
+from trilogy.core.enums import ComparisonOperator, FunctionType
+from trilogy.core.models.build import (
+    BuildComparison,
+    BuildConcept,
+    BuildConditional,
+    BuildDatasource,
+    BuildFunction,
+    BuildParenthetical,
 )
+from trilogy.core.models.core import DataType
 
 # Define a generic type that ensures start and end are the same type
 T = TypeVar("T", int, date, datetime)
 
 
 def reduce_expression(
-    var: Concept, group_tuple: list[tuple[ComparisonOperator, T]]
+    var: BuildConcept, group_tuple: list[tuple[ComparisonOperator, T]]
 ) -> bool:
     # Track ranges
     lower_check: T
@@ -95,27 +94,27 @@ def reduce_expression(
 
 
 def simplify_conditions(
-    conditions: list[Comparison | Conditional | Parenthetical],
+    conditions: list[BuildComparison | BuildConditional | BuildParenthetical],
 ) -> bool:
     # Group conditions by variable
-    grouped: dict[Concept, list[tuple[ComparisonOperator, datetime | int | date]]] = (
-        defaultdict(list)
-    )
+    grouped: dict[
+        BuildConcept, list[tuple[ComparisonOperator, datetime | int | date]]
+    ] = defaultdict(list)
     for condition in conditions:
-        if not isinstance(condition, Comparison):
+        if not isinstance(condition, BuildComparison):
             return False
         if not isinstance(
-            condition.left, (int, date, datetime, Function)
-        ) and not isinstance(condition.right, (int, date, datetime, Function)):
+            condition.left, (int, date, datetime, BuildFunction)
+        ) and not isinstance(condition.right, (int, date, datetime, BuildFunction)):
             return False
-        if not isinstance(condition.left, Concept) and not isinstance(
-            condition.right, Concept
+        if not isinstance(condition.left, BuildConcept) and not isinstance(
+            condition.right, BuildConcept
         ):
             return False
         vars = [condition.left, condition.right]
-        concept = [x for x in vars if isinstance(x, Concept)][0]
-        comparison = [x for x in vars if not isinstance(x, Concept)][0]
-        if isinstance(comparison, Function):
+        concept = [x for x in vars if isinstance(x, BuildConcept)][0]
+        comparison = [x for x in vars if not isinstance(x, BuildConcept)][0]
+        if isinstance(comparison, BuildFunction):
             if not comparison.operator == FunctionType.CONSTANT:
                 return False
             first_arg = comparison.arguments[0]
@@ -159,26 +158,19 @@ def is_fully_covered(
 
     # Check for gaps
     current_end = start
-    print(ranges)
     for r_start, r_end in ranges:
-        print(r_start, r_end)
         # If there's a gap between the current range and the previous coverage
-        print(r_start - current_end)
         if (r_start - current_end) > increment:  # type: ignore
-            print("gap")
             return False
-        print("okay")
         # Extend the current coverage
         current_end = max(current_end, r_end)
 
     # If the loop ends and we haven't reached the end, return False
-    print(current_end, end)
-    print(current_end >= end)
     return current_end >= end
 
 
-def get_union_sources(datasources: list[Datasource], concepts: list[Concept]):
-    candidates: list[Datasource] = []
+def get_union_sources(datasources: list[BuildDatasource], concepts: list[BuildConcept]):
+    candidates: list[BuildDatasource] = []
     for x in datasources:
         if all([c.address in x.output_concepts for c in concepts]):
             if (
@@ -187,7 +179,7 @@ def get_union_sources(datasources: list[Datasource], concepts: list[Concept]):
             ):
                 candidates.append(x)
 
-    assocs: dict[str, list[Datasource]] = defaultdict(list[Datasource])
+    assocs: dict[str, list[BuildDatasource]] = defaultdict(list[BuildDatasource])
     for x in candidates:
         if not x.non_partial_for:
             continue
@@ -195,8 +187,9 @@ def get_union_sources(datasources: list[Datasource], concepts: list[Concept]):
             continue
         merge_key = x.non_partial_for.concept_arguments[0]
         assocs[merge_key.address].append(x)
-    final: list[list[Datasource]] = []
+    final: list[list[BuildDatasource]] = []
     for _, dses in assocs.items():
+
         conditions = [c.non_partial_for.conditional for c in dses if c.non_partial_for]
         if simplify_conditions(conditions):
             final.append(dses)

trilogy/core/processing/node_generators/select_merge_node.py

@@ -1,18 +1,19 @@
+from functools import reduce
 from typing import List, Optional
 
 import networkx as nx
 
 from trilogy.constants import logger
-from trilogy.core.enums import PurposeLineage
+from trilogy.core.enums import Derivation
 from trilogy.core.graph_models import concept_to_node
-from trilogy.core.models import (
-    Concept,
-    Datasource,
-    Environment,
-    Grain,
-    LooseConceptList,
-    WhereClause,
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildDatasource,
+    BuildGrain,
+    BuildWhereClause,
+    LooseBuildConceptList,
 )
+from trilogy.core.models.build_environment import BuildEnvironment
 from trilogy.core.processing.node_generators.select_helpers.datasource_injection import (
     get_union_sources,
 )
@@ -33,16 +34,17 @@ def extract_address(node: str):
 
 
 def get_graph_partial_nodes(
-    g: nx.DiGraph, conditions: WhereClause | None
+    g: nx.DiGraph, conditions: BuildWhereClause | None
 ) -> dict[str, list[str]]:
-    datasources: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
-        g, "datasource"
+    datasources: dict[str, BuildDatasource | list[BuildDatasource]] = (
+        nx.get_node_attributes(g, "datasource")
     )
     partial: dict[str, list[str]] = {}
     for node in g.nodes:
         if node in datasources:
             ds = datasources[node]
             if not isinstance(ds, list):
+
                 if ds.non_partial_for and conditions == ds.non_partial_for:
                     partial[node] = []
                     continue
@@ -55,36 +57,58 @@ def get_graph_partial_nodes(
     return partial
 
 
-def get_graph_grain_length(g: nx.DiGraph) -> dict[str, int]:
-    datasources: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
-        g, "datasource"
+def get_graph_exact_match(
+    g: nx.DiGraph, conditions: BuildWhereClause | None
+) -> set[str]:
+    datasources: dict[str, BuildDatasource | list[BuildDatasource]] = (
+        nx.get_node_attributes(g, "datasource")
     )
-    grain_length: dict[str, int] = {}
+    exact: set[str] = set()
     for node in g.nodes:
         if node in datasources:
+            ds = datasources[node]
+            if not isinstance(ds, list):
+                if ds.non_partial_for and conditions == ds.non_partial_for:
+                    exact.add(node)
+                    continue
+            else:
+                continue
+
+    return exact
+
+
+def get_graph_grains(g: nx.DiGraph) -> dict[str, list[str]]:
+    datasources: dict[str, BuildDatasource | list[BuildDatasource]] = (
+        nx.get_node_attributes(g, "datasource")
+    )
+    grain_length: dict[str, list[str]] = {}
+    for node in g.nodes:
+        if node in datasources:
+            base: set[str] = set()
             lookup = datasources[node]
             if not isinstance(lookup, list):
                 lookup = [lookup]
             assert isinstance(lookup, list)
-            grain_length[node] = sum(len(x.grain.components) for x in lookup)
+            grain_length[node] = reduce(
+                lambda x, y: x.union(y.grain.components), lookup, base  # type: ignore
+            )
     return grain_length
 
 
 def create_pruned_concept_graph(
     g: nx.DiGraph,
-    all_concepts: List[Concept],
-    datasources: list[Datasource],
+    all_concepts: List[BuildConcept],
+    datasources: list[BuildDatasource],
     accept_partial: bool = False,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
     depth: int = 0,
 ) -> nx.DiGraph:
     orig_g = g
     g = g.copy()
-
     union_options = get_union_sources(datasources, all_concepts)
     for ds_list in union_options:
         node_address = "ds~" + "-".join([x.name for x in ds_list])
-        common: set[Concept] = set.intersection(
+        common: set[BuildConcept] = set.intersection(
             *[set(x.output_concepts) for x in ds_list]
         )
         g.add_node(node_address, datasource=ds_list)
@@ -92,9 +116,9 @@ def create_pruned_concept_graph(
             g.add_edge(node_address, concept_to_node(c))
 
     target_addresses = set([c.address for c in all_concepts])
-    concepts: dict[str, Concept] = nx.get_node_attributes(orig_g, "concept")
-    datasource_map: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
-        orig_g, "datasource"
+    concepts: dict[str, BuildConcept] = nx.get_node_attributes(orig_g, "concept")
+    datasource_map: dict[str, BuildDatasource | list[BuildDatasource]] = (
+        nx.get_node_attributes(orig_g, "datasource")
     )
     relevant_concepts_pre = {
         n: x.address
@@ -149,6 +173,7 @@ def create_pruned_concept_graph(
             if len(neighbors) > 1:
                 relevant_concepts.append(n)
             roots[root] = set()
+
     g.remove_nodes_from(
         [
             n
@@ -163,6 +188,8 @@ def create_pruned_concept_graph(
             f"{padding(depth)}{LOGGER_PREFIX} cannot resolve root graph - no subgraphs after node prune"
         )
         return None
+    # from trilogy.hooks.graph_hook import GraphHook
+    # GraphHook().query_graph_built(g)
     if subgraphs and len(subgraphs) != 1:
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} cannot resolve root graph - subgraphs are split - have {len(subgraphs)} from {subgraphs}"
@@ -183,15 +210,26 @@ def create_pruned_concept_graph(
 
 
 def resolve_subgraphs(
-    g: nx.DiGraph, conditions: WhereClause | None
+    g: nx.DiGraph, relevant: list[BuildConcept], conditions: BuildWhereClause | None
 ) -> dict[str, list[str]]:
+    """When we have multiple distinct subgraphs within our matched
+    nodes that can satisfy a query, resolve which one of those we should
+    ultimately ues.
+    This should generally return one subgraph for each
+    unique set of sub concepts that can be referenced,
+    discarding duplicates.
+    Duplicate subgraphs will be resolved based on which
+    ones are most 'optimal' to use, a hueristic
+    that can evolve in the future but is currently based on
+    cardinality."""
     datasources = [n for n in g.nodes if n.startswith("ds~")]
     subgraphs: dict[str, list[str]] = {
         ds: list(set(list(nx.all_neighbors(g, ds)))) for ds in datasources
     }
     partial_map = get_graph_partial_nodes(g, conditions)
-    grain_length = get_graph_grain_length(g)
-    concepts: dict[str, Concept] = nx.get_node_attributes(g, "concept")
+    exact_map = get_graph_exact_match(g, conditions)
+    grain_length = get_graph_grains(g)
+    concepts: dict[str, BuildConcept] = nx.get_node_attributes(g, "concept")
     non_partial_map = {
         ds: [concepts[c].address for c in subgraphs[ds] if c not in partial_map[ds]]
         for ds in datasources
@@ -200,7 +238,27 @@ def resolve_subgraphs(
         ds: [concepts[c].address for c in subgraphs[ds]] for ds in datasources
     }
     pruned_subgraphs = {}
+
+    def score_node(input: str):
+        logger.debug(f"scoring node {input}")
+        grain = grain_length[input]
+        # first - go for lowest grain
+        # but if the object we want is in the grain, treat that as "free"
+        # ex - pick source with grain(product_id) over grain(order_id)
+        # when going for product_id
+        score = (
+            len(list(grain)) - sum([1 for x in concept_map[input] if x in grain]),
+            # then check if it's an exact condition match
+            0 if input in exact_map else 0.5,
+            # last, number of concepts
+            len(subgraphs[input]),
+            input,
+        )
+        logger.debug(score)
+        return score
+
     for key, nodes in subgraphs.items():
+
         value = non_partial_map[key]
         all_concepts = concept_map[key]
         is_subset = False
@@ -219,31 +277,66 @@ def resolve_subgraphs(
                     logger.debug(
                         f"Dropping subgraph {key} with {value} as it is a subset of {other_key} with {other_value}"
                     )
-                    break
                 elif len(value) == len(other_value) and len(all_concepts) == len(
                     other_all_concepts
                 ):
                     matches.add(other_key)
                     matches.add(key)
-        if matches:
-            is_subset = key is not min(matches, key=lambda x: (grain_length[x], x))
+        if matches and not is_subset:
+            min_node = min(matches, key=score_node)
+            logger.debug(f"minimum source score is {min_node}")
+            is_subset = key is not min(matches, key=score_node)
         if not is_subset:
             pruned_subgraphs[key] = nodes
+
+    final_nodes: set[str] = set([n for v in pruned_subgraphs.values() for n in v])
+    relevant_concepts_pre = {
+        n: x.address
+        for n in g.nodes()
+        # filter out synonyms
+        if (x := concepts.get(n, None)) and x.address in relevant
+    }
+    for node in final_nodes:
+        keep = True
+        if node.startswith("c~") and node not in relevant_concepts_pre:
+            keep = (
+                sum(
+                    [
+                        1 if node in sub_nodes else 0
+                        for _, sub_nodes in pruned_subgraphs.items()
+                    ]
+                )
+                > 1
+            )
+        if not keep:
+            logger.debug(f"Pruning node {node} as irrelevant after subgraph resolution")
+            pruned_subgraphs = {
+                k: [n for n in v if n != node] for k, v in pruned_subgraphs.items()
+            }
+
     return pruned_subgraphs
 
 
 def create_datasource_node(
-    datasource: Datasource,
-    all_concepts: List[Concept],
+    datasource: BuildDatasource,
+    all_concepts: List[BuildConcept],
     accept_partial: bool,
-    environment: Environment,
+    environment: BuildEnvironment,
     depth: int,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
 ) -> tuple[StrategyNode, bool]:
-    target_grain = Grain.from_concepts(all_concepts, environment=environment)
+    logger.info(all_concepts)
+    target_grain = BuildGrain.from_concepts(all_concepts, environment=environment)
     force_group = False
     if not datasource.grain.issubset(target_grain):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX}_DS_NODE Select node must be wrapped in group, {datasource.grain} not subset of target grain {target_grain}"
+        )
         force_group = True
+    else:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX}_DS_NODE Select node grain {datasource.grain} is subset of target grain {target_grain}, no group required"
+        )
     if not datasource.grain.components:
         force_group = True
     partial_concepts = [
@@ -251,14 +344,17 @@ def create_datasource_node(
         for c in datasource.columns
         if not c.is_complete and c.concept.address in all_concepts
     ]
-    partial_lcl = LooseConceptList(concepts=partial_concepts)
+    partial_lcl = LooseBuildConceptList(concepts=partial_concepts)
     nullable_concepts = [
         c.concept
         for c in datasource.columns
         if c.is_nullable and c.concept.address in all_concepts
     ]
-    nullable_lcl = LooseConceptList(concepts=nullable_concepts)
+    nullable_lcl = LooseBuildConceptList(concepts=nullable_concepts)
     partial_is_full = conditions and (conditions == datasource.non_partial_for)
+
+    datasource_conditions = datasource.where.conditional if datasource.where else None
+
     return (
         SelectNode(
             input_concepts=[c.concept for c in datasource.columns],
@@ -272,8 +368,8 @@ def create_datasource_node(
             nullable_concepts=[c for c in all_concepts if c in nullable_lcl],
             accept_partial=accept_partial,
             datasource=datasource,
-            grain=Grain.from_concepts(all_concepts),
-            conditions=datasource.where.conditional if datasource.where else None,
+            grain=datasource.grain,
+            conditions=datasource_conditions,
             preexisting_conditions=(
                 conditions.conditional if partial_is_full and conditions else None
             ),
@@ -287,16 +383,16 @@ def create_select_node(
     subgraph: list[str],
     accept_partial: bool,
     g,
-    environment: Environment,
+    environment: BuildEnvironment,
     depth: int,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
 ) -> StrategyNode:
 
     all_concepts = [
         environment.concepts[extract_address(c)] for c in subgraph if c.startswith("c~")
     ]
 
-    if all([c.derivation == PurposeLineage.CONSTANT for c in all_concepts]):
+    if all([c.derivation == Derivation.CONSTANT for c in all_concepts]):
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} All concepts {[x.address for x in all_concepts]} are constants, returning constant node"
         )
@@ -311,10 +407,10 @@ def create_select_node(
             force_group=False,
         )
 
-    datasource: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
-        g, "datasource"
-    )[ds_name]
-    if isinstance(datasource, Datasource):
+    datasource: dict[str, BuildDatasource | list[BuildDatasource]] = (
+        nx.get_node_attributes(g, "datasource")[ds_name]
+    )
+    if isinstance(datasource, BuildDatasource):
         bcandidate, force_group = create_datasource_node(
             datasource,
             all_concepts,
@@ -325,6 +421,9 @@ def create_select_node(
         )
 
     elif isinstance(datasource, list):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} generating union node parents with condition {conditions}"
+        )
         from trilogy.core.processing.nodes.union_node import UnionNode
 
         force_group = False
@@ -340,6 +439,7 @@ def create_select_node(
             )
             parents.append(subnode)
             force_group = force_group or fg
+        logger.info(f"{padding(depth)}{LOGGER_PREFIX} generating union node")
         bcandidate = UnionNode(
             output_concepts=all_concepts,
             input_concepts=all_concepts,
@@ -373,15 +473,15 @@ def create_select_node(
 
 
 def gen_select_merge_node(
-    all_concepts: List[Concept],
+    all_concepts: List[BuildConcept],
     g: nx.DiGraph,
-    environment: Environment,
+    environment: BuildEnvironment,
     depth: int,
     accept_partial: bool = False,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
 ) -> Optional[StrategyNode]:
-    non_constant = [c for c in all_concepts if c.derivation != PurposeLineage.CONSTANT]
-    constants = [c for c in all_concepts if c.derivation == PurposeLineage.CONSTANT]
+    non_constant = [c for c in all_concepts if c.derivation != Derivation.CONSTANT]
+    constants = [c for c in all_concepts if c.derivation == Derivation.CONSTANT]
     if not non_constant and constants:
         return ConstantNode(
             output_concepts=constants,
@@ -398,7 +498,7 @@ def gen_select_merge_node(
             non_constant,
             accept_partial=attempt,
             conditions=conditions,
-            datasources=list(environment.datasources.values()),
+            datasources=list([x for x in environment.datasources.values()]),
             depth=depth,
         )
         if pruned_concept_graph:
@@ -411,9 +511,12 @@ def gen_select_merge_node(
         logger.info(f"{padding(depth)}{LOGGER_PREFIX} no covering graph found.")
         return None
 
-    sub_nodes = resolve_subgraphs(pruned_concept_graph, conditions)
+    sub_nodes = resolve_subgraphs(
+        pruned_concept_graph, relevant=non_constant, conditions=conditions
+    )
 
     logger.info(f"{padding(depth)}{LOGGER_PREFIX} fetching subgraphs {sub_nodes}")
+
     parents = [
         create_select_node(
             k,
@@ -444,6 +547,10 @@ def gen_select_merge_node(
 
     if len(parents) == 1:
         return parents[0]
+    logger.info(
+        f"{padding(depth)}{LOGGER_PREFIX} Multiple parent DS nodes resolved - {[type(x) for x in parents]}, wrapping in merge"
+    )
+
     preexisting_conditions = None
     if conditions and all(
         [
@@ -461,15 +568,5 @@ def gen_select_merge_node(
         parents=parents,
         preexisting_conditions=preexisting_conditions,
     )
-    target_grain = Grain.from_concepts(all_concepts)
-    if not base.resolve().grain.issubset(target_grain):
-        return GroupNode(
-            output_concepts=all_concepts,
-            input_concepts=all_concepts,
-            environment=environment,
-            parents=[base],
-            depth=depth,
-            preexisting_conditions=preexisting_conditions,
-            partial_concepts=base.partial_concepts,
-        )
+
     return base

trilogy/core/processing/node_generators/select_node.py

@@ -1,12 +1,12 @@
 from trilogy.constants import logger
-from trilogy.core.enums import PurposeLineage
+from trilogy.core.enums import Derivation
 from trilogy.core.exceptions import NoDatasourceException
-from trilogy.core.models import (
-    Concept,
-    Environment,
-    LooseConceptList,
-    WhereClause,
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildWhereClause,
+    LooseBuildConceptList,
 )
+from trilogy.core.models.build_environment import BuildEnvironment
 from trilogy.core.processing.node_generators.select_merge_node import (
     gen_select_merge_node,
 )
@@ -19,23 +19,23 @@ LOGGER_PREFIX = "[GEN_SELECT_NODE]"
 
 
 def gen_select_node(
-    concept: Concept,
-    local_optional: list[Concept],
-    environment: Environment,
+    concept: BuildConcept,
+    local_optional: list[BuildConcept],
+    environment: BuildEnvironment,
     g,
     depth: int,
     accept_partial: bool = False,
     fail_if_not_found: bool = True,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
 ) -> StrategyNode | None:
     all_concepts = [concept] + local_optional
-    all_lcl = LooseConceptList(concepts=all_concepts)
-    materialized_lcl = LooseConceptList(
+    all_lcl = LooseBuildConceptList(concepts=all_concepts)
+    materialized_lcl = LooseBuildConceptList(
         concepts=[
             x
             for x in all_concepts
             if x.address in environment.materialized_concepts
-            or x.derivation == PurposeLineage.CONSTANT
+            or x.derivation == Derivation.CONSTANT
         ]
     )
     if materialized_lcl != all_lcl:

trilogy/core/processing/node_generators/union_node.py

@@ -2,26 +2,29 @@ from typing import List
 
 from trilogy.constants import logger
 from trilogy.core.enums import FunctionType, Purpose
-from trilogy.core.models import Concept, Function, WhereClause
+from trilogy.core.models.build import BuildConcept, BuildFunction, BuildWhereClause
 from trilogy.core.processing.nodes import History, StrategyNode, UnionNode
 from trilogy.core.processing.utility import padding
 
 LOGGER_PREFIX = "[GEN_UNION_NODE]"
 
 
-def is_union(c: Concept):
-    return isinstance(c.lineage, Function) and c.lineage.operator == FunctionType.UNION
+def is_union(c: BuildConcept):
+    return (
+        isinstance(c.lineage, BuildFunction)
+        and c.lineage.operator == FunctionType.UNION
+    )
 
 
 def gen_union_node(
-    concept: Concept,
-    local_optional: List[Concept],
+    concept: BuildConcept,
+    local_optional: List[BuildConcept],
     environment,
     g,
     depth: int,
     source_concepts,
     history: History | None = None,
-    conditions: WhereClause | None = None,
+    conditions: BuildWhereClause | None = None,
 ) -> StrategyNode | None:
     all_unions = [x for x in local_optional if is_union(x)] + [concept]
 
@@ -30,15 +33,13 @@ def gen_union_node(
     base = keys.pop()
     remaining = [x for x in all_unions if x.address != base.address]
     arguments = []
-    if isinstance(base.lineage, Function):
+    if isinstance(base.lineage, BuildFunction):
         arguments = base.lineage.concept_arguments
     for arg in arguments:
-        relevant_parents: list[Concept] = []
+        relevant_parents: list[BuildConcept] = []
         for other_union in remaining:
             assert other_union.lineage
-            potential_parents = [
-                z for z in other_union.lineage.arguments if isinstance(z, Concept)
-            ]
+            potential_parents = [z for z in other_union.lineage.concept_arguments]
             relevant_parents += [
                 x for x in potential_parents if x.keys and arg.address in x.keys
             ]