cmake_integration.py

"""
Mechanism for integrating code generation into walberla's CMake system.
CMake needs to determine which C++ source files are generated by which Python script.
The list of files should be available fast, without running the generation process itself.
Thus all code generation function are registered at a single class that manages this process.

Usage example:
    from pystencils_walberla.cmake_integration import codegen
    codegen.register(['MyClass.h', 'MyClass.cpp'], functionReturningTwoStringsForHeaderAndCpp)

"""
import json
import sys
import os
import warnings

__all__ = ['CodeGeneration', 'ManualCodeGenerationContext']


class CodeGeneration:
    def __init__(self):
        expected_files, cmake_vars = parse_json_args()
        self.context = CodeGenerationContext(cmake_vars)
        self.expected_files = expected_files

    def __enter__(self):
        return self.context

    def __exit__(self, *args):
        if self.expected_files and (set(self.context.files_written) != set(self.expected_files)):
            only_in_cmake = set(self.expected_files) - set(self.context.files_written)
            only_generated = set(self.context.files_written) - set(self.expected_files)
            error_message = "Generated files specified not correctly in cmake with 'waLBerla_python_file_generates'\n"
            if only_in_cmake:
                error_message += "Files only specified in CMake {}\n".format([os.path.basename(p) for p in only_in_cmake])
            if only_generated:
                error_message += "Unexpected generated files {}\n".format([os.path.basename(p) for p in only_generated])
            raise ValueError(error_message)


def parse_json_args():
    default = {'EXPECTED_FILES': [],
               'CMAKE_VARS': {'WALBERLA_BUILD_WITH_OPENMP': False,
                              'WALBERLA_OPTIMIZE_FOR_LOCALHOST': False,
                              'WALBERLA_DOUBLE_ACCURACY': True,
                              'WALBERLA_BUILD_WITH_MPI': True,
                              'WALBERLA_BUILD_WITH_CUDA': False}
               }

    if len(sys.argv) == 2:
        try:
            parsed = json.loads(sys.argv[1])
        except json.JSONDecodeError:
            warnings.warn("Could not parse JSON arguments")
            parsed = default
    else:
        parsed = default
    expected_files = parsed['EXPECTED_FILES']
    cmake_vars = {}
    for key, value in parsed['CMAKE_VARS'].items():
        if value in ("ON", "1", "YES"):
            value = True
        elif value in ("OFF", "0", "NO"):
            value = False
        cmake_vars[key] = value
    return expected_files, cmake_vars


class CodeGenerationContext:
    def __init__(self, cmake_vars):
        self.files_written = []
        self.openmp = cmake_vars['WALBERLA_BUILD_WITH_OPENMP']
        self.optimize_for_localhost = cmake_vars['WALBERLA_OPTIMIZE_FOR_LOCALHOST']
        self.mpi = cmake_vars['WALBERLA_BUILD_WITH_MPI']
        self.double_accuracy = cmake_vars['WALBERLA_DOUBLE_ACCURACY']
        self.cuda = cmake_vars['WALBERLA_BUILD_WITH_CUDA']

    def write_file(self, name, content):
        self.files_written.append(os.path.abspath(name))
        with open(name, 'w') as f:
            f.write(content)


class ManualCodeGenerationContext:
    """Context for testing - does not actually write files but puts them into a public dict
    Environment parameters like if OpenMP, MPI or CPU-specific optimization should be used can be explicitly passed
    to constructor instead of getting them from CMake
    """
    def __init__(self, openmp=False, optimize_for_localhost=False, mpi=True, double_accuracy=True):
        self.openmp = openmp
        self.optimize_for_localhost = optimize_for_localhost
        self.mpi = mpi
        self.double_accuracy = double_accuracy
        self.files = dict()
        self.cuda = False

    def write_file(self, name, content):
        self.files[name] = content

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass