Source code for flow360_schema.framework.entity.entity_selector

"""
Entity selector models

Defines a minimal, stable schema for selecting entities by rules.
"""

import logging
import re
from collections import deque
from collections.abc import Callable
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Literal, cast

import pydantic as pd
from typing_extensions import Self

from flow360_schema.framework.base_model import Flow360BaseModel
from flow360_schema.framework.entity.entity_base import EntityBase
from flow360_schema.framework.entity.entity_utils import (
    DEFAULT_NOT_MERGED_TYPES,
    deduplicate_entities,
    generate_uuid,
)
from flow360_schema.framework.glob_pattern import compile_glob

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from flow360_schema.framework.entity.entity_list import EntityList

# These corresponds to the private_attribute_entity_type_name of supported entity types.
TargetClass = Literal["Surface", "Edge", "GenericVolume", "GeometryBodyGroup"]

EntityNode = EntityBase | dict[str, Any]


class Predicate(Flow360BaseModel):
    """
    Single predicate in a selector.
    """

    # For now only name matching is supported
    attribute: Literal["name"] = pd.Field("name", description="The attribute to match/filter on.")
    operator: Literal[
        "any_of",
        "not_any_of",
        "matches",
        "not_matches",
    ] = pd.Field()
    value: str | list[str] = pd.Field()
    # Applies only to matches/notMatches; default to glob if not specified explicitly.
    non_glob_syntax: Literal["regex"] | None = pd.Field(
        None,
        description="If specified, the pattern (`value`) will be treated "
        "as a non-glob pattern with the specified syntax.",
    )

    @pd.field_validator("value", mode="before")
    @classmethod
    def _normalize_membership_value(cls, value: Any, info: pd.ValidationInfo) -> Any:
        """Normalize single-string membership input to a one-item list."""
        if info.data.get("operator") in ("any_of", "not_any_of") and isinstance(value, str):
            return [value]
        return value

    @pd.model_validator(mode="after")
    def _reject_glob_outside_the_dialect(self) -> Self:
        """A glob spelling the dialect does not support fails here, not silently at first use."""
        if self.operator in ("matches", "not_matches") and self.non_glob_syntax is None:
            compile_glob(cast(str, self.value))
        return self


class EntitySelector(Flow360BaseModel):
    """Entity selector for an EntityList.

    - target_class chooses the entity pool
    - logic combines child predicates (AND = intersection, OR = union)
    """

    target_class: TargetClass = pd.Field()
    description: str | None = pd.Field(None, description="Customizable description of the selector.")
    selector_id: str = pd.Field(
        default_factory=generate_uuid,
        description="[Internal] Unique identifier for the selector.",
        frozen=True,
    )
    # Unique name for global reuse
    name: str = pd.Field(description="Unique name for this selector.")
    logic: Literal["AND", "OR"] = pd.Field("AND")
    children: list[Predicate] = pd.Field()

    @pd.validate_call
    def match(
        self,
        pattern: str,
        *,
        attribute: Literal["name"] = "name",
    ) -> Self:
        """Append a matches predicate and return self for chaining.

        Glob dialect: flow360_schema.framework.glob_pattern.
        """
        self.children.append(
            Predicate(  # type: ignore[call-arg]
                attribute=attribute,
                operator="matches",
                value=pattern,
            )
        )
        return self

    @pd.validate_call
    def not_match(
        self,
        pattern: str,
        *,
        attribute: Literal["name"] = "name",
    ) -> Self:
        """Append a not-matches predicate and return self for chaining.

        Glob dialect: flow360_schema.framework.glob_pattern.
        """
        self.children.append(
            Predicate(  # type: ignore[call-arg]
                attribute=attribute,
                operator="not_matches",
                value=pattern,
            )
        )
        return self

    @pd.validate_call
    def any_of(self, values: list[str], *, attribute: Literal["name"] = "name") -> Self:
        """Append an in predicate and return self for chaining."""
        self.children.append(Predicate(attribute=attribute, operator="any_of", value=values))  # type: ignore[call-arg]
        return self

    @pd.validate_call
    def not_any_of(self, values: list[str], *, attribute: Literal["name"] = "name") -> Self:
        """Append a notIn predicate and return self for chaining."""
        self.children.append(Predicate(attribute=attribute, operator="not_any_of", value=values))  # type: ignore[call-arg]
        return self


