feat(ini): allow custom parsers/serializers for ini options

Signed-off-by: Sylvain Leclerc <[email protected]>

antarest/study/storage/rawstudy/ini_reader.py

@@ -15,13 +15,25 @@
 import typing as t
 from abc import ABC, abstractmethod
 from pathlib import Path
+from typing import Callable, Optional
 
 from typing_extensions import override
 
 from antarest.core.model import JSON
 
+PrimitiveType = t.Union[str, int, float, bool]
+ValueParser = Callable[[str], PrimitiveType]
+SelectionPredicate = Callable[[str], bool]
 
-def convert_value(value: str) -> t.Union[str, int, float, bool]:
+
+def _lower_case(input: str) -> str:
+    return input.lower()
+
+
+LOWER_CASE_PARSER: ValueParser = _lower_case
+
+
+def _convert_value(value: str) -> PrimitiveType:
     """Convert value to the appropriate type for JSON."""
 
     try:
@@ -38,7 +50,44 @@ def convert_value(value: str) -> t.Union[str, int, float, bool]:
                 return value
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
+class OptionMatcher:
+    """
+    Defines a location in INI file data.
+    A None section means all sections.
+    """
+
+    section: Optional[str]
+    key: Optional[str]
+
+
+def any_section_option_matcher(key: str) -> OptionMatcher:
+    """
+    Return a matcher which will match the provided key in any section.
+    """
+    return OptionMatcher(section=None, key=key)
+
+
+class ValueParsers:
+    _parsers: t.Dict[OptionMatcher, ValueParser]
+
+    def __init__(self, parsers: t.Dict[OptionMatcher, ValueParser]):
+        self._parsers = parsers
+
+    def find_parser(self, section: str, key: str) -> ValueParser:
+        if not self._parsers:
+            return _convert_value
+        possible_keys = [
+            OptionMatcher(section=section, key=key),
+            OptionMatcher(section=None, key=key),
+        ]
+        for k in possible_keys:
+            if parser := self._parsers.get(k, None):
+                return parser
+        return _convert_value
+
+
+@dataclasses.dataclass(frozen=True)
 class IniFilter:
     """
     Filter sections and options in an INI file based on regular expressions.
@@ -115,8 +164,8 @@ def read(self, path: t.Any, **kwargs: t.Any) -> JSON:
         Parse `.ini` file to json object.
 
         Args:
-            path: Path to `.ini` file or file-like object.
-            kwargs: Additional options used for reading.
+            path:    Path to `.ini` file or file-like object.
+            options: Additional options used for reading.
 
         Returns:
             Dictionary of parsed `.ini` file which can be converted to JSON.
@@ -152,11 +201,17 @@ class IniReader(IReader):
     This class is not compatible with standard `.ini` readers.
     """
 
-    def __init__(self, special_keys: t.Sequence[str] = (), section_name: str = "settings") -> None:
+    def __init__(
+        self,
+        special_keys: t.Sequence[str] = (),
+        section_name: str = "settings",
+        value_parsers: t.Dict[OptionMatcher, ValueParser] | None = None,
+    ) -> None:
         super().__init__()
 
         # Default section name to use if `.ini` file has no section.
         self._special_keys = set(special_keys)
+        self._value_parsers = ValueParsers(value_parsers or {})
 
         # List of keys which should be parsed as list.
         self._section_name = section_name
