Coverage for /usr/local/lib/python3.7/dist-packages/pyrocko/squirrel/base.py : 92%

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

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

from __future__ import absolute_import, print_function 

import sys 

import os 

import re 

import math 

import threading 

import sqlite3 

import logging 

from collections import defaultdict 

from pyrocko.io_common import FileLoadError 

from pyrocko.guts import Object, Int, List, Tuple, String, Timestamp, Dict 

from pyrocko import util 

from pyrocko.progress import progress 

from . import model, io, cache, dataset 

from .model import to_kind_id, to_kind, separator, WaveformOrder 

from .client import fdsn, catalog 

from . import client, environment, error, pile 

logger = logging.getLogger('pyrocko.squirrel.base') 

guts_prefix = 'pf' 

class GeneratorWithLen(object): 

def __init__(self, gen, length): 

self.gen = gen 

self.length = length 

def __len__(self): 

return self.length 

def __iter__(self): 

return self.gen 

def make_task(*args, **kwargs): 

kwargs['logger'] = logger 

return progress.task(*args, **kwargs) 

def lpick(condition, seq): 

ft = [], [] 

for ele in seq: 

ft[int(bool(condition(ele)))].append(ele) 

return ft 

def execute_get1(connection, sql, args): 

return list(connection.execute(sql, args))[0] 

g_icount = 0 

g_lock = threading.Lock() 

def make_unique_name(): 

with g_lock: 

global g_icount 

name = '%i_%i' % (os.getpid(), g_icount) 

g_icount += 1 

return name 

def codes_fill(n, codes): 

return codes[:n] + ('*',) * (n-len(codes)) 

c_kind_to_ncodes = { 

'station': 4, 

'channel': 5, 

'response': 5, 

'waveform': 6, 

'event': 1, 

'waveform_promise': 6, 

'undefined': 1} 

c_inflated = ['', '*', '*', '*', '*', '*'] 

c_offsets = [0, 2, 1, 1, 1, 1, 0] 

def codes_inflate(codes): 

codes = codes[:6] 

inflated = list(c_inflated) 

ncodes = len(codes) 

offset = c_offsets[ncodes] 

inflated[offset:offset+ncodes] = codes 

return inflated 

def codes_patterns_for_kind(kind, codes): 

if not codes: 

return [] 

if kind in ('event', 'undefined'): 

return [codes] 

cfill = codes_inflate(codes)[:c_kind_to_ncodes[kind]] 

if kind == 'station': 

cfill2 = list(cfill) 

cfill2[3] = '[*]' 

return [cfill, cfill2] 

return [cfill] 

def group_channels(channels): 

groups = defaultdict(list) 

for channel in channels: 

codes = channel.codes 

gcodes = codes[:-1] + (codes[-1][:-1],) 

groups[gcodes].append(channel) 

return groups 

def pyrocko_station_from_channel_group(group, extra_args): 

list_of_args = [channel._get_pyrocko_station_args() for channel in group] 

args = util.consistency_merge(list_of_args + extra_args) 

from pyrocko import model as pmodel 

return pmodel.Station( 

network=args[0], 

station=args[1], 

location=args[2], 

lat=args[3], 

lon=args[4], 

elevation=args[5], 

depth=args[6], 

channels=[ch.get_pyrocko_channel() for ch in group]) 

def blocks(tmin, tmax, deltat, nsamples_block=100000): 

tblock = deltat * nsamples_block 

iblock_min = int(math.floor(tmin / tblock)) 

iblock_max = int(math.ceil(tmax / tblock)) 

for iblock in range(iblock_min, iblock_max): 

yield iblock * tblock, (iblock+1) * tblock 

def gaps(avail, tmin, tmax): 

assert tmin < tmax 

data = [(tmax, 1), (tmin, -1)] 

for (tmin_a, tmax_a) in avail: 

assert tmin_a < tmax_a 

data.append((tmin_a, 1)) 

data.append((tmax_a, -1)) 

data.sort() 

s = 1 

gaps = [] 

tmin_g = None 

for t, x in data: 

if s == 1 and x == -1: 

tmin_g = t 

elif s == 0 and x == 1 and tmin_g is not None: 

tmax_g = t 

if tmin_g != tmax_g: 

gaps.append((tmin_g, tmax_g)) 

s += x 

return gaps 

def order_key(order): 

return (order.codes, order.tmin, order.tmax) 

class Selection(object): 

