improve: extract method for safe loading yaml file and avoid using PyYaml's FullLoader (#4031)
Bowen Liang
2 years ago
committed by
GitHub
parent
296887754f
commit
3fda2245a4
@ -0,0 +1,34 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import yaml
|
||||
from yaml import YAMLError
|
||||
|
||||
|
||||
def load_yaml_file(file_path: str, ignore_error: bool = False) -> dict:
|
||||
"""
|
||||
Safe loading a YAML file to a dict
|
||||
:param file_path: the path of the YAML file
|
||||
:param ignore_error:
|
||||
if True, return empty dict if error occurs and the error will be logged in warning level
|
||||
if False, raise error if error occurs
|
||||
:return: a dict of the YAML content
|
||||
"""
|
||||
try:
|
||||
if not file_path or not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f'Failed to load YAML file {file_path}: file not found')
|
||||
|
||||
with open(file_path, encoding='utf-8') as file:
|
||||
try:
|
||||
return yaml.safe_load(file)
|
||||
except Exception as e:
|
||||
raise YAMLError(f'Failed to load YAML file {file_path}: {e}')
|
||||
except FileNotFoundError as e:
|
||||
logging.debug(f'Failed to load YAML file {file_path}: {e}')
|
||||
return {}
|
||||
except Exception as e:
|
||||
if ignore_error:
|
||||
logging.warning(f'Failed to load YAML file {file_path}: {e}')
|
||||
return {}
|
||||
else:
|
||||
raise e
|
||||
@ -0,0 +1,34 @@
|
||||
from textwrap import dedent
|
||||
|
||||
import pytest
|
||||
|
||||
from core.utils.position_helper import get_position_map
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prepare_example_positions_yaml(tmp_path, monkeypatch) -> str:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
tmp_path.joinpath("example_positions.yaml").write_text(dedent(
|
||||
"""\
|
||||
- first
|
||||
- second
|
||||
# - commented
|
||||
- third
|
||||
|
||||
- 9999999999999
|
||||
- forth
|
||||
"""))
|
||||
return str(tmp_path)
|
||||
|
||||
|
||||
def test_position_helper(prepare_example_positions_yaml):
|
||||
position_map = get_position_map(
|
||||
folder_path=prepare_example_positions_yaml,
|
||||
file_name='example_positions.yaml')
|
||||
assert len(position_map) == 4
|
||||
assert position_map == {
|
||||
'first': 0,
|
||||
'second': 1,
|
||||
'third': 2,
|
||||
'forth': 3,
|
||||
}
|
||||
@ -0,0 +1,74 @@
|
||||
from textwrap import dedent
|
||||
|
||||
import pytest
|
||||
from yaml import YAMLError
|
||||
|
||||
from core.tools.utils.yaml_utils import load_yaml_file
|
||||
|
||||
EXAMPLE_YAML_FILE = 'example_yaml.yaml'
|
||||
INVALID_YAML_FILE = 'invalid_yaml.yaml'
|
||||
NON_EXISTING_YAML_FILE = 'non_existing_file.yaml'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prepare_example_yaml_file(tmp_path, monkeypatch) -> str:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
file_path = tmp_path.joinpath(EXAMPLE_YAML_FILE)
|
||||
file_path.write_text(dedent(
|
||||
"""\
|
||||
address:
|
||||
city: Example City
|
||||
country: Example Country
|
||||
age: 30
|
||||
gender: male
|
||||
languages:
|
||||
- Python
|
||||
- Java
|
||||
- C++
|
||||
empty_key:
|
||||
"""))
|
||||
return str(file_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prepare_invalid_yaml_file(tmp_path, monkeypatch) -> str:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
file_path = tmp_path.joinpath(INVALID_YAML_FILE)
|
||||
file_path.write_text(dedent(
|
||||
"""\
|
||||
address:
|
||||
city: Example City
|
||||
country: Example Country
|
||||
age: 30
|
||||
gender: male
|
||||
languages:
|
||||
- Python
|
||||
- Java
|
||||
- C++
|
||||
"""))
|
||||
return str(file_path)
|
||||
|
||||
|
||||
def test_load_yaml_non_existing_file():
|
||||
assert load_yaml_file(file_path=NON_EXISTING_YAML_FILE) == {}
|
||||
assert load_yaml_file(file_path='') == {}
|
||||
|
||||
|
||||
def test_load_valid_yaml_file(prepare_example_yaml_file):
|
||||
yaml_data = load_yaml_file(file_path=prepare_example_yaml_file)
|
||||
assert len(yaml_data) > 0
|
||||
assert yaml_data['age'] == 30
|
||||
assert yaml_data['gender'] == 'male'
|
||||
assert yaml_data['address']['city'] == 'Example City'
|
||||
assert set(yaml_data['languages']) == {'Python', 'Java', 'C++'}
|
||||
assert yaml_data.get('empty_key') is None
|
||||
assert yaml_data.get('non_existed_key') is None
|
||||
|
||||
|
||||
def test_load_invalid_yaml_file(prepare_invalid_yaml_file):
|
||||
# yaml syntax error
|
||||
with pytest.raises(YAMLError):
|
||||
load_yaml_file(file_path=prepare_invalid_yaml_file)
|
||||
|
||||
# ignore error
|
||||
assert load_yaml_file(file_path=prepare_invalid_yaml_file, ignore_error=True) == {}
|
||||
Loading…
Reference in New Issue