# -*- coding: utf-8 -*-

"""Predicates for checking nodes' variants."""

from typing import Tuple, Type, Union

from .utils import node_predicate
from ..typing import NodePredicate
from ....dsl import BaseEntity, CentralDogma, Fragment, GeneModification, Hgvs, ProteinModification, Variant

__all__ = [
    'has_variant',
    'has_protein_modification',
    'has_gene_modification',
    'has_fragment',
    'has_hgvs',
]


@node_predicate
def has_variant(node: BaseEntity) -> bool:
    """Return true if the node has any variants."""
    return isinstance(node, CentralDogma) and node.variants


def _variant_checker(variant_cls: Union[Type[Variant], Tuple[Type[Variant], ...]]) -> NodePredicate:
    @node_predicate
    def _node_has_variant(node: BaseEntity) -> bool:
        """Return true if the node has at least one of the given variant."""
        return isinstance(node, CentralDogma) and node.variants and any(
            isinstance(variant, variant_cls)
            for variant in node.variants
        )

    return _node_has_variant


has_protein_modification = _variant_checker(ProteinModification)
has_gene_modification = _variant_checker(GeneModification)
has_hgvs = _variant_checker(Hgvs)
has_fragment = _variant_checker(Fragment)

# -*- coding: utf-8 -*-

"""Predicates for checking nodes' variants."""

from functools import wraps
from typing import Tuple, Type, Union

from .utils import node_predicate
from ..typing import NodePredicate
from ....dsl import BaseEntity, CentralDogma, Fragment, GeneModification, Hgvs, ProteinModification, Variant

__all__ = [
    'has_variant',
    'has_protein_modification',
    'has_gene_modification',
    'has_fragment',
    'has_hgvs',
]


@node_predicate
def has_variant(node: BaseEntity) -> bool:
    """Return true if the node has any variants."""
    return isinstance(node, CentralDogma) and node.variants


def _variant_checker(variant_cls: Union[Type[Variant], Tuple[Type[Variant], ...]]) -> NodePredicate:
    @node_predicate
    @wraps(node_has_variant)
    def _rv(node: BaseEntity):
        return node_has_variant(node, variant_cls)

    return _rv


def node_has_variant(node: BaseEntity, variant_cls) -> bool:
    """Return true if the node has at least one of the given variant."""
    return isinstance(node, CentralDogma) and node.variants and any(
        isinstance(variant, variant_cls)
        for variant in node.variants
    )


has_protein_modification = _variant_checker(ProteinModification)
has_gene_modification = _variant_checker(GeneModification)
has_hgvs = _variant_checker(Hgvs)
has_fragment = _variant_checker(Fragment)

Fix docs