diff --git a/setup.py b/setup.py index 4367f81f..e817c6a1 100644 --- a/setup.py +++ b/setup.py @@ -108,7 +108,8 @@ ] extras_require = setuptools_args['extras_require'] = { - 'test': ['pytest'], + 'test': ['pytest', 'toml'], + 'with-toml': ['toml'], } if 'setuptools' in sys.modules: diff --git a/traitlets/config/loader.py b/traitlets/config/loader.py index d6d90ee4..e2246524 100644 --- a/traitlets/config/loader.py +++ b/traitlets/config/loader.py @@ -9,6 +9,11 @@ import re import sys import json +try: + import toml + HAS_TOML = True +except ImportError: + HAS_TOML = None import warnings from ..utils import cast_unicode @@ -543,6 +548,7 @@ def _find_file(self): """Try to find the file by searching the paths.""" self.full_filename = filefind(self.filename, self.path) + class JSONFileConfigLoader(FileConfigLoader): """A JSON file loader for config @@ -598,6 +604,69 @@ def __exit__(self, exc_type, exc_value, traceback): f.write(json_config) +class TOMLFileConfigLoader(FileConfigLoader): + """A TOML file loader for config + + Can also act as a context manager that rewrite the configuration file to disk on exit. + + Example:: + + with TOMLFileConfigLoader('myapp.toml','/home/jupyter/configurations/') as c: + c.MyNewConfigurable.new_value = 'Updated' + + """ + + def __init__(self, filename, **kw): + """Wrapper for checking (optional) toml module import""" + if not HAS_TOML: + raise ConfigLoaderError('toml module is not found. In order to use toml configuration' + 'files, please either install traitlets with corresponding option' + ' (pip install "traitlets[with-toml]") or simply add it with' + '"pip install toml"') + super(TOMLFileConfigLoader, self).__init__(filename, **kw) + + def load_config(self): + """Load the config from a file and return it as a Config object.""" + self.clear() + try: + self._find_file() + except IOError as e: + raise ConfigFileNotFound(str(e)) + dct = self._read_file_as_dict() + self.config = self._convert_to_config(dct) + return self.config + + def _read_file_as_dict(self): + with open(self.full_filename) as f: + return toml.load(f) + + def _convert_to_config(self, dictionary): + if 'version' in dictionary: + version = dictionary.pop('version') + else: + version = 1 + + if version == 1: + return Config(dictionary) + else: + raise ValueError('Unknown version of TOML config file: {version}'.format(version=version)) + + def __enter__(self): + self.load_config() + return self.config + + def __exit__(self, exc_type, exc_value, traceback): + """ + Exit the context manager but do not handle any errors. + + In case of any error, we do not want to write the potentially broken + configuration to disk. + """ + self.config.version = 1 + toml_config = toml.dumps(self.config) + with open(self.full_filename, 'w') as f: + f.write(toml_config) + class PyFileConfigLoader(FileConfigLoader): """A config loader for pure python files. diff --git a/traitlets/config/tests/test_loader.py b/traitlets/config/tests/test_loader.py index 74289dc9..10d38f28 100644 --- a/traitlets/config/tests/test_loader.py +++ b/traitlets/config/tests/test_loader.py @@ -18,6 +18,7 @@ LazyConfigValue, PyFileConfigLoader, JSONFileConfigLoader, + TOMLFileConfigLoader, KeyValueConfigLoader, ArgParseConfigLoader, KVArgParseConfigLoader, @@ -69,6 +70,27 @@ } """ +toml_file = """ +# This is a TOML document. + +version = 1 + +a = 10 +b = 20 + +[Foo] + # Indentation (tabs and/or spaces) is allowed but not required + [Foo.Bam] + value = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + + [Foo.Bar] + value = 10 + +[D] + [D.C] + value = 'hi there' +""" + import logging log = logging.getLogger('devnull') log.setLevel(0) @@ -102,6 +124,32 @@ def test_json(self): config = cl.load_config() self._check_conf(config) + def test_toml(self): + fd, fname = mkstemp('.toml', prefix='μnïcø∂e') + f = os.fdopen(fd, 'w') + f.write(toml_file) + f.close() + # Unlink the file + cl = TOMLFileConfigLoader(fname, log=log) + config = cl.load_config() + self._check_conf(config) + + def test_optional_toml(self): + import traitlets.config.loader as loader + from traitlets.config.loader import ConfigLoaderError + loader.HAS_TOML = False + fd, fname = mkstemp('.toml', prefix='μnïcø∂e') + f = os.fdopen(fd, 'w') + f.write(toml_file) + f.close() + error_raised = False + try: + cl = TOMLFileConfigLoader(fname, log=log) + except ConfigLoaderError: + error_raised = True + loader.HAS_TOML = True + assert error_raised is True + def test_context_manager(self): fd, fname = mkstemp('.json', prefix='μnïcø∂e')