Skip to content
Snippets Groups Projects
conditional.py 1.61 KiB
Newer Older
Frederik Hennig's avatar
Frederik Hennig committed
from __future__ import annotations
from typing import TYPE_CHECKING, Sequence, Optional

if TYPE_CHECKING:
    from ..context import SfgContext

from jinja2.filters import do_indent
from pystencils.typing import TypedSymbol

from .basic_nodes import SfgCallTreeNode, SfgCallTreeLeaf

class SfgCondition(SfgCallTreeLeaf):
    pass

class SfgCustomCondition(SfgCondition):
    def __init__(self, cond_text: str):
        self._cond_text = cond_text

    def required_symbols(self) -> set(TypedSymbol):
        return set()
    
    def defined_symbols(self) -> set(TypedSymbol):
        return set()

    def get_code(self, ctx: SfgContext) -> str:
        return self._cond_text
    

# class IntEven(SfgCondition):
#     def __init__(self, )


class SfgBranch(SfgCallTreeNode):
    def __init__(self, cond: SfgCondition, branch_true: SfgCallTreeNode, branch_false: Optional[SfgCallTreeNode] = None):
        self._cond = cond
        self._branch_true = branch_true
        self._branch_false = branch_false
    
    @property
    def children(self) -> Sequence[SfgCallTreeNode]:
        if self._branch_false is not None:
            return (self._branch_true, self._branch_false)
        else:
            return (self._branch_true,)
        
    def get_code(self, ctx: SfgContext) -> str:
        code = f"if({self._cond.get_code(ctx)}) {{\n"
        code += ctx.codestyle.indent(self._branch_true.get_code(ctx))
        code += "\n}"
        
        if self._branch_false is not None:
            code += "else {\n"
            code += ctx.codestyle.indent(self._branch_false.get_code(ctx))
            code += "\n}"

        return code