config_manager.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. # coding: utf-8
  2. """Manager to read and modify config data in JSON files."""
  3. # Copyright (c) Jupyter Development Team.
  4. # Distributed under the terms of the Modified BSD License.
  5. import errno
  6. import glob
  7. import io
  8. import json
  9. import os
  10. import copy
  11. from six import PY3
  12. from traitlets.config import LoggingConfigurable
  13. from traitlets.traitlets import Unicode, Bool
  14. def recursive_update(target, new):
  15. """Recursively update one dictionary using another.
  16. None values will delete their keys.
  17. """
  18. for k, v in new.items():
  19. if isinstance(v, dict):
  20. if k not in target:
  21. target[k] = {}
  22. recursive_update(target[k], v)
  23. if not target[k]:
  24. # Prune empty subdicts
  25. del target[k]
  26. elif v is None:
  27. target.pop(k, None)
  28. else:
  29. target[k] = v
  30. def remove_defaults(data, defaults):
  31. """Recursively remove items from dict that are already in defaults"""
  32. # copy the iterator, since data will be modified
  33. for key, value in list(data.items()):
  34. if key in defaults:
  35. if isinstance(value, dict):
  36. remove_defaults(data[key], defaults[key])
  37. if not data[key]: # prune empty subdicts
  38. del data[key]
  39. else:
  40. if value == defaults[key]:
  41. del data[key]
  42. class BaseJSONConfigManager(LoggingConfigurable):
  43. """General JSON config manager
  44. Deals with persisting/storing config in a json file with optionally
  45. default values in a {section_name}.d directory.
  46. """
  47. config_dir = Unicode('.')
  48. read_directory = Bool(True)
  49. def ensure_config_dir_exists(self):
  50. """Will try to create the config_dir directory."""
  51. try:
  52. os.makedirs(self.config_dir, 0o755)
  53. except OSError as e:
  54. if e.errno != errno.EEXIST:
  55. raise
  56. def file_name(self, section_name):
  57. """Returns the json filename for the section_name: {config_dir}/{section_name}.json"""
  58. return os.path.join(self.config_dir, section_name+'.json')
  59. def directory(self, section_name):
  60. """Returns the directory name for the section name: {config_dir}/{section_name}.d"""
  61. return os.path.join(self.config_dir, section_name+'.d')
  62. def get(self, section_name, include_root=True):
  63. """Retrieve the config data for the specified section.
  64. Returns the data as a dictionary, or an empty dictionary if the file
  65. doesn't exist.
  66. When include_root is False, it will not read the root .json file,
  67. effectively returning the default values.
  68. """
  69. paths = [self.file_name(section_name)] if include_root else []
  70. if self.read_directory:
  71. pattern = os.path.join(self.directory(section_name), '*.json')
  72. # These json files should be processed first so that the
  73. # {section_name}.json take precedence.
  74. # The idea behind this is that installing a Python package may
  75. # put a json file somewhere in the a .d directory, while the
  76. # .json file is probably a user configuration.
  77. paths = sorted(glob.glob(pattern)) + paths
  78. self.log.debug('Paths used for configuration of %s: \n\t%s', section_name, '\n\t'.join(paths))
  79. data = {}
  80. for path in paths:
  81. if os.path.isfile(path):
  82. with io.open(path, encoding='utf-8') as f:
  83. recursive_update(data, json.load(f))
  84. return data
  85. def set(self, section_name, data):
  86. """Store the given config data.
  87. """
  88. filename = self.file_name(section_name)
  89. self.ensure_config_dir_exists()
  90. if self.read_directory:
  91. # we will modify data in place, so make a copy
  92. data = copy.deepcopy(data)
  93. defaults = self.get(section_name, include_root=False)
  94. remove_defaults(data, defaults)
  95. # Generate the JSON up front, since it could raise an exception,
  96. # in order to avoid writing half-finished corrupted data to disk.
  97. json_content = json.dumps(data, indent=2)
  98. if PY3:
  99. f = io.open(filename, 'w', encoding='utf-8')
  100. else:
  101. f = open(filename, 'wb')
  102. with f:
  103. f.write(json_content)
  104. def update(self, section_name, new_data):
  105. """Modify the config section by recursively updating it with new_data.
  106. Returns the modified config data as a dictionary.
  107. """
  108. data = self.get(section_name)
  109. recursive_update(data, new_data)
  110. self.set(section_name, data)
  111. return data