1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Manage separate databases and caches for different user projects. 

8 

9Squirrel based applications can either use the user's global database which 

10lives in Pyrocko's global cache directory (by default under 

11``'$HOME/.pyrocko/cache/squirrel'``) or a project specific local database which 

12can be conveniently created in the top level directory of a user's project 

13under ``'.squirrel'`` or ``'squirrel'``. The choice of database and associated 

14cache directory locations is referred to here as the Squirrel environment. This 

15module provides functions to create local environments and to look for a usable 

16environment in the hierarchy of a user's project directory. 

17''' 

18 

19from __future__ import absolute_import, print_function 

20 

21import os 

22 

23from pyrocko.guts import String, load 

24from pyrocko import config 

25from pyrocko.has_paths import HasPaths, Path 

26 

27from . import error 

28 

29guts_prefix = 'squirrel' 

30 

31 

32g_db_filename = 'nuts.sqlite' 

33g_cache_dirname = 'cache' 

34g_config_filename = 'config.yaml' 

35 

36 

37def get_config_path(squirrel_path): 

38 return os.path.join(squirrel_path, g_config_filename) 

39 

40 

41def get_squirrel_path(path=None): 

42 if path is None: 

43 path = os.curdir 

44 

45 path = os.path.abspath(path) 

46 if os.path.isdir(path) and os.path.exists( 

47 os.path.join(path, g_db_filename)): 

48 return path 

49 

50 while True: 

51 for entry in ['squirrel', '.squirrel']: 

52 candidate = os.path.join(path, entry) 

53 if os.path.isdir(candidate) \ 

54 and os.path.exists(os.path.join(candidate, g_db_filename)): 

55 

56 return candidate 

57 

58 path_new = os.path.dirname(path) 

59 if path_new == path: 

60 break 

61 

62 path = path_new 

63 

64 return os.path.join(config.config().cache_dir, 'squirrel') 

65 

66 

67def get_environment(path=None): 

68 ''' 

69 Get default Squirrel environment relevant for a given file system path. 

70 

71 :param path: 

72 Directory path to use as starting point for detection of the Squirrel 

73 environment. By default, the current working directory is used as 

74 starting point. When searching for a usable environment the directory 

75 ``'.squirrel'`` or ``'squirrel'`` in the current (or starting point) 

76 directory is used if it exists, otherwise the parent directories are 

77 search upwards for the existence of such a directory. If no such 

78 directory is found, the user's global Squirrel environment, usually 

79 ``'$HOME/.pyrocko/cache/squirrel'``, is used. 

80 

81 :returns: 

82 :py:class:`Environment` object containing the detected database 

83 and cache directory paths. 

84 ''' 

85 if path is None: 

86 path = os.curdir 

87 

88 squirrel_path = get_squirrel_path(path) 

89 

90 if os.path.exists(get_config_path(squirrel_path)): 

91 return Environment.load(squirrel_path) 

92 else: 

93 return Environment.make(squirrel_path) 

94 

95 

96def init_environment(path=None): 

97 ''' 

98 Initialize empty Squirrel environment. 

99 

100 :param path: 

101 Path to the directory where the new environment's ``'.squirrel'`` 

102 directory should be created. If set to ``None``, the current working 

103 directory is used. 

104 

105 If a ``'.squirrel'`` directory already exists at the given location, 

106 :py:exc:`~pyrocko.squirrel.error.SquirrelError` is raised. 

107 ''' 

108 if path is None: 

109 path = os.curdir 

110 

111 squirrel_path = os.path.join(path, '.squirrel') 

112 try: 

113 os.mkdir(squirrel_path) 

114 except OSError: 

115 raise error.SquirrelError( 

116 'Cannot create squirrel directory: %s' % squirrel_path) 

117 

118 from .base import Squirrel 

119 env = Environment.make(squirrel_path) 

120 env.dump(filename=get_config_path(squirrel_path)) 

121 sq = Squirrel(env) 

122 del sq 

123 

124 

125class Environment(HasPaths): 

126 ''' 

127 Configuration object providing paths to database and cache. 

128 ''' 

129 

130 database_path = Path.T(optional=True) 

131 cache_path = Path.T(optional=True) 

132 persistent = String.T(optional=True) 

133 

134 @classmethod 

135 def make(cls, squirrel_path): 

136 env = cls( 

137 database_path=g_db_filename, 

138 cache_path=g_cache_dirname) 

139 

140 env.set_basepath(squirrel_path) 

141 return env 

142 

143 @classmethod 

144 def load(cls, squirrel_path): 

145 path = get_config_path(squirrel_path) 

146 try: 

147 env = load(filename=path) 

148 except OSError: 

149 raise error.SquirrelError( 

150 'Cannot read environment config file: %s' % path) 

151 

152 if not isinstance(env, Environment): 

153 raise error.SquirrelError( 

154 'Invalid environment config file "%s".' % path) 

155 

156 env.set_basepath(squirrel_path) 

157 return env 

158 

159 

160__all__ = [ 

161 'get_squirrel_path', 

162 'get_environment', 

163 'init_environment', 

164 'Environment']