Source code for pyrocko.squirrel.environment

# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------

'''
Manage separate databases and caches for different user projects.

Squirrel based applications can either use the user's global database which
lives in Pyrocko's global cache directory (by default under
``'$HOME/.pyrocko/cache/squirrel'``) or a project specific local database which
can be conveniently created in the top level directory of a user's project
under ``'.squirrel'`` or ``'squirrel'``. The choice of database and associated
cache directory locations is referred to here as the Squirrel environment. This
module provides functions to create local environments and to look for a usable
environment in the hierarchy of a user's project directory.
'''

import os
import logging

from pyrocko.guts import String, load
from pyrocko import config
from pyrocko.has_paths import HasPaths, Path

from . import error
from .database import close_database

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

guts_prefix = 'squirrel'


g_db_filename = 'nuts.sqlite'
g_cache_dirname = 'cache'
g_config_filename = 'config.yaml'




def get_config_path(squirrel_path):
    '''
    Get path of environment's configuration file, given its directory.

    :param squirrel_path:
        Path to the Squirrel environment directory (``'.squirrel'`` or the
        global cache directory).
    :type squirrel_path:
        str
    '''
    return os.path.join(squirrel_path, g_config_filename)





def get_squirrel_path(path=None):
    '''
    Detect the Squirrel environment directory relevant for a given path.

    :param path:
        Directory path to use as starting point for the search. By
        default, the current working directory is used. Searches the
        given (or current) directory and its parents upwards for a
        ``'squirrel'`` or ``'.squirrel'`` directory containing an already
        initialized database; falls back to the user's global Squirrel
        cache directory if none is found.
    :type path:
        str

    :returns:
        Detected environment directory path.
    :rtype:
        str
    '''
    if path is None:
        path = os.curdir

    path = os.path.abspath(path)
    if os.path.isdir(path) and os.path.exists(
            os.path.join(path, g_db_filename)):
        return path

    while True:
        for entry in ['squirrel', '.squirrel']:
            candidate = os.path.join(path, entry)
            if os.path.isdir(candidate) \
                    and os.path.exists(os.path.join(candidate, g_db_filename)):

                return candidate

        path_new = os.path.dirname(path)
        if path_new == path:
            break

        path = path_new

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





def get_environment(path=None):
    '''
    Get default Squirrel environment relevant for a given file system path.

    :param path:
        Directory path to use as starting point for detection of the Squirrel
        environment. By default, the current working directory is used as
        starting point. When searching for a usable environment the directory
        ``'.squirrel'`` or ``'squirrel'`` in the current (or starting point)
        directory is used if it exists, otherwise the parent directories are
        search upwards for the existence of such a directory. If no such
        directory is found, the user's global Squirrel environment, usually
        ``'$HOME/.pyrocko/cache/squirrel'``, is used.

    :returns:
        :py:class:`Environment` object containing the detected database
        and cache directory paths.
    '''
    if path is None:
        path = os.curdir

    squirrel_path = get_squirrel_path(path)

    if os.path.exists(get_config_path(squirrel_path)):
        return Environment.load(squirrel_path)
    else:
        return Environment.make(squirrel_path)





def init_environment(path=None):
    '''
    Initialize empty Squirrel environment.

    :param path:
        Path to the directory where the new environment's ``'.squirrel'``
        directory should be created. If set to ``None``, the current working
        directory is used.

    If a ``'.squirrel'`` directory already exists at the given location,
    :py:exc:`~pyrocko.squirrel.error.SquirrelError` is raised.
    '''
    if path is None:
        path = os.curdir

    squirrel_path = os.path.join(path, '.squirrel')
    try:
        logger.info(
            'Creating squirrel environment directory: %s'
            % os.path.abspath(squirrel_path))
        os.mkdir(squirrel_path)
    except OSError:
        raise error.SquirrelError(
            'Cannot create squirrel directory: %s' % squirrel_path)

    from .base import Squirrel
    env = Environment.make(squirrel_path)
    env.dump(filename=get_config_path(squirrel_path))
    sq = Squirrel(env)
    database = sq.get_database()
    del sq
    close_database(database)





class Environment(HasPaths):
    '''
    Configuration object providing paths to database and cache.
    '''

    database_path = Path.T(
        optional=True,
        help='Path to the SQLite database file holding the index.')
    cache_path = Path.T(
        optional=True,
        help='Path to the directory used for cached content, e.g. '
             'downloaded waveforms.')
    persistent = String.T(
        optional=True,
        help='Name of a persistent selection to reuse, or ``None`` to use '
             'a temporary selection (see '
             ':py:class:`~pyrocko.squirrel.selection.Selection`).')



    @classmethod
    def make(cls, squirrel_path):
        '''
        Create a new environment configuration for a given directory.

        :param squirrel_path:
            Path to the Squirrel environment directory. Used as base path
            for :py:attr:`database_path`/:py:attr:`cache_path`.
        :type squirrel_path:
            str

        :rtype:
            :py:class:`Environment`
        '''
        env = cls(
            database_path=g_db_filename,
            cache_path=g_cache_dirname)

        env.set_basepath(squirrel_path)
        return env




    @classmethod
    def load(cls, squirrel_path):
        '''
        Load environment configuration from a directory's config file.

        :param squirrel_path:
            Path to the Squirrel environment directory containing the
            configuration file (see :py:func:`get_config_path`).
        :type squirrel_path:
            str

        :rtype:
            :py:class:`Environment`

        :raises SquirrelError:
            If the config file cannot be read or does not contain a valid
            :py:class:`Environment`.
        '''
        path = get_config_path(squirrel_path)
        try:
            env = load(filename=path)
        except OSError:
            raise error.SquirrelError(
                'Cannot read environment config file: %s' % path)

        if not isinstance(env, Environment):
            raise error.SquirrelError(
                'Invalid environment config file "%s".' % path)

        env.set_basepath(squirrel_path)
        return env




__all__ = [
    'get_squirrel_path',
    'get_environment',
    'init_environment',
    'Environment']