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 

22import logging 

23 

24from pyrocko.guts import String, load 

25from pyrocko import config 

26from pyrocko.has_paths import HasPaths, Path 

27 

28from . import error 

29from .database import close_database 

30 

31logger = logging.getLogger('psq.environment') 

32 

33guts_prefix = 'squirrel' 

34 

35 

36g_db_filename = 'nuts.sqlite' 

37g_cache_dirname = 'cache' 

38g_config_filename = 'config.yaml' 

39 

40 

41def get_config_path(squirrel_path): 

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

43 

44 

45def get_squirrel_path(path=None): 

46 if path is None: 

47 path = os.curdir 

48 

49 path = os.path.abspath(path) 

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

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

52 return path 

53 

54 while True: 

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

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

57 if os.path.isdir(candidate) \ 

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

59 

60 return candidate 

61 

62 path_new = os.path.dirname(path) 

63 if path_new == path: 

64 break 

65 

66 path = path_new 

67 

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

69 

70 

71def get_environment(path=None): 

72 ''' 

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

74 

75 :param path: 

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

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

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

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

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

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

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

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

84 

85 :returns: 

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

87 and cache directory paths. 

88 ''' 

89 if path is None: 

90 path = os.curdir 

91 

92 squirrel_path = get_squirrel_path(path) 

93 

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

95 return Environment.load(squirrel_path) 

96 else: 

97 return Environment.make(squirrel_path) 

98 

99 

100def init_environment(path=None): 

101 ''' 

102 Initialize empty Squirrel environment. 

103 

104 :param path: 

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

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

107 directory is used. 

108 

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

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

111 ''' 

112 if path is None: 

113 path = os.curdir 

114 

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

116 try: 

117 logger.info( 

118 'Creating squirrel environment directory: %s' 

119 % os.path.abspath(squirrel_path)) 

120 os.mkdir(squirrel_path) 

121 except OSError: 

122 raise error.SquirrelError( 

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

124 

125 from .base import Squirrel 

126 env = Environment.make(squirrel_path) 

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

128 sq = Squirrel(env) 

129 database = sq.get_database() 

130 del sq 

131 close_database(database) 

132 

133 

134class Environment(HasPaths): 

135 ''' 

136 Configuration object providing paths to database and cache. 

137 ''' 

138 

139 database_path = Path.T(optional=True) 

140 cache_path = Path.T(optional=True) 

141 persistent = String.T(optional=True) 

142 

143 @classmethod 

144 def make(cls, squirrel_path): 

145 env = cls( 

146 database_path=g_db_filename, 

147 cache_path=g_cache_dirname) 

148 

149 env.set_basepath(squirrel_path) 

150 return env 

151 

152 @classmethod 

153 def load(cls, squirrel_path): 

154 path = get_config_path(squirrel_path) 

155 try: 

156 env = load(filename=path) 

157 except OSError: 

158 raise error.SquirrelError( 

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

160 

161 if not isinstance(env, Environment): 

162 raise error.SquirrelError( 

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

164 

165 env.set_basepath(squirrel_path) 

166 return env 

167 

168 

169__all__ = [ 

170 'get_squirrel_path', 

171 'get_environment', 

172 'init_environment', 

173 'Environment']