diff --git a/madeleine/base.py b/madeleine/base.py index 0c12424..c5dfa69 100644 --- a/madeleine/base.py +++ b/madeleine/base.py @@ -3,14 +3,13 @@ from collections.abc import Mapping, Iterable from enum import Enum from pathlib import Path from objtools.registry import ClassRegistry -import json import random -import yaml class FileFormat(Enum): JSON = 'json' YAML = 'yaml' + TOML = 'toml' class ComponentRegistry(ClassRegistry): @@ -60,11 +59,34 @@ class Component(metaclass=ComponentMetaclass, register=False): if not isinstance(path, Path): path = Path(path) with path.open() as f: - method = json.load if fmt == FileFormat.JSON else yaml.safe_load + if fmt == FileFormat.JSON: + import json + data = json.load(f) + elif fmt == FileFormat.YAML: + try: + import yaml + except ImportError as e: + raise ImportError( + '{}\nYou may need to install madeleine with YAML support: ' + 'pip install madeleine[yaml]'.format(e.message), + ) + data = yaml.safe_load(f) + elif fmt == FileFormat.TOML: + try: + import toml + except ImportError as e: + raise ImportError( + '{}\nYou may need to install madeleine with TOML support: ' + 'pip install madeleine[toml]'.format(e.message), + ) + data = toml.load(f) + else: + raise NotImplementedError + if cls is Component: # Guess the component type when calling Component.from_path - return Component._make(method(f)) - return cls(**method(f)) + return Component._make(data) + return cls(**data) @staticmethod def _make(data): diff --git a/requirements.txt b/requirements.txt index b8303d0..a6f46e9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1 @@ objtools>=0.1 -pyyaml>=5.1 diff --git a/setup.py b/setup.py index 4c3c793..34f5fd2 100755 --- a/setup.py +++ b/setup.py @@ -23,6 +23,8 @@ setup( install_requires=requirements, extras_require={ 'dev': dev_requirements, + 'yaml': ['pyyaml>=5.1'], + 'toml': ['toml>=0.10'], }, tests_require=dev_requirements, test_suite='tests',