Skip to content
Snippets Groups Projects
test_classes.py 1.58 KiB
Newer Older
Frederik Hennig's avatar
Frederik Hennig committed
# type: ignore
from pystencilssfg import SourceFileGenerator, SfgConfiguration, SfgComposer
from pystencilssfg.configuration import SfgCodeStyle
Frederik Hennig's avatar
Frederik Hennig committed
from pystencils.types import PsCustomType
from pystencilssfg.ir.source_components import SfgClass, SfgMemberVariable, SfgConstructor, SfgMethod
Frederik Hennig's avatar
Frederik Hennig committed

from pystencils import fields, kernel

sfg_config = SfgConfiguration(
    output_directory="out/test_classes",
    outer_namespace="gen_code",
    codestyle=SfgCodeStyle(
        code_style="Mozilla",
        force_clang_format=True
    )
Frederik Hennig's avatar
Frederik Hennig committed
)

f, g = fields("f, g(1): double[2D]")

with SourceFileGenerator(sfg_config) as ctx:
    sfg = SfgComposer(ctx)
Frederik Hennig's avatar
Frederik Hennig committed

    @kernel
    def assignments():
        f[0,0] @= 3 * g[0,0]

    khandle = sfg.kernels.create(assignments)

    cls = SfgClass("MyClass")
    cls.default.append_member(SfgMethod(
Frederik Hennig's avatar
Frederik Hennig committed
        "callKernel",
        sfg.call(khandle)
    cls.default.append_member(SfgMethod(
        "inlineConst",
        sfg.seq(
            "return -1.0;"
        ),
Frederik Hennig's avatar
Frederik Hennig committed
        return_type="double",
        inline=True,
        const=True
    ))

    cls.default.append_member(SfgMethod(
        "awesomeMethod",
        sfg.seq(
            "return 2.0f;"
        ),
Frederik Hennig's avatar
Frederik Hennig committed
        return_type="float",
        inline=False,
        const=True
    ))

    cls.default.append_member(
Frederik Hennig's avatar
Frederik Hennig committed
        SfgMemberVariable(
Frederik Hennig's avatar
Frederik Hennig committed
            "stuff", PsCustomType("std::vector< int > &")
    cls.default.append_member(
Frederik Hennig's avatar
Frederik Hennig committed
        SfgConstructor(
Frederik Hennig's avatar
Frederik Hennig committed
            [sfg.var("stuff", PsCustomType("std::vector< int > &"))],
            ["stuff_(stuff)"]
Frederik Hennig's avatar
Frederik Hennig committed
        )
    )

    sfg.context.add_class(cls)