@@ -313,10 +368,12 @@ def _handle_option(self, ini_filter: IniFilter, section: str, key: str, value: s
     def _append_option(self, section: str, key: str, value: str) -> None:
         self._curr_sections.setdefault(section, {})
         values = self._curr_sections[section]
+        parser = self._value_parsers.find_parser(section, key)
+        parsed = parser(value)
         if key in self._special_keys:
-            values.setdefault(key, []).append(convert_value(value))
+            values.setdefault(key, []).append(parsed)
         else:
-            values[key] = convert_value(value)
+            values[key] = parsed
         self._curr_option = key
 
 

antarest/study/storage/rawstudy/ini_writer.py

@@ -12,24 +12,52 @@
 
 import ast
 import configparser
-import typing as t
 from pathlib import Path
+from typing import Callable, Dict, List, Optional, Union
 
 from typing_extensions import override
 
 from antarest.core.model import JSON
+from antarest.study.storage.rawstudy.ini_reader import OptionMatcher
+
+PrimitiveType = Union[str, int, float, bool]
+ValueSerializer = Callable[[str], PrimitiveType]
+
+
+def _lower_case(input: str) -> str:
+    return input.lower()
+
+
+LOWER_CASE_SERIALIZER: ValueSerializer = _lower_case
 
 
 class IniConfigParser(configparser.RawConfigParser):
-    def __init__(self, special_keys: t.Optional[t.List[str]] = None) -> None:
+    def __init__(
+        self,
+        special_keys: Optional[List[str]] = None,
+        value_serializers: Optional[Dict[OptionMatcher, ValueSerializer]] = None,
+    ) -> None:
         super().__init__()
         self.special_keys = special_keys
+        self._value_serializers = value_serializers or {}
 
     # noinspection SpellCheckingInspection
     @override
     def optionxform(self, optionstr: str) -> str:
         return optionstr
 
+    def _get_serializer(self, section: str, key: str) -> Optional[ValueSerializer]:
+        if not self._value_serializers:
+            return None
+        possible_keys = [
+            OptionMatcher(section=section, key=key),
+            OptionMatcher(section=None, key=key),
+        ]
+        for k in possible_keys:
+            if parser := self._value_serializers.get(k, None):
+                return parser
+        return None
+
     def _write_line(  # type:ignore
         self,
         delimiter,
@@ -41,6 +69,9 @@ def _write_line(  # type:ignore
         value = self._interpolation.before_write(  # type:ignore
             self, section_name, key, value
         )
+        if self._value_serializers:
+            if serializer := self._get_serializer(section_name, key):
+                value = serializer(value)
         if value is not None or not self._allow_no_value:  # type:ignore
             value = delimiter + str(value).replace("\n", "\n\t")
         else:
@@ -70,8 +101,13 @@ class IniWriter:
     Standard INI writer.
     """
 
-    def __init__(self, special_keys: t.Optional[t.List[str]] = None):
+    def __init__(
+        self,
+        special_keys: Optional[List[str]] = None,
+        value_serializers: Optional[Dict[OptionMatcher, ValueSerializer]] = None,
+    ):
         self.special_keys = special_keys
+        self._value_serializers = value_serializers or {}
 
     def write(self, data: JSON, path: Path) -> None:
         """
@@ -81,7 +117,7 @@ def write(self, data: JSON, path: Path) -> None:
             data: JSON content.
             path: path to `.ini` file.
         """
-        config_parser = IniConfigParser(special_keys=self.special_keys)
+        config_parser = IniConfigParser(special_keys=self.special_keys, value_serializers=self._value_serializers)
         config_parser.read_dict(data)
         with path.open("w") as fp:
             config_parser.write(fp)

tests/storage/repository/antares_io/reader/test_ini_reader.py

@@ -14,7 +14,12 @@
 import textwrap
 from pathlib import Path
 
-from antarest.study.storage.rawstudy.ini_reader import IniReader, SimpleKeyValueReader
+from antarest.study.storage.rawstudy.ini_reader import (
+    IniReader,
+    OptionMatcher,
+    SimpleKeyValueReader,
+    any_section_option_matcher,
+)
 
 
 class TestIniReader:
@@ -324,6 +329,33 @@ def test_read__filtered_option(self, tmp_path) -> None:
         expected = {"part1": {"bar": "hello"}, "part2": {"bar": "salut"}}
         assert actual == expected
 
+    def test_read__with_custom_parser(self, tmp_path):
+        path = Path(tmp_path) / "test.ini"
+        path.write_text(
+            textwrap.dedent(
+                """
+                [part1]
+                bar = Hello
+                
+                [part2]
+                bar = Hello
+                """
+            )
+        )
+
+        def to_lower(input: str) -> str:
+            return input.lower()
+
+        value_parsers = {OptionMatcher("part2", "bar"): to_lower}
+        actual = IniReader(value_parsers=value_parsers).read(path)
+        expected = {"part1": {"bar": "Hello"}, "part2": {"bar": "hello"}}
+        assert actual == expected
+
+        value_parsers = {any_section_option_matcher("bar"): to_lower}
+        actual = IniReader(value_parsers=value_parsers).read(path)
+        expected = {"part1": {"bar": "hello"}, "part2": {"bar": "hello"}}
+        assert actual == expected
+
 
 class TestSimpleKeyValueReader:
     def test_read(self) -> None:

tests/storage/repository/antares_io/writer/test_ini_writer.py

@@ -15,6 +15,7 @@
 
 import pytest
 
+from antarest.study.storage.rawstudy.ini_reader import OptionMatcher, any_section_option_matcher
 from antarest.study.storage.rawstudy.ini_writer import IniWriter
 
 
@@ -59,3 +60,31 @@ def test_write(tmp_path: str, ini_cleaner: Callable) -> None:
     writer.write(json_data, path)
 
     assert ini_cleaner(ini_content) == ini_cleaner(path.read_text())
+
+
+@pytest.mark.unit_test
+def test_write_with_custom_serializer(tmp_path: str, ini_cleaner: Callable) -> None:
+    path = Path(tmp_path) / "test.ini"
+
+    serializers = {any_section_option_matcher("group"): lambda x: x.lower()}
+    writer = IniWriter(value_serializers=serializers)
+
+    expected = """
+        [part1]
+        group = gas
+
+        [part2]
+        group = gas
+
+        [part3]
+        other = Gas
+    """
+
+    json_data = {
+        "part1": {"group": "Gas"},
+        "part2": {"group": "Gas"},
+        "part3": {"other": "Gas"},
+    }
+    writer.write(json_data, path)
+
+    assert ini_cleaner(path.read_text()) == ini_cleaner(expected)