''' 

Database backed file selection (base class for :py:class:`Squirrel`). 

:param database: 

Database instance or file path to database. 

:type database: 

:py:class:`Database` or :py:class:`str` 

:param persistent: 

If given a name, create a persistent selection. 

:type persistent: 

:py:class:`str` 

In the Squirrel framework, a selection is conceptually a list of files to 

be made available in the application. Instead of using 

:py:class:`Selection` directly, user applications should usually use its 

subclass :py:class:`Squirrel` which adds content indices to the selection 

and provides high level data querying. 

By default, a temporary table in the database is created to hold the names 

of the files in the selection. This table is only visible inside the 

application which created it. If a name is given to ``persistent``, a named 

selection is created, which is visible also in other applications using the 

same database. 

Besides the filename references, desired content kind masks and file format 

indications are stored in the selection's database table to make the user 

choice regarding these options persistent on a per-file basis. Book-keeping 

on whether files are unknown, known or if modification checks are forced is 

also handled in the selection's file-state table. 

Paths of files can be added to the selection using the :py:meth:`add` 

method and removed with :py:meth:`remove`. :py:meth:`undig_grouped` can be 

used to iterate over all content known to the selection. 

''' 

def __init__(self, database, persistent=None): 

self._conn = None 

if not isinstance(database, Database): 

database = get_database(database) 

if persistent is not None: 

assert isinstance(persistent, str) 

if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', persistent): 

raise error.SquirrelError( 

'invalid persistent selection name: %s' % persistent) 

self.name = 'psel_' + persistent 

else: 

self.name = 'sel_' + make_unique_name() 

self._persistent = persistent is not None 

self._database = database 

self._conn = self._database.get_connection() 

self._sources = [] 

self._names = { 

'db': 'main' if self._persistent else 'temp', 

'file_states': self.name + '_file_states', 

'bulkinsert': self.name + '_bulkinsert'} 

self._conn.execute(self._register_table(self._sql( 

''' 

CREATE TABLE IF NOT EXISTS %(db)s.%(file_states)s ( 

file_id integer PRIMARY KEY, 

file_state integer, 

kind_mask integer, 

format text) 

'''))) 

self._conn.execute(self._sql( 

''' 

CREATE INDEX 

IF NOT EXISTS %(db)s.%(file_states)s_index_file_state 

ON %(file_states)s (file_state) 

''')) 

def __del__(self): 

if hasattr(self, '_conn') and self._conn: 

if not self._persistent: 

self._delete() 

else: 

self._conn.commit() 

def _register_table(self, s): 

return self._database._register_table(s) 

def _sql(self, s): 

return s % self._names 

def get_database(self): 

''' 

Get the database to which this selection belongs. 

:returns: :py:class:`Database` object 

''' 

return self._database 

def _delete(self): 

''' 

Destroy the tables assoctiated with this selection. 

''' 

self._conn.execute(self._sql( 

'DROP TABLE %(db)s.%(file_states)s')) 

self._conn.commit() 

def add( 

self, 

paths, 

kind_mask=model.g_kind_mask_all, 

format='detect'): 

''' 

Add files to the selection. 

:param paths: 

Paths to files to be added to the selection. 

:type paths: 

iterator yielding :py:class:`str` objects 

''' 

if isinstance(paths, str): 

paths = [paths] 

task = make_task('Gathering file names') 

paths = task(paths) 

paths = util.short_to_list(200, paths) 

if isinstance(paths, list) and len(paths) <= 200: 

# short non-iterator paths: can do without temp table 

self._conn.executemany( 

''' 

INSERT OR IGNORE INTO files 

VALUES (NULL, ?, NULL, NULL, NULL) 

''', ((x,) for x in paths)) 

task = make_task('Preparing database', 3) 

task.update(0, condition='pruning stale information') 

self._conn.executemany(self._sql( 

''' 

DELETE FROM %(db)s.%(file_states)s 

WHERE file_id IN ( 

SELECT files.file_id 

FROM files 

WHERE files.path == ? ) 

AND kind_mask != ? OR format != ? 

'''), ( 

(path, kind_mask, format) for path in paths)) 

task.update(1, condition='adding file names to selection') 

self._conn.executemany(self._sql( 

''' 

INSERT OR IGNORE INTO %(db)s.%(file_states)s 

SELECT files.file_id, 0, ?, ? 

FROM files 

WHERE files.path = ? 

'''), ((kind_mask, format, path) for path in paths)) 

task.update(2, condition='updating file states') 

self._conn.executemany(self._sql( 

''' 

UPDATE %(db)s.%(file_states)s 

SET file_state = 1 

WHERE file_id IN ( 

SELECT files.file_id 

FROM files 

WHERE files.path == ? ) 

AND file_state != 0 

'''), ((path,) for path in paths)) 

task.update(3) 

task.done() 

else: 

self._conn.execute(self._sql( 

''' 

CREATE TEMP TABLE temp.%(bulkinsert)s 

(path text) 

''')) 

self._conn.executemany(self._sql( 

'INSERT INTO temp.%(bulkinsert)s VALUES (?)'), 

((x,) for x in paths)) 

task = make_task('Preparing database', 5)