Coverage for src/km3pipe/style/__init__.py: 98%

1# Filename: style.py 

2# pylint: disable=locally-disabled 

3""" 

4The KM3Pipe style definitions. 

5 

6""" 

7 

8from os.path import dirname, join, exists 

9from itertools import cycle 

10 

11__author__ = "Tamas Gal" 

12__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration." 

13__credits__ = [] 

14__license__ = "MIT" 

15__maintainer__ = "Tamas Gal" 

16__email__ = "tgal@km3net.de" 

17__status__ = "Development" 

18 

19STYLE_DIR = join(dirname(dirname(__file__)), "stylelib") 

20 

21 

22def get_style_path(style): 

23 return STYLE_DIR + "/" + style + ".mplstyle" 

24 

25 

26def use(style="km3pipe"): 

27 import matplotlib.pyplot as plt 

28 

29 for s in (get_style_path("km3pipe-" + style), get_style_path(style), style): 

30 if exists(s): 

31 plt.style.use(s) 

32 return 

33 print("Could not find style: '{0}'".format(style)) 

34 

35 

36class ColourCycler(object): 

37 """Basic colour cycler. 

38 

39 Instantiate with `cc = ColourCycler()` and use it in plots 

40 like `plt.plot(xs, ys, c=next(cc))`. 

41 """ 

42 

43 def __init__(self, palette="km3pipe"): 

44 self.colours = {} 

45 self.refresh_styles() 

46 self.choose(palette) 

47 

48 def choose(self, palette): 

49 """Pick a palette""" 

50 try: 

51 self._cycler = cycle(self.colours[palette]) 

52 except KeyError: 

53 raise KeyError( 

54 "Chose one of the following colour palettes: {0}".format(self.available) 

55 ) 

56 

57 def refresh_styles(self): 

58 """Load all available styles""" 

59 import matplotlib.pyplot as plt 

60 

61 self.colours = {} 

62 for style in plt.style.available: 

63 try: 

64 style_colours = plt.style.library[style]["axes.prop_cycle"] 

65 self.colours[style] = [c["color"] for c in list(style_colours)] 

66 except KeyError: 

67 continue 

68 

69 self.colours["km3pipe"] = [ 

70 "#ff7869", 

71 "#4babe1", 

72 "#96ad3e", 

73 "#e4823d", 

74 "#5d72b2", 

75 "#e2a3c2", 

76 "#fd9844", 

77 "#e480e7", 

78 ] 

79 

80 @property 

81 def available(self): 

82 """Return a list of available styles""" 

83 return list(self.colours.keys()) 

84 

85 def __next__(self): 

86 """Return the next colour in current palette""" 

87 return next(self._cycler) 

88 

89 def next(self): 

90 """Python 2 compatibility for iterators""" 

91 return self.__next__()