Source code for SDF.file_io.convert

import logging
import os
import warnings
from contextlib import contextmanager
from enum import Enum
from io import BytesIO, StringIO
from typing import List, Optional, Callable, Union, BinaryIO, TextIO

from SDF.data_model import SDFObject
from SDF.file_io.force_sdf import load_from_force_sdf, sdf_to_force_sdf, is_convertible_to_force_sdf, force_sdf_to_mat
from SDF.file_io.jpk import load_from_jpk
from SDF.file_io.lsm import load_from_lsm
from SDF.file_io.mfp import load_from_mfp
from SDF.file_io.oif import load_from_oib, load_from_oif
from SDF.file_io.sdf import load_from_sdf, write_sdf

logger = logging.getLogger(__name__)


class OutputFormat(Enum):
    """Possible output formats for SDF converters"""
    SDF = 0
    FORCE = 1
    MAT = 2


class InputFormat(Enum):
    """Possible input formats for SDF converters"""
    SDF = 0
    JPK = 1
    MFP = 2
    OIB = 3
    OIF = 4
    LSM = 5
    FORCE_SDF = 6
    JPK_MAP = 7


class SDFConverter:
    """Handles SDF file conversion"""
    input_format: Optional[InputFormat]
    output_directory: Optional[str]

    def __init__(self, input_format: Optional[InputFormat] = None, output_directory: Optional[str] = None,
                 generate_mat: bool = False, generate_force: bool = True, dry_run: bool = False) -> None:
        self.input_format = input_format
        self.output_directory = output_directory
        self.generate_mat = generate_mat
        self.generate_force = generate_force
        self.dry_run = dry_run

        # output directory
        if self.output_directory is not None and not os.path.isdir(self.output_directory):
            logger.info(f"Output directory '{self.output_directory}' does not exist. Creating it.")
            os.makedirs(self.output_directory)

    @staticmethod
    def get_load_function(input_format: InputFormat) -> Callable[[str], SDFObject]:
        """Determines the load function for a input format"""
        if input_format is InputFormat.MFP:
            return load_from_mfp
        elif input_format is InputFormat.OIB:
            return load_from_oib
        elif input_format is InputFormat.OIF:
            return load_from_oif
        elif input_format is InputFormat.JPK:
            return load_from_jpk
        elif input_format is InputFormat.LSM:
            return load_from_lsm
        elif input_format is InputFormat.SDF:
            return load_from_sdf
        elif input_format is InputFormat.FORCE_SDF:
            return load_from_force_sdf
        elif input_format is InputFormat.JPK_MAP:
            return load_from_jpk
        else:
            raise RuntimeError(f"Unknown input format: {input_format}")

    def infer_output_formats(self, sdf: SDFObject, input_format: InputFormat) -> List[OutputFormat]:
        """Infers the output formats for writing the SDF object"""
        output_formats = []
        if input_format is not InputFormat.SDF:
            output_formats.append(OutputFormat.SDF)

        if is_convertible_to_force_sdf(sdf):
            if self.generate_force:
                output_formats.append(OutputFormat.FORCE)
            if self.generate_mat:
                output_formats.append(OutputFormat.MAT)

        return output_formats

    @staticmethod
    def infer_input_format(original_file: str) -> InputFormat:
        """Infers the input format from the filename"""
        logger.info(f"Inferring input format for file '{original_file}'")
        if original_file.endswith(".jpk-force"):
            input_format = InputFormat.JPK
        elif original_file.endswith(".jpk-proc-force"):
            input_format = InputFormat.JPK
        elif original_file.endswith(".jpk-force-map"):
            input_format = InputFormat.JPK_MAP
        elif original_file.endswith(".jpk-nt-force"):
            input_format = InputFormat.JPK
        elif original_file.endswith(".ibw"):
            input_format = InputFormat.MFP
        elif original_file.endswith(".force-sdf"):
            input_format = InputFormat.FORCE_SDF
        elif original_file.endswith(".force.sdf"):
            input_format = InputFormat.FORCE_SDF
        elif original_file.endswith(".sdf"):  # has to be checked after `.force.sdf`
            input_format = InputFormat.SDF
        elif original_file.endswith(".lsm"):
            input_format = InputFormat.LSM
        elif original_file.endswith(".oif"):
            input_format = InputFormat.OIF
        elif original_file.endswith(".oib"):
            input_format = InputFormat.OIB
        else:
            raise RuntimeError(f"Failed to infer input format from filename '{original_file}'. "
                               f"If the format should be obvious, please contact the maintainers.")

        logger.info(f"Determined input file to be of type {input_format.name}")
        return input_format

    def convert_file(self, input_file: str) -> None:
        """Converts the input file"""
        with warnings_to_logging():
            logger.info(f"Starting conversion of file {input_file}.")

            if self.input_format is None:
                input_format = SDFConverter.infer_input_format(input_file)
            else:
                input_format = self.input_format

            logger.info(f"Reading file as type {input_format.name}")
            sdf = SDFConverter.get_load_function(input_format)(input_file)

            output_formats = self.infer_output_formats(sdf, input_format)
            if OutputFormat.SDF in output_formats:
                output_file = self.__generate_output_file(input_file, OutputFormat.SDF)
                logger.info(f"Saving to {output_file}")
                write_sdf(sdf, output_file)

            if OutputFormat.FORCE in output_formats or OutputFormat.MAT in output_formats:
                logger.info("Reading as ForceSDF object")
                f_sdf = sdf_to_force_sdf(sdf)

                if OutputFormat.FORCE in output_formats:
                    output_file = self.__generate_output_file(input_file, OutputFormat.FORCE)
                    logger.info(f"Saving to {output_file}")
                    write_sdf(f_sdf, output_file)

                if OutputFormat.MAT in output_formats:
                    output_file = self.__generate_output_file(input_file, OutputFormat.MAT)
                    logger.info(f"Saving to {output_file}")
                    force_sdf_to_mat(f_sdf, output_file)

    def __generate_output_file(self, original_file: str, output_format: OutputFormat) -> Union[str, TextIO, BinaryIO]:
        directory, name = os.path.split(os.path.abspath(original_file))
        basename, original_extension = os.path.splitext(name)

        if self.output_directory is None:
            _output_directory = directory
        else:
            _output_directory = self.output_directory

        if output_format is OutputFormat.SDF:
            extension = ".sdf"
            mode = str
        elif output_format is OutputFormat.FORCE:
            extension = ".force.sdf"
            mode = str
        elif output_format is OutputFormat.MAT:
            extension = ".mat"
            mode = bytes
        else:
            raise ValueError(f"Unknown output format: {output_format}")

        if self.dry_run:
            if mode is bytes:
                return BytesIO()
            elif mode is str:
                return StringIO()
            else:
                raise ValueError(f"Unknown mode: {mode}")

        return f"{_output_directory}/{basename}{extension}"


@contextmanager
def warnings_to_logging():
    """Context manager that converts warnings into logging messages"""
    with warnings.catch_warnings():
        orig_formatwarning = warnings.formatwarning

        logging.captureWarnings(True)  # send warnings to logger "py.warnings"

        def formatwarning(message, category, filename, lineno, line=None):
            # same interface as warnings.formatwarning
            return f"{message} ({category.__name__} caught in {filename}:{lineno})"

        warnings.formatwarning = formatwarning

        try:
            yield  # do whatever is done inside the context manager
        finally:
            # restore previous behavior
            logging.captureWarnings(False)
            warnings.formatwarning = orig_formatwarning