########## SELECTOR CLASSES ##########


[docs] class SurfaceSelector(EntitySelector): """ Pattern-based selector for Surface entities. Stores matching rules, enabling reusable SimulationParams templates. Example ------- >>> # Simple: match surfaces by glob pattern >>> fl.SurfaceSelector(name="wing_surfaces").match("wing*") >>> # With OR logic: match surfaces matching either pattern >>> sel = fl.SurfaceSelector( ... name="wing_and_tail", logic="OR", ... description="Wing and tail surfaces" ... ).match("wing*").match("tail*") ==== """ target_class: Literal["Surface"] = pd.Field("Surface", frozen=True) children: list[Predicate] = pd.Field(default_factory=list)
[docs] class EdgeSelector(EntitySelector): """ Pattern-based selector for Edge entities. Stores matching rules, enabling reusable SimulationParams templates. Example ------- >>> # Simple: match edges by glob pattern >>> fl.EdgeSelector(name="leading_edges").match("leading_edge*") >>> # With OR logic: match edges by name list >>> sel = fl.EdgeSelector( ... name="selected_edges", logic="OR", ... description="Edges for refinement" ... ).any_of(["edge1", "edge2"]) ==== """ target_class: Literal["Edge"] = pd.Field("Edge", frozen=True) children: list[Predicate] = pd.Field(default_factory=list)
[docs] class VolumeSelector(EntitySelector): """ Pattern-based selector for volume zone entities. Stores matching rules, enabling reusable SimulationParams templates. Example ------- >>> # Simple: match volumes by glob pattern >>> fl.VolumeSelector(name="fluid_zones").match("fluid*") >>> # With OR logic: match specific zones by name >>> sel = fl.VolumeSelector( ... name="porous_zones", logic="OR", ... description="Zones for porous medium" ... ).any_of(["zone1", "zone2"]) ==== """ target_class: Literal["GenericVolume"] = pd.Field("GenericVolume", frozen=True) children: list[Predicate] = pd.Field(default_factory=list)
[docs] class BodyGroupSelector(EntitySelector): """ Pattern-based selector for body group entities. Stores matching rules, enabling reusable SimulationParams templates. Example ------- >>> # Simple: match body groups by glob pattern >>> fl.BodyGroupSelector(name="rotor_bodies").match("rotor*") >>> # With OR logic: match specific bodies by name >>> sel = fl.BodyGroupSelector( ... name="rotating_bodies", logic="OR", ... description="Bodies for rotation" ... ).any_of(["body1", "body2"]) ==== """ target_class: Literal["GeometryBodyGroup"] = pd.Field("GeometryBodyGroup", frozen=True) children: list[Predicate] = pd.Field(default_factory=list)
########## EXPANSION IMPLEMENTATION ########## @lru_cache(maxsize=2048) def _compile_regex_cached(pattern: str) -> re.Pattern[str]: return re.compile(pattern) def _get_node_attribute(entity: Any, attribute: str) -> Any: """Return attribute value from either dicts or entity objects.""" if isinstance(entity, dict): return entity.get(attribute) return getattr(entity, attribute, None) def _get_attribute_value(entity: Any, attribute: str) -> str | None: """Return the scalar string value of an attribute, or None if absent/unsupported. Only scalar string attributes are supported by this matcher layer for now. """ val = _get_node_attribute(entity, attribute) if isinstance(val, str): return val return None def _build_value_matcher(predicate: Predicate) -> Callable[[str | None], bool]: """ Build a fast predicate(value: Optional[str])->bool matcher. Precompiles regex/glob and converts membership lists to sets for speed. """ operator = predicate.operator value = predicate.value non_glob_syntax = predicate.non_glob_syntax negate = False if operator in ("not_any_of", "not_matches"): negate = True base_operator = { "not_any_of": "any_of", "not_matches": "matches", }.get(operator) else: base_operator = operator if base_operator == "any_of": values = {value} if isinstance(value, str) else set(value or []) def base_match(val: str | None) -> bool: return val in values elif base_operator == "matches": pattern = _compile_regex_cached(value) if non_glob_syntax == "regex" else compile_glob(value) # type: ignore[arg-type] def base_match(val: str | None) -> bool: return isinstance(val, str) and (pattern.fullmatch(val) is not None) else: def base_match(_val: str | None) -> bool: # type: ignore[misc] return False if negate: return lambda val: not base_match(val) return base_match def _build_index(pool: list[EntityNode], attribute: str) -> dict[str, list[int]]: """Build an index for in lookups on a given attribute.""" value_to_indices: dict[str, list[int]] = {} for idx, item in enumerate(pool): val = _get_node_attribute(item, attribute) if isinstance(val, str): value_to_indices.setdefault(val, []).append(idx) return value_to_indices def _apply_or_selector( pool: list[EntityNode], ordered_children: list[Predicate], ) -> list[Any]: indices: set[int] = set() for predicate in ordered_children: attribute = predicate.attribute matcher = _build_value_matcher(predicate) for i, item in enumerate(pool): if i in indices: continue if matcher(_get_attribute_value(item, attribute)): indices.add(i) if len(indices) >= len(pool): break if len(indices) * 4 < len(pool): return [pool[i] for i in sorted(indices)] return [pool[i] for i in range(len(pool)) if i in indices] def _apply_and_selector( pool: list[EntityNode], ordered_children: list[Predicate], indices_by_attribute: dict[str, dict[str, list[int]]], ) -> list[Any]: candidate_indices: set[int] | None = None def _matched_indices_for_predicate(predicate: Predicate, current_candidates: set[int] | None) -> set[int]: operator = predicate.operator attribute = predicate.attribute if operator == "any_of": idx_map = indices_by_attribute.get(attribute) if idx_map is not None: result: set[int] = set() for v in predicate.value or []: result.update(idx_map.get(v, [])) return result matcher = _build_value_matcher(predicate) matched: set[int] = set() if current_candidates is None: for i, item in enumerate(pool): if matcher(_get_attribute_value(item, attribute)): matched.add(i) return matched for i in current_candidates: if matcher(_get_attribute_value(pool[i], attribute)): matched.add(i) return matched for predicate in ordered_children: matched = _matched_indices_for_predicate(predicate, candidate_indices) candidate_indices = matched if candidate_indices is None else candidate_indices.intersection(matched) if not candidate_indices: return [] assert candidate_indices is not None if len(candidate_indices) * 4 < len(pool): return [pool[i] for i in sorted(candidate_indices)] return [pool[i] for i in range(len(pool)) if i in candidate_indices] def _apply_single_selector(pool: list[EntityNode], selector: EntitySelector) -> list[EntityNode]: """Apply one selector over a pool of entities (dicts or objects). Implementation notes for future readers: - We assume selector_dict conforms to the EntitySelector schema (no validation here for speed). - We respect the default of logic="AND" when absent. - For performance: * Reorder predicates under AND so that cheap/selective operations run first. * Build a name->indices index to accelerate equals/in where beneficial. * Precompile regex/glob matchers once per predicate. * Short-circuit when the candidate set becomes empty. - Result ordering is stable (by original pool index) to keep the operation idempotent. """ logic = selector.logic children = selector.children # Fast path: empty predicates -> return nothing. Empty children is actually misuse. if not children: return [] # Predicate ordering (AND only): cheap/selective first def _cost(predicate: Predicate) -> int: op = predicate.operator order = { "any_of": 0, "matches": 1, "not_any_of": 2, "not_matches": 3, } return order.get(op, 10) ordered_children = children if logic == "OR" else sorted(children, key=_cost) # Optional per-attribute indices for in attributes_needing_index = {p.attribute for p in ordered_children if p.operator == "any_of"} indices_by_attribute: dict[str, dict[str, list[int]]] = ( {attr: _build_index(pool, attr) for attr in attributes_needing_index} if attributes_needing_index else {} ) if logic == "OR": # Favor a full scan for OR to preserve predictable union behavior # and avoid over-indexing that could complicate ordering. result = _apply_or_selector(pool, ordered_children) else: result = _apply_and_selector(pool, ordered_children, indices_by_attribute) if not result: name = selector.name target_class = selector.target_class logger.warning( "Entity selector '%s' (target_class=%s) matched 0 entities. " "Please check if the entity name or pattern is correct.", name, target_class, ) return result def _get_selector_cache_key(selector: EntitySelector) -> tuple[str, str]: """ Return the cache key for a selector: requires unique name. We mandate a unique identifier per selector; use ("name", target_class, name) for stable global reuse. If neither `name` is provided, fall back to a structural key so different unnamed selectors won't collide. """ # selector_id is always present and unique by schema. return ("selector_id", selector.selector_id) def _process_selectors( registry: Any, selectors_list: list[Any], selector_cache: dict[str, Any], *, expansion_map: dict[str, list[str]] | None = None, ) -> tuple[dict[str, list[EntityNode]], list[str]]: """Process selectors and return additions grouped by class. This function iterates over a list of materialized selectors (EntitySelector-like objects or plain dicts) and applies them over the entity registry. Results are cached in `selector_cache` to avoid re-computation for the same selector. Parameters: registry: EntityRegistry instance containing entities. selectors_list: List of selector definitions (materialized; no string tokens). selector_cache: Cache for selector results. expansion_map: Type expansion mapping. If None, uses DEFAULT_TARGET_CLASS_EXPANSION_MAP from the local schema module (lazy import). Returns: Tuple of (additions_by_class dict, ordered_target_classes list). """ # Set default expansion map via lazy import from the local schema module. if expansion_map is None: from flow360_schema.framework.entity.entity_expansion_config import ( DEFAULT_TARGET_CLASS_EXPANSION_MAP, ) expansion_map = DEFAULT_TARGET_CLASS_EXPANSION_MAP additions_by_class: dict[str, list[EntityNode]] = {} ordered_target_classes: list[str] = [] for item in selectors_list: if not isinstance(item, EntitySelector): raise TypeError(f"[Internal] selectors_list must contain EntitySelector objects. Got: {type(item)}") selector: EntitySelector = item target_class = selector.target_class # Use expansion map to get multiple type names expanded_type_names = expansion_map.get(target_class, [target_class]) entities = registry.find_by_type_name(expanded_type_names) if not entities: continue cache_key = _get_selector_cache_key(selector) additions = selector_cache.get(cache_key) # type: ignore[call-overload] if additions is None: additions = _apply_single_selector(entities, selector) selector_cache[cache_key] = additions # type: ignore[index] if target_class not in additions_by_class: additions_by_class[target_class] = [] ordered_target_classes.append(target_class) additions_by_class[target_class].extend(additions) return additions_by_class, ordered_target_classes def _merge_entities( existing: list[EntityNode], additions_by_class: dict[str, list[EntityNode]], ordered_target_classes: list[str], merge_mode: Literal["merge", "replace"], not_merged_types: frozenset[str] = DEFAULT_NOT_MERGED_TYPES, ) -> list[Any]: """Merge existing entities with selector additions based on merge mode. Type filtering is intentionally out of scope here. Callers are expected to filter selector additions before passing them into this merge helper. """ candidates: list[EntityNode] = [] if merge_mode == "merge": # explicit first, then selector additions candidates.extend(existing) for target_class in ordered_target_classes: candidates.extend(additions_by_class.get(target_class, [])) else: # replace: drop explicit items of targeted classes classes_to_update = set(ordered_target_classes) for item in existing: entity_type = _get_node_attribute(item, "private_attribute_entity_type_name") if entity_type not in classes_to_update: candidates.append(item) for target_class in ordered_target_classes: candidates.extend(additions_by_class.get(target_class, [])) # Deduplication logic (same as materialize_entities_and_selectors_in_place) return deduplicate_entities( candidates, not_merged_types=not_merged_types, ) def expand_entity_list_selectors( registry: Any, entity_list: Any, *, selector_cache: dict[str, Any] | None = None, merge_mode: Literal["merge", "replace"] = "merge", expansion_map: dict[str, list[str]] | None = None, ) -> list[EntityNode]: """ Expand selectors in a single EntityList within an EntityRegistry context. Parameters ---------- expansion_map : Optional type expansion mapping for selectors. Notes ----- - This function does NOT modify the input EntityList. - selector_cache can be shared across multiple calls to reuse selector results. - This helper performs raw selector expansion only. It does not apply any EntityList-specific type filtering. """ stored_entities = list(getattr(entity_list, "stored_entities", []) or []) raw_selectors = list(getattr(entity_list, "selectors", []) or []) if selector_cache is None: selector_cache = {} if not raw_selectors: return stored_entities additions_by_class, ordered_target_classes = _process_selectors( registry, raw_selectors, selector_cache, expansion_map=expansion_map, ) return _merge_entities( stored_entities, additions_by_class, ordered_target_classes, merge_mode, ) def resolve_entity_list_selectors( registry: Any, entity_list: "EntityList", *, selector_cache: dict[str, Any] | None = None, merge_mode: Literal["merge", "replace"] = "merge", expansion_map: dict[str, list[str]] | None = None, ) -> list[EntityBase]: """ Expand selectors for a concrete EntityList and filter selector additions to its valid types. Explicit ``stored_entities`` are assumed to have already been pre-filtered by EntityList's ``before`` validator. Only selector additions are type-filtered here before they are merged into the resolved entity list. """ from flow360_schema.framework.entity.entity_list import ( _filter_entities_by_valid_types, _get_valid_entity_type_info, ) stored_entities = list(getattr(entity_list, "stored_entities", []) or []) raw_selectors = list(getattr(entity_list, "selectors", []) or []) if selector_cache is None: selector_cache = {} if not raw_selectors: return stored_entities additions_by_class, ordered_target_classes = _process_selectors( registry, raw_selectors, selector_cache, expansion_map=expansion_map, ) raw_selector_entities = [] filtered_additions_by_class: dict[str, list[EntityNode]] = {} for target_class in ordered_target_classes: additions = additions_by_class.get(target_class, []) entity_additions = [item for item in additions if isinstance(item, EntityBase)] raw_selector_entities.extend(entity_additions) filtered_additions_by_class[target_class] = cast( list[EntityNode], _filter_entities_by_valid_types( entity_list.__class__, entity_additions, ), ) merged_entities = _merge_entities( stored_entities, filtered_additions_by_class, ordered_target_classes, merge_mode, ) if raw_selector_entities and not merged_entities: valid_types, _ = _get_valid_entity_type_info(entity_list.__class__) valid_type_name_list = [valid_type.__name__ for valid_type in valid_types] raise ValueError(f"Can not find any valid entity of type {valid_type_name_list} from the input.") return merged_entities def collect_and_tokenize_selectors_in_place( params_as_dict: dict[str, Any], ) -> dict[str, Any]: """ Collect selector definitions from a fresh JSON params dict and replace them with tokens (`selector_id`). Intended usage: - Input is the result of `model_dump(mode="json")` before selector tokenization. - This is a single-pass preprocessing step for upload/validation payload preparation. - It is not intended to be re-run on a params dict whose `selectors` entries are already tokens. Behavior: 1. Traverse `params_as_dict` to find inline `EntitySelector` definitions. 2. Move these definitions into `private_attribute_asset_cache["used_selectors"]`. 3. Replace each selector definition in a `selectors` list with its `selector_id` token. """ known_selectors = {} # Pre-populate from existing AssetCache if any asset_cache = params_as_dict.setdefault("private_attribute_asset_cache", {}) if isinstance(asset_cache, dict): cached_selectors = asset_cache.get("used_selectors") if isinstance(cached_selectors, list): for s in cached_selectors: selector_id = s.get("selector_id") if selector_id is None: selector_id = generate_uuid() s["selector_id"] = selector_id known_selectors[selector_id] = s queue: deque[dict[str, Any] | list[Any]] = deque([params_as_dict]) while queue: node = queue.popleft() if isinstance(node, dict): selectors = node.get("selectors") new_selectors: list[Any] = [] for item in selectors or (): selector_id = item.get("selector_id") if selector_id is None: selector_id = generate_uuid() item["selector_id"] = selector_id known_selectors[selector_id] = item new_selectors.append(selector_id) if selectors is not None: node["selectors"] = new_selectors # Recurse for value in node.values(): if value is asset_cache: continue if isinstance(value, (dict, list)): queue.append(value) elif isinstance(node, list): for item in node: if isinstance(item, (dict, list)): queue.append(item) # Update AssetCache if isinstance(asset_cache, dict): asset_cache.pop("used_selectors", None) asset_cache["used_selectors"] = list(known_selectors.values()) return params_as_dict