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 

28from .database import close_database 

29 

30guts_prefix = 'squirrel' 

31 

32 

33g_db_filename = 'nuts.sqlite' 

34g_cache_dirname = 'cache' 

35g_config_filename = 'config.yaml' 

36 

37 

38def get_config_path(squirrel_path): 

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

40 

41 

42def get_squirrel_path(path=None): 

43 if path is None: 

44 path = os.curdir 

45 

46 path = os.path.abspath(path) 

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

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

49 return path 

50 

51 while True: 

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

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

54 if os.path.isdir(candidate) \ 

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

56 

57 return candidate 

58 

59 path_new = os.path.dirname(path) 

60 if path_new == path: 

61 break 

62 

63 path = path_new 

64 

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

66 

67 

68def get_environment(path=None): 

69 ''' 

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

71 

72 :param path: 

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

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

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

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

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

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

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

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

81 

82 :returns: 

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

84 and cache directory paths. 

85 ''' 

86 if path is None: 

87 path = os.curdir 

88 

89 squirrel_path = get_squirrel_path(path) 

90 

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

92 return Environment.load(squirrel_path) 

93 else: 

94 return Environment.make(squirrel_path) 

95 

96 

97def init_environment(path=None): 

98 ''' 

99 Initialize empty Squirrel environment. 

100 

101 :param path: 

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

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

104 directory is used. 

105 

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

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

108 ''' 

109 if path is None: 

110 path = os.curdir 

111 

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

113 try: 

114 os.mkdir(squirrel_path) 

115 except OSError: 

116 raise error.SquirrelError( 

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

118 

119 from .base import Squirrel 

120 env = Environment.make(squirrel_path) 

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

122 sq = Squirrel(env) 

123 database = sq.get_database() 

124 del sq 

125 close_database(database) 

126 

127 

128class Environment(HasPaths): 

129 ''' 

130 Configuration object providing paths to database and cache. 

131 ''' 

132 

133 database_path = Path.T(optional=True) 

134 cache_path = Path.T(optional=True) 

135 persistent = String.T(optional=True) 

136 

137 @classmethod 

138 def make(cls, squirrel_path): 

139 env = cls( 

140 database_path=g_db_filename, 

141 cache_path=g_cache_dirname) 

142 

143 env.set_basepath(squirrel_path) 

144 return env 

145 

146 @classmethod 

147 def load(cls, squirrel_path): 

148 path = get_config_path(squirrel_path) 

149 try: 

150 env = load(filename=path) 

151 except OSError: 

152 raise error.SquirrelError( 

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

154 

155 if not isinstance(env, Environment): 

156 raise error.SquirrelError( 

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

158 

159 env.set_basepath(squirrel_path) 

160 return env 

161 

162 

163__all__ = [ 

164 'get_squirrel_path', 

165 'get_environment', 

166 'init_environment', 

167 'Environment']