Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/environment.py: 93%

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 

19import os 

20import logging 

21 

22from pyrocko.guts import String, load 

23from pyrocko import config 

24from pyrocko.has_paths import HasPaths, Path 

25 

26from . import error 

27from .database import close_database 

28 

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

30 

31guts_prefix = 'squirrel' 

32 

33 

34g_db_filename = 'nuts.sqlite' 

35g_cache_dirname = 'cache' 

36g_config_filename = 'config.yaml' 

37 

38 

39def get_config_path(squirrel_path): 

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

41 

42 

43def get_squirrel_path(path=None): 

44 if path is None: 

45 path = os.curdir 

46 

47 path = os.path.abspath(path) 

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

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

50 return path 

51 

52 while True: 

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

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

55 if os.path.isdir(candidate) \ 

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

57 

58 return candidate 

59 

60 path_new = os.path.dirname(path) 

61 if path_new == path: 

62 break 

63 

64 path = path_new 

65 

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

67 

68 

69def get_environment(path=None): 

70 ''' 

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

72 

73 :param path: 

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

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

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

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

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

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

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

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

82 

83 :returns: 

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

85 and cache directory paths. 

86 ''' 

87 if path is None: 

88 path = os.curdir 

89 

90 squirrel_path = get_squirrel_path(path) 

91 

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

93 return Environment.load(squirrel_path) 

94 else: 

95 return Environment.make(squirrel_path) 

96 

97 

98def init_environment(path=None): 

99 ''' 

100 Initialize empty Squirrel environment. 

101 

102 :param path: 

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

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

105 directory is used. 

106 

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

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

109 ''' 

110 if path is None: 

111 path = os.curdir 

112 

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

114 try: 

115 logger.info( 

116 'Creating squirrel environment directory: %s' 

117 % os.path.abspath(squirrel_path)) 

118 os.mkdir(squirrel_path) 

119 except OSError: 

120 raise error.SquirrelError( 

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

122 

123 from .base import Squirrel 

124 env = Environment.make(squirrel_path) 

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

126 sq = Squirrel(env) 

127 database = sq.get_database() 

128 del sq 

129 close_database(database) 

130 

131 

132class Environment(HasPaths): 

133 ''' 

134 Configuration object providing paths to database and cache. 

135 ''' 

136 

137 database_path = Path.T(optional=True) 

138 cache_path = Path.T(optional=True) 

139 persistent = String.T(optional=True) 

140 

141 @classmethod 

142 def make(cls, squirrel_path): 

143 env = cls( 

144 database_path=g_db_filename, 

145 cache_path=g_cache_dirname) 

146 

147 env.set_basepath(squirrel_path) 

148 return env 

149 

150 @classmethod 

151 def load(cls, squirrel_path): 

152 path = get_config_path(squirrel_path) 

153 try: 

154 env = load(filename=path) 

155 except OSError: 

156 raise error.SquirrelError( 

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

158 

159 if not isinstance(env, Environment): 

160 raise error.SquirrelError( 

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

162 

163 env.set_basepath(squirrel_path) 

164 return env 

165 

166 

167__all__ = [ 

168 'get_squirrel_path', 

169 'get_environment', 

170 'init_environment', 

171 'Environment']