Coverage for sacred/sacred/config/config_files.py: 97%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

35 statements  

1#!/usr/bin/env python 

2# coding=utf-8 

3 

4import os 

5import pickle 

6 

7import json 

8 

9import sacred.optional as opt 

10from sacred.serializer import flatten, restore 

11 

12__all__ = ("load_config_file", "save_config_file") 

13 

14 

15class Handler: 

16 def __init__(self, load, dump, mode): 

17 self.load = load 

18 self.dump = dump 

19 self.mode = mode 

20 

21 

22HANDLER_BY_EXT = { 

23 ".json": Handler( 

24 lambda fp: restore(json.load(fp)), 

25 lambda obj, fp: json.dump(flatten(obj), fp, sort_keys=True, indent=2), 

26 "", 

27 ), 

28 ".pickle": Handler(pickle.load, pickle.dump, "b"), 

29} 

30 

31yaml_extensions = (".yaml", ".yml") 

32if opt.has_yaml: 

33 

34 def load_yaml(filename): 

35 return opt.yaml.load(filename, Loader=opt.yaml.FullLoader) 

36 

37 yaml_handler = Handler(load_yaml, opt.yaml.dump, "") 

38 

39 for extension in yaml_extensions: 

40 HANDLER_BY_EXT[extension] = yaml_handler 

41 

42 

43def get_handler(filename): 

44 _, extension = os.path.splitext(filename) 

45 if extension in yaml_extensions and not opt.has_yaml: 

46 raise KeyError( 

47 'Configuration file "{}" cannot be loaded as ' 

48 "you do not have PyYAML installed.".format(filename) 

49 ) 

50 try: 

51 return HANDLER_BY_EXT[extension] 

52 except KeyError: 

53 raise ValueError( 

54 'Configuration file "{}" has invalid or unsupported extension ' 

55 '"{}".'.format(filename, extension) 

56 ) 

57 

58 

59def load_config_file(filename): 

60 handler = get_handler(filename) 

61 with open(filename, "r" + handler.mode) as f: 

62 return handler.load(f) 

63 

64 

65def save_config_file(config, filename): 

66 handler = get_handler(filename) 

67 with open(filename, "w" + handler.mode) as f: 

68 handler.dump(config, f)