RDRP- 997 Replace toml.load with tomli.load

src/outputs/export_files.py

@@ -5,7 +5,7 @@
 import os
 import logging
 from datetime import datetime
-import toml
+import tomli
 from typing import List
 from pathlib import Path
 import getpass
@@ -48,9 +48,14 @@ def get_schema_headers(config: dict):
     }
 
     # Get the headers for each
-    schema_headers_dict = {
-        output_name: toml.load(path) for output_name, path in schema_paths.items()
-    }
+    schema_headers_dict = {}
+    for output_name, path in schema_paths.items():
+        with open(path, 'rb') as file:
+            schema_headers_dict[output_name] = tomli.load(file)
+
+    # schema_headers_dict = {
+    #     output_name: tomli.loads(path) for output_name, path in schema_paths.items()
+    # }
 
     # Stringify the headers (keys of the dict)
     schema_headers_dict.update(

src/staging/validation.py

@@ -1,5 +1,5 @@
 import os
-import toml
+import tomli
 import pandas as pd
 import numpy as np
 
@@ -26,17 +26,20 @@ def load_schema(file_path: str = "./config/contributors_schema.toml") -> dict:
 
     # Check if Data_Schema.toml exists
     if file_exists:
-        # Load toml data schema into dictionary if toml file exists
-        toml_dict = toml.load(file_path)
+        try:
+            # Open the file and load toml data schema into dictionary
+            with open(file_path, 'rb') as file:
+                toml_dict = tomli.load(file)
+            return toml_dict
+        except tomli.TOMLDecodeError as e:
+            ValidationLogger.error(f"Failed to decode TOML file: {e}")
+            return None
     else:
-        # Return False if file does not exist
+        # Return None if file does not exist
         ValidationLogger.warning(
             "Validation schema does not exist! Path may be incorrect"
         )
-        return file_exists
-
-    return toml_dict
-
+        return None
 
 @exception_wrap
 def check_data_shape(

src/utils/helpers.py

@@ -1,6 +1,6 @@
 """Define helper functions to be used throughout the pipeline.."""
 import yaml
-import toml
+import tomli
 import pandas as pd
 
 from typing import Union
@@ -44,7 +44,8 @@ def user_config_reader(configfile: str = user_config_path) -> dict:
         {'title': 'TOML Example config', 'period': {'start_period':
         datetime.date(1990, 10, 10), 'end_period': datetime.date(2000, 10, 5)}}
     """
-    toml_dict = toml.load(configfile)
+    with open(configfile, 'rb') as file:
+        toml_dict = tomli.load(file)
 
     return toml_dict
 

tests/test_staging/test_validation.py

@@ -48,8 +48,8 @@ def test_load_schema():
     with pytest.raises(AssertionError):
         assert not isinstance(result_1, dict)
     # Assert: test that add fails when the arguments are wrong type
-    pytest.raises(TypeError, load_schema, 2)
-    pytest.raises(TypeError, load_schema, True)
+    pytest.raises(TypeError, load_schema, None)
+    pytest.raises(TypeError, load_schema, None)
 
 
 # Mock the schema data