From 1fff44fbeb5a71db7b807efd26fc99a93de7de3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Morais?=
 <146729917+SMoraisAnsys@users.noreply.github.com>
Date: Tue, 10 Dec 2024 16:15:09 +0100
Subject: [PATCH] BUILD: Modify TOML dependency (#5548)

Co-authored-by: Maxime Rey <87315832+MaxJPRey@users.noreply.github.com>
---
 pyproject.toml                                |  3 +-
 .../aedt/core/generic/general_methods.py      | 17 ++++-------
 src/ansys/aedt/core/modules/material_lib.py   |  2 +-
 tests/unit/test_utils.py                      | 29 +++++++++++++++++++
 4 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 53955317fed..33af53982f2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,8 @@ dependencies = [
     "pyedb>=0.4.0; python_version == '3.7'",
     "pyedb>=0.24.0; python_version > '3.7'",
     "pyedb!=0.28.0; python_version > '3.7'",
-    "pytomlpp; python_version < '3.12'",
+    "tomli; python_version < '3.11'",
+    "tomli-w",
     "rpyc>=6.0.0,<6.1",
     "pyyaml",
     "defusedxml>=0.7,<8.0"
diff --git a/src/ansys/aedt/core/generic/general_methods.py b/src/ansys/aedt/core/generic/general_methods.py
index dc7d593af7a..fcb8ac29873 100644
--- a/src/ansys/aedt/core/generic/general_methods.py
+++ b/src/ansys/aedt/core/generic/general_methods.py
@@ -514,11 +514,10 @@ def read_toml(file_path):  # pragma: no cover
     dict
         Parsed TOML file as a dictionary.
     """
-    current_python_version = sys.version_info[:2]
-    if current_python_version < (3, 12):
-        import pytomlpp as tomllib
-    else:
+    try:
         import tomllib
+    except (ImportError, ModuleNotFoundError):
+        import tomli as tomllib
 
     with open_file(file_path, "rb") as fb:
         return tomllib.load(fb)
@@ -1298,11 +1297,7 @@ def is_digit(c):
 
 @pyaedt_function_handler()
 def _create_toml_file(input_dict, full_toml_path):
-    current_python_version = sys.version_info[:2]
-    if current_python_version < (3, 12):
-        import pytomlpp as tomllib
-    else:
-        import tomllib
+    import tomli_w
 
     if not os.path.exists(os.path.dirname(full_toml_path)):
         os.makedirs(os.path.dirname(full_toml_path))
@@ -1322,8 +1317,8 @@ def _dict_toml(d):
         return new_dict
 
     new_dict = _dict_toml(input_dict)
-    with open_file(full_toml_path, "w") as fp:
-        tomllib.dump(new_dict, fp)
+    with open_file(full_toml_path, "wb") as fp:
+        tomli_w.dump(new_dict, fp)
     settings.logger.info(f"{full_toml_path} correctly created.")
     return True
 
diff --git a/src/ansys/aedt/core/modules/material_lib.py b/src/ansys/aedt/core/modules/material_lib.py
index 2ab3dfb7749..1d72f91cf7d 100644
--- a/src/ansys/aedt/core/modules/material_lib.py
+++ b/src/ansys/aedt/core/modules/material_lib.py
@@ -733,7 +733,7 @@ def _aedmattolibrary(self, matname):
 
     @pyaedt_function_handler(full_json_path="output_file")
     def export_materials_to_file(self, output_file):
-        """Export all materials to a JSON  or TOML file.
+        """Export all materials to a JSON or TOML file.
 
         Parameters
         ----------
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 915e4c33b80..49bc5638e03 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -43,6 +43,7 @@
 SETTINGS_RELEASE_ON_EXCEPTION = settings.release_on_exception
 SETTINGS_ENABLE_ERROR_HANDLER = settings.enable_error_handler
 ERROR_MESSAGE = "Dummy message."
+TOML_DATA = {"key_0": "dummy", "key_1": 12, "key_2": [1, 2], "key_3": {"key_4": 42}}
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -237,3 +238,31 @@ def test_settings_check_allowed_env_variables():
         allowed_env_var_expected.remove("ANS_NODEPCHECK")
 
     assert sorted(allowed_env_var_expected) == sorted(env_variables)
+
+
+def test_read_toml(tmp_path):
+    """Test loading a TOML file."""
+    from ansys.aedt.core.generic.general_methods import read_toml
+
+    file_path = tmp_path / "dummy.toml"
+    content = """
+    key_0 = 'dummy'
+    key_1 = 12
+    key_2 = [1,2]
+    [key_3]
+    key_4 = 42
+    """
+    file_path.write_text(content, encoding="utf-8")
+
+    res = read_toml(file_path)
+    assert TOML_DATA == res
+
+
+def test_write_toml(tmp_path):
+    """Test writing a TOML file."""
+    from ansys.aedt.core.generic.general_methods import _create_toml_file
+
+    file_path = tmp_path / "dummy.toml"
+    _create_toml_file(TOML_DATA, file_path)
+
+    assert file_path.exists()