<commit_before>from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel
import itertools
from itertools import chain

import threading
import torch
from torch.cuda._utils import _get_device_index
import pdb

def _find_tensors(obj):
    r"""
    Recursively find all tensors contained in the specified object.
    """
    if isinstance(obj, torch.Tensor):
        return [obj]
    if isinstance(obj, (list, tuple)):
        return itertools.chain(*map(_find_tensors, obj))
    if isinstance(obj, dict):
        return itertools.chain(*map(_find_tensors, obj.values()))
    return []

def get_a_var(obj):
    if isinstance(obj, torch.Tensor):
        return obj

    if isinstance(obj, list) or isinstance(obj, tuple):
        for result in map(get_a_var, obj):
            if isinstance(result, torch.Tensor):
                return result
    if isinstance(obj, dict):
        for result in map(get_a_var, obj.items()):
            if isinstance(result, torch.Tensor):
                return result
    return None

class LightningDataParallel(DataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)

        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            # lightning
            if self.module.training:
                return self.module.training_step(*inputs[0], **kwargs[0])
            else:
                return self.module.validation_step(*inputs[0], **kwargs[0])

        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)

    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

class LightningDistributedDataParallel(DistributedDataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

    def forward(self, *inputs, **kwargs):
        self._sync_params()
        if self.device_ids:
            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
            if len(self.device_ids) == 1:
                # --------------
                # LIGHTNING MOD
                # --------------
                # normal
                # output = self.module(*inputs[0], **kwargs[0])

                # lightning
                if self.module.training:
                    output = self.module.training_step(*inputs[0], **kwargs[0])
                else:
                    output = self.module.validation_step(*inputs[0], **kwargs[0])
            else:
                outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
                output = self.gather(outputs, self.output_device)
        else:
            output = self.module(*inputs, **kwargs)

        if torch.is_grad_enabled():
            # We'll return the output object verbatim since it is a freeform
            # object. We need to find any tensors in this object, though,
            # because we need to figure out which parameters were used during
            # this forward pass, to ensure we short circuit reduction for any
            # unused parameters. Only if `find_unused_parameters` is set.
            if self.find_unused_parameters:
                self.reducer.prepare_for_backward(list(_find_tensors(output)))
            else:
                self.reducer.prepare_for_backward([])
        return output

def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
    r"""Applies each `module` in :attr:`modules` in parallel on arguments
    contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
    on each of :attr:`devices`.

    Args:
        modules (Module): modules to be parallelized
        inputs (tensor): inputs to the modules
        devices (list of int or torch.device): CUDA devices

    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
    :attr:`devices` (if given) should all have same length. Moreover, each
    element of :attr:`inputs` can either be a single object as the only argument
    to a module, or a collection of positional arguments.
    """
    assert len(modules) == len(inputs)
    if kwargs_tup is not None:
        assert len(modules) == len(kwargs_tup)
    else:
        kwargs_tup = ({},) * len(modules)
    if devices is not None:
        assert len(modules) == len(devices)
    else:
        devices = [None] * len(modules)
    devices = list(map(lambda x: _get_device_index(x, True), devices))
    lock = threading.Lock()
    results = {}
    grad_enabled = torch.is_grad_enabled()

    def _worker(i, module, input, kwargs, device=None):
        torch.set_grad_enabled(grad_enabled)
        if device is None:
            device = get_a_var(input).get_device()
        try:
            with torch.cuda.device(device):
                # this also avoids accidental slicing of `input` if it is a Tensor
                if not isinstance(input, (list, tuple)):
                    input = (input,)

                # ---------------
                # CHANGE
                pdb.set_trace()
                if module.training:
                    output = module.training_step(*input, **kwargs)
                else:
                    output = module.validation_step(*input, **kwargs)
                # ---------------

            with lock:
                results[i] = output
        except Exception as e:
            with lock:
                results[i] = e

    if len(modules) > 1:
        threads = [threading.Thread(target=_worker,
                                    args=(i, module, input, kwargs, device))
                   for i, (module, input, kwargs, device) in
                   enumerate(zip(modules, inputs, kwargs_tup, devices))]

        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
    else:
        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])

    outputs = []
    for i in range(len(inputs)):
        output = results[i]
        if isinstance(output, Exception):
            raise output
        outputs.append(output)
    return outputs
<commit_msg>set dp as default backend<commit_after>from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel
import itertools
from itertools import chain

import threading
import torch
from torch.cuda._utils import _get_device_index
import pdb

def _find_tensors(obj):
    r"""
    Recursively find all tensors contained in the specified object.
    """
    if isinstance(obj, torch.Tensor):
        return [obj]
    if isinstance(obj, (list, tuple)):
        return itertools.chain(*map(_find_tensors, obj))
    if isinstance(obj, dict):
        return itertools.chain(*map(_find_tensors, obj.values()))
    return []

def get_a_var(obj):
    if isinstance(obj, torch.Tensor):
        return obj

    if isinstance(obj, list) or isinstance(obj, tuple):
        for result in map(get_a_var, obj):
            if isinstance(result, torch.Tensor):
                return result
    if isinstance(obj, dict):
        for result in map(get_a_var, obj.items()):
            if isinstance(result, torch.Tensor):
                return result
    return None

class LightningDataParallel(DataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)

        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            # lightning
            if self.module.training:
                return self.module.training_step(*inputs[0], **kwargs[0])
            else:
                return self.module.validation_step(*inputs[0], **kwargs[0])

        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)

    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

class LightningDistributedDataParallel(DistributedDataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

    def forward(self, *inputs, **kwargs):
        self._sync_params()
        if self.device_ids:
            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
            if len(self.device_ids) == 1:
                # --------------
                # LIGHTNING MOD
                # --------------
                # normal
                # output = self.module(*inputs[0], **kwargs[0])

                # lightning
                if self.module.training:
                    output = self.module.training_step(*inputs[0], **kwargs[0])
                else:
                    output = self.module.validation_step(*inputs[0], **kwargs[0])
            else:
                outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
                output = self.gather(outputs, self.output_device)
        else:
            output = self.module(*inputs, **kwargs)

        if torch.is_grad_enabled():
            # We'll return the output object verbatim since it is a freeform
            # object. We need to find any tensors in this object, though,
            # because we need to figure out which parameters were used during
            # this forward pass, to ensure we short circuit reduction for any
            # unused parameters. Only if `find_unused_parameters` is set.
            if self.find_unused_parameters:
                self.reducer.prepare_for_backward(list(_find_tensors(output)))
            else:
                self.reducer.prepare_for_backward([])
        return output

def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
    r"""Applies each `module` in :attr:`modules` in parallel on arguments
    contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
    on each of :attr:`devices`.

    Args:
        modules (Module): modules to be parallelized
        inputs (tensor): inputs to the modules
        devices (list of int or torch.device): CUDA devices

    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
    :attr:`devices` (if given) should all have same length. Moreover, each
    element of :attr:`inputs` can either be a single object as the only argument
    to a module, or a collection of positional arguments.
    """
    assert len(modules) == len(inputs)
    if kwargs_tup is not None:
        assert len(modules) == len(kwargs_tup)
    else:
        kwargs_tup = ({},) * len(modules)
    if devices is not None:
        assert len(modules) == len(devices)
    else:
        devices = [None] * len(modules)
    devices = list(map(lambda x: _get_device_index(x, True), devices))
    lock = threading.Lock()
    results = {}
    grad_enabled = torch.is_grad_enabled()

    def _worker(i, module, input, kwargs, device=None):
        torch.set_grad_enabled(grad_enabled)
        if device is None:
            device = get_a_var(input).get_device()
        try:
            with torch.cuda.device(device):
                # this also avoids accidental slicing of `input` if it is a Tensor
                if not isinstance(input, (list, tuple)):
                    input = (input,)

                # ---------------
                # CHANGE
                if module.training:
                    output = module.training_step(*input, **kwargs)
                else:
                    output = module.validation_step(*input, **kwargs)
                # ---------------

            with lock:
                results[i] = output
        except Exception as e:
            with lock:
                results[i] = e

    if len(modules) > 1:
        threads = [threading.Thread(target=_worker,
                                    args=(i, module, input, kwargs, device))
                   for i, (module, input, kwargs, device) in
                   enumerate(zip(modules, inputs, kwargs_tup, devices))]

        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
    else:
        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])

    outputs = []
    for i in range(len(inputs)):
        output = results[i]
        if isinstance(output, Exception):
            raise output
        outputs.append(output)
    return outputs
<|endoftext|>


<commit_before>"""
Misc database methods
"""

import os
import logging
from datetime import datetime,timedelta
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math
import uuid
import shutil
from io import BytesIO

# load IntegrityError
try:
    import MySQLdb
except ImportError:
    try:
        import pymysql as MySQLdb
    except:
        class MySQLdb:
            class IntegrityError(Exception):
                pass

import iceprod.core.functions
from iceprod.core.util import Node_Resources
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server.dbmethods import dbmethod,_Methods_Base,filtered_input,datetime2str,str2datetime

logger = logging.getLogger('dbmethods.misc')

class misc(_Methods_Base):
    """
    misc DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    ### site to site ###

    @dbmethod
    def misc_site_to_site_upload(self,src,dest,callback=None):
        if callback:
            callback()
        # TODO: actually write this method

    @dbmethod
    def misc_get_tables_for_task(self,task_ids,callback=None):
        """
        Get all tables necessary to run task(s).

        :param task_ids: Either a single, or an iterable of task_ids.
        :returns: (via callback) dict of table entries.
        """
        if isinstance(task_ids,str):
            task_ids = [task_ids]
        elif not isinstance(task_ids,Iterable):
            callback(Exception('task_ids not Iterable'))
        else:
            cb = partial(self._misc_get_tables_for_task_blocking,task_ids,
                         callback=callback)
            self.db.non_blocking_task(cb)
    def _misc_get_tables_for_task_blocking(self,task_ids,callback=None):
        conn,archive_conn = self.db._dbsetup()
        tables = {}

        sql = 'select depends from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if not ret:
            callback({})
            return
        task_ids = set(task_ids)
        for row in ret:
            for d in row[0].split(','):
                d = d.strip()
                if d:
                    task_ids.add(d)

        sql = 'select * from search where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if not ret:
            callback({})
            return

        search_table = {}
        keys = []
        for row in ret:
            search_table[row[0]] = self._list_to_dict('search',row)
            if not keys:
                keys = search_table[row[0]].keys()
        tables['search'] = {'keys':keys,'values':ret}

        sql = 'select * from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        task_rel_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('task',row)
                task_rel_ids.add(row2['task_rel_id'])
                if not keys:
                    keys = row2.keys()
            tables['task'] = {'keys':keys,'values':ret}

        job_ids = set(search_table[id]['job_id'] for id in task_ids)
        sql = 'select * from job where job_id in ('
        sql += ','.join('?' for _ in job_ids)
        sql += ')'
        bindings = tuple(job_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('job',row).keys()
                    break
            tables['job'] = {'keys':keys,'values':ret}

        dataset_ids = set(search_table[id]['dataset_id'] for id in task_ids)
        sql = 'select * from dataset where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        categoryvalue_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('dataset',row)
                if row2['categoryvalue_ids']:
                    for cv_id in row2['categoryvalue_ids'].split(','):
                        categoryvalue_ids.add(cv_id)
                if not keys:
                    keys = row2.keys()
            tables['dataset'] = {'keys':keys,'values':ret}

        sql = 'select * from task_rel where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('task_rel',row).keys()
                    break
            tables['task_rel'] = {'keys':keys,'values':ret}

        sql = 'select * from config where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('config',row).keys()
                    break
            tables['config'] = {'keys':keys,'values':ret}

        categorydef_ids = set()
        if categoryvalue_ids:
            sql = 'select * from categoryvalue where categoryvalue_id in ('
            sql += ','.join('?' for _ in categoryvalue_ids)
            sql += ')'
            bindings = tuple(categoryvalue_ids)
            try:
                ret = self.db._db_read(conn,sql,bindings,None,None,None)
            except Exception as e:
                ret = e
            if isinstance(ret,Exception):
                callback(ret)
                return
            if ret:
                keys = []
                for row in ret:
                    row2 = self._list_to_dict('categoryvalue',row)
                    categorydef_ids.add(row2['categorydef_id'])
                    if not keys:
                        keys = row2.keys()
                tables['categoryvalue'] = {'keys':keys,'values':ret}

        if categorydef_ids:
            sql = 'select * from categorydef where categorydef_id in ('
            sql += ','.join('?' for _ in categorydef_ids)
            sql += ')'
            bindings = tuple(categorydef_ids)
            try:
                ret = self.db._db_read(conn,sql,bindings,None,None,None)
            except Exception as e:
                ret = e
            if isinstance(ret,Exception):
                callback(ret)
                return
            if ret:
                keys = []
                for row in ret:
                    if not keys:
                        keys = self._list_to_dict('categorydef',row).keys()
                        break
                tables['categorydef'] = {'keys':keys,'values':ret}

        callback(tables)

    @dbmethod
    def misc_update_tables(self, tables, callback=None):
        """
        Update the DB tables with the incoming information.

        :param tables: A dict of {table_name:{keys:[],values:[[]]}}
        :returns: (via callback) success or failure
        """
        if not tables:
            callback()
            return
        if not isinstance(tables,dict):
            callback(Exception('tables not a dict'))
        else:
            cb = partial(self._misc_update_tables_blocking,tables,
                         callback=callback)
            self.db.non_blocking_task(cb)
    def _misc_update_tables_blocking(self, tables, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            for name in tables:
                sql = 'replace into %s ('%filtered_input(name)
                sql += ','.join(str(filtered_input(k)) for k in tables[name]['keys'])
                sql += ') values ('
                sql += ','.join('?' for _ in tables[name]['keys'])
                sql += ')'
                for values in tables[name]['values']:
                    bindings = tuple(values)
                    try:
                        ret = self.db._db_write(conn,sql,bindings,None,None,None)
                    except Exception as e:
                        ret = e
                    if isinstance(ret,Exception):
                        callback(ret)
                        return
        except Exception as e:
            logger.warn('error updating tables', exc_info=True)
            callback(e)
        callback()

    @dbmethod
    def misc_update_master_db(self, table, index, timestamp, sql,
                       bindings, callback=None):
        """
        Update the DB with incoming information (query provided).

        :param table: The table affected.
        :param index: That table's index id.
        :param timestamp: An ISO 8601 UTC timestamp.
        :param sql: An sql statement.
        :param bindings: Bindings for the sql statement.
        :returns: (via callback) success or failure
        """
        cb = partial(self._misc_update_db_blocking, table, index,
                     timestamp, sql, bindings, callback=callback)
        self.db.blocking_task('update',cb)
    def _misc_update_db_blocking(self, table, index, timestamp, sql,
                                      bindings, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            sql2 = 'select timestamp from master_update_history '
            sql2 += 'where table_name = ? and update_index = ?'
            bindings2 = (table,index)
            ret = self.db._db_read(conn,sql2,bindings2,None,None,None)
            if ret and ret[0][0] >= timestamp:
                logger.info('newer data already present for %s %s %s',
                            table, index, timestamp)
                callback(None)
                return
            ret = self.db._db_write(conn,sql,tuple(bindings),None,None,None)
        except MySQLdb.IntegrityError:
            logger.warn('dropping history for %r', sql, exc_info=True)
        except Exception as e:
            logger.warn('error finding old timestamp', exc_info=True)
            ret = e
        else:
            sql2 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
            bindings2 = (table,index,timestamp)
            try:
                ret = self.db._db_write(conn,sql2,bindings2,None,None,None)
            except Exception as e:
                logger.warn('error updating history', exc_info=True)
                ret = e
        if callback:
            callback(ret)
<commit_msg>try fixing the master update error<commit_after>"""
Misc database methods
"""

import os
import logging
from datetime import datetime,timedelta
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math
import uuid
import shutil
from io import BytesIO

# load IntegrityError
try:
    import MySQLdb
except ImportError:
    try:
        import pymysql as MySQLdb
    except:
        class MySQLdb:
            class IntegrityError(Exception):
                pass

import iceprod.core.functions
from iceprod.core.util import Node_Resources
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server.dbmethods import dbmethod,_Methods_Base,filtered_input,datetime2str,str2datetime

logger = logging.getLogger('dbmethods.misc')

class misc(_Methods_Base):
    """
    misc DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    ### site to site ###

    @dbmethod
    def misc_site_to_site_upload(self,src,dest,callback=None):
        if callback:
            callback()
        # TODO: actually write this method

    @dbmethod
    def misc_get_tables_for_task(self,task_ids,callback=None):
        """
        Get all tables necessary to run task(s).

        :param task_ids: Either a single, or an iterable of task_ids.
        :returns: (via callback) dict of table entries.
        """
        if isinstance(task_ids,str):
            task_ids = [task_ids]
        elif not isinstance(task_ids,Iterable):
            callback(Exception('task_ids not Iterable'))
        else:
            cb = partial(self._misc_get_tables_for_task_blocking,task_ids,
                         callback=callback)
            self.db.non_blocking_task(cb)
    def _misc_get_tables_for_task_blocking(self,task_ids,callback=None):
        conn,archive_conn = self.db._dbsetup()
        tables = {}

        sql = 'select depends from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if not ret:
            callback({})
            return
        task_ids = set(task_ids)
        for row in ret:
            for d in row[0].split(','):
                d = d.strip()
                if d:
                    task_ids.add(d)

        sql = 'select * from search where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if not ret:
            callback({})
            return

        search_table = {}
        keys = []
        for row in ret:
            search_table[row[0]] = self._list_to_dict('search',row)
            if not keys:
                keys = search_table[row[0]].keys()
        tables['search'] = {'keys':keys,'values':ret}

        sql = 'select * from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        task_rel_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('task',row)
                task_rel_ids.add(row2['task_rel_id'])
                if not keys:
                    keys = row2.keys()
            tables['task'] = {'keys':keys,'values':ret}

        job_ids = set(search_table[id]['job_id'] for id in task_ids)
        sql = 'select * from job where job_id in ('
        sql += ','.join('?' for _ in job_ids)
        sql += ')'
        bindings = tuple(job_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('job',row).keys()
                    break
            tables['job'] = {'keys':keys,'values':ret}

        dataset_ids = set(search_table[id]['dataset_id'] for id in task_ids)
        sql = 'select * from dataset where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        categoryvalue_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('dataset',row)
                if row2['categoryvalue_ids']:
                    for cv_id in row2['categoryvalue_ids'].split(','):
                        categoryvalue_ids.add(cv_id)
                if not keys:
                    keys = row2.keys()
            tables['dataset'] = {'keys':keys,'values':ret}

        sql = 'select * from task_rel where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('task_rel',row).keys()
                    break
            tables['task_rel'] = {'keys':keys,'values':ret}

        sql = 'select * from config where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        try:
            ret = self.db._db_read(conn,sql,bindings,None,None,None)
        except Exception as e:
            ret = e
        if isinstance(ret,Exception):
            callback(ret)
            return
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('config',row).keys()
                    break
            tables['config'] = {'keys':keys,'values':ret}

        categorydef_ids = set()
        if categoryvalue_ids:
            sql = 'select * from categoryvalue where categoryvalue_id in ('
            sql += ','.join('?' for _ in categoryvalue_ids)
            sql += ')'
            bindings = tuple(categoryvalue_ids)
            try:
                ret = self.db._db_read(conn,sql,bindings,None,None,None)
            except Exception as e:
                ret = e
            if isinstance(ret,Exception):
                callback(ret)
                return
            if ret:
                keys = []
                for row in ret:
                    row2 = self._list_to_dict('categoryvalue',row)
                    categorydef_ids.add(row2['categorydef_id'])
                    if not keys:
                        keys = row2.keys()
                tables['categoryvalue'] = {'keys':keys,'values':ret}

        if categorydef_ids:
            sql = 'select * from categorydef where categorydef_id in ('
            sql += ','.join('?' for _ in categorydef_ids)
            sql += ')'
            bindings = tuple(categorydef_ids)
            try:
                ret = self.db._db_read(conn,sql,bindings,None,None,None)
            except Exception as e:
                ret = e
            if isinstance(ret,Exception):
                callback(ret)
                return
            if ret:
                keys = []
                for row in ret:
                    if not keys:
                        keys = self._list_to_dict('categorydef',row).keys()
                        break
                tables['categorydef'] = {'keys':keys,'values':ret}

        callback(tables)

    @dbmethod
    def misc_update_tables(self, tables, callback=None):
        """
        Update the DB tables with the incoming information.

        :param tables: A dict of {table_name:{keys:[],values:[[]]}}
        :returns: (via callback) success or failure
        """
        if not tables:
            callback()
            return
        if not isinstance(tables,dict):
            callback(Exception('tables not a dict'))
        else:
            cb = partial(self._misc_update_tables_blocking,tables,
                         callback=callback)
            self.db.non_blocking_task(cb)
    def _misc_update_tables_blocking(self, tables, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            for name in tables:
                sql = 'replace into %s ('%filtered_input(name)
                sql += ','.join(str(filtered_input(k)) for k in tables[name]['keys'])
                sql += ') values ('
                sql += ','.join('?' for _ in tables[name]['keys'])
                sql += ')'
                for values in tables[name]['values']:
                    bindings = tuple(values)
                    try:
                        ret = self.db._db_write(conn,sql,bindings,None,None,None)
                    except Exception as e:
                        ret = e
                    if isinstance(ret,Exception):
                        callback(ret)
                        return
        except Exception as e:
            logger.warn('error updating tables', exc_info=True)
            callback(e)
        callback()

    @dbmethod
    def misc_update_master_db(self, table, index, timestamp, sql,
                       bindings, callback=None):
        """
        Update the DB with incoming information (query provided).

        :param table: The table affected.
        :param index: That table's index id.
        :param timestamp: An ISO 8601 UTC timestamp.
        :param sql: An sql statement.
        :param bindings: Bindings for the sql statement.
        :returns: (via callback) success or failure
        """
        cb = partial(self._misc_update_db_blocking, table, index,
                     timestamp, sql, bindings, callback=callback)
        self.db.blocking_task('update',cb)
    def _misc_update_db_blocking(self, table, index, timestamp, sql,
                                      bindings, callback=None):
        conn,archive_conn = self.db._dbsetup()
        try:
            sql2 = 'select timestamp from master_update_history '
            sql2 += 'where table_name = ? and update_index = ?'
            bindings2 = (table,index)
            ret = self.db._db_read(conn,sql2,bindings2,None,None,None)
            prev_timestamp = None
            for row in ret:
                prev_timestamp = row[0]
            if prev_timestamp and prev_timestamp >= timestamp:
                logger.info('newer data already present for %s %s %s',
                            table, index, timestamp)
                callback(None)
                return
            ret = self.db._db_write(conn,sql,tuple(bindings),None,None,None)
        except MySQLdb.IntegrityError:
            logger.warn('dropping history for %r', sql, exc_info=True)
        except Exception as e:
            logger.warn('error updating master', exc_info=True)
            ret = e
        else:
            sql2 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
            bindings2 = (table,index,timestamp)
            try:
                ret = self.db._db_write(conn,sql2,bindings2,None,None,None)
            except Exception as e:
                logger.warn('error updating update_history', exc_info=True)
                ret = e
        if callback:
            callback(ret)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#!/usr/bin/env python

# Example using a character LCD plate.
import time, os, json
import Adafruit_CharLCD as LCD
from char_ui import *

# Initialize the LCD
lcd = LCD.Adafruit_CharLCDPlate()

# activate backlight and clear display and init everything
lcd.set_color(1.0, 1.0, 1.0)
lcd.clear()
lcd.message("Please wait...")
time.sleep(1)

# configuration file
conf_file = "conf.json"
conf = json.load(open(conf_file))

# check if frame/video save folder exists
if not os.path.exists(conf["directory"]):
    os.makedirs(conf["directory"])

# define special chars
for char in special_chars:
    lcd.create_char(i[0], i[1])

# start program at main_menu
print("[INFO] Press Ctrl-C to quit.")
submenu = main_menu
try:
    while hasattr(submenu, '__call__'):
        submenu = submenu()
except KeyboardInterrupt:
    print("[INFO] KeyboardInterrupt, exiting...")
<commit_msg>Minor update.<commit_after>#!/usr/bin/env python
#!/usr/bin/env python

# Example using a character LCD plate.
import time, os, json
import Adafruit_CharLCD as LCD
from char_ui import *

# Initialize the LCD
lcd = LCD.Adafruit_CharLCDPlate()

# activate backlight and clear display and init everything
lcd.set_color(1.0, 1.0, 1.0)
lcd.clear()
lcd.message("Please wait...")
time.sleep(1)

# configuration file
conf_file = "conf.json"
conf = json.load(open(conf_file))

# check if frame/video save folder exists
if not os.path.exists(conf["directory"]):
    os.makedirs(conf["directory"])

# define special chars
for char in special_chars:
    lcd.create_char(char[0], char[1])

# start program at main_menu
print("[INFO] Press Ctrl-C to quit.")
submenu = main_menu
try:
    while hasattr(submenu, '__call__'):
        submenu = submenu(lcd)
except KeyboardInterrupt:
    print("[INFO] KeyboardInterrupt, exiting...")
<|endoftext|>


<commit_before># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyExecnet(PythonPackage):
    """execnet provides a share-nothing model with channel-send/receive
    communication for distributing execution across many Python interpreters
    across version, platform and network barriers."""

    homepage = "http://codespeak.net/execnet"
    url      = "https://pypi.io/packages/source/e/execnet/execnet-1.4.1.tar.gz"

    version('1.4.1', sha256='f66dd4a7519725a1b7e14ad9ae7d3df8e09b2da88062386e08e941cafc0ef3e6')

    depends_on('py-setuptools',  type='build')
    depends_on('py-setuptools-scm',  type='build')
    depends_on('py-apipkg@1.4:', type=('build', 'run'))
<commit_msg>Add latest version of py-execnet (#13317)<commit_after># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyExecnet(PythonPackage):
    """execnet provides a share-nothing model with channel-send/receive
    communication for distributing execution across many Python interpreters
    across version, platform and network barriers."""

    homepage = "http://codespeak.net/execnet"
    url      = "https://pypi.io/packages/source/e/execnet/execnet-1.7.1.tar.gz"

    version('1.7.1', sha256='cacb9df31c9680ec5f95553976c4da484d407e85e41c83cb812aa014f0eddc50')
    version('1.4.1', sha256='f66dd4a7519725a1b7e14ad9ae7d3df8e09b2da88062386e08e941cafc0ef3e6')

    depends_on('python@2.7:2.8,3.4:', type=('build', 'run'))
    depends_on('py-setuptools',  type='build')
    depends_on('py-setuptools-scm',  type='build')
    depends_on('py-apipkg@1.4:', type=('build', 'run'))
<|endoftext|>


<commit_before># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyPyfaidx(PythonPackage):
    """pyfaidx: efficient pythonic random access to fasta subsequences"""

    homepage = "https://pypi.org/project/pyfaidx/"
    url      = "https://pypi.io/packages/source/p/pyfaidx/pyfaidx-0.5.5.2.tar.gz"

    version('0.5.5.2', sha256='9ac22bdc7b9c5d995d32eb9dc278af9ba970481636ec75c0d687d38c26446caa')

    depends_on('py-setuptools', type='build')
<commit_msg>py-pyfaidx: tweak dependencies; fixes #12445 (#12446)<commit_after># Copyright 2013-2019 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyPyfaidx(PythonPackage):
    """pyfaidx: efficient pythonic random access to fasta subsequences"""

    homepage = "https://pypi.org/project/pyfaidx/"
    url      = "https://pypi.io/packages/source/p/pyfaidx/pyfaidx-0.5.5.2.tar.gz"

    version('0.5.5.2', sha256='9ac22bdc7b9c5d995d32eb9dc278af9ba970481636ec75c0d687d38c26446caa')

    depends_on('py-setuptools@0.7:', type=('build', 'run'))
    depends_on('py-six', type=('build', 'run'))
    depends_on('py-ordereddict', type=('build', 'run'), when='^python@:2.6')
    depends_on('py-argparse', type=('build', 'run'), when='^python@:2.6')
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys, os
import argparse
import logging
import json
from deepharvest.deepharvest_nuxeo import DeepHarvestNuxeo
import s3stash.s3tools
from s3stash.nxstashref_image import NuxeoStashImage
from s3stash.nxstashref_file import NuxeoStashFile
from s3stash.nxstash_thumb import NuxeoStashThumb
from s3stash.nxstash_mediajson import NuxeoStashMediaJson

IMAGE_BUCKET = 'ucldc-private-files/jp2000'
IMAGE_REGION = 'us-west-2'
FILE_BUCKET = 'ucldc-nuxeo-ref-media'
FILE_REGION = 'us-west-2'
THUMB_BUCKET = 'static.ucldc.cdlib.org/ucldc-nuxeo-thumb-media'
THUMB_REGION = 'us-east-1'
MEDIAJSON_BUCKET = 'static.ucldc.cdlib.org/media_json'
MEDIAJSON_REGION = 'us-east-1'

_loglevel_ = 'INFO'

class Stash(object):
    '''
        stash various files on s3 for a Nuxeo collection
        in preparation for harvesting into Calisphere
    '''
    def __init__(self, path, pynuxrc, replace=False, loglevel=_loglevel_):
        self.logger = logging.getLogger(__name__)

        self.path = path
        self.pynuxrc = pynuxrc
        self.replace = replace

        self.dh = DeepHarvestNuxeo(self.path, '', pynuxrc=self.pynuxrc)

        self.objects = self.dh.fetch_objects()

    def images(self):
        ''' stash Nuxeo image files on s3 '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashImage(obj['path'], IMAGE_BUCKET, IMAGE_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashImage(c['path'], IMAGE_BUCKET, IMAGE_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def files(self):
        ''' stash Nuxeo files of type 'file', 'audio', or 'video' for a collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashFile(obj['path'], FILE_BUCKET, FILE_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashFile(c['path'], FILE_BUCKET, FILE_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def thumbnails(self):
        ''' stash thumbnail images for Nuxeo files of type 'file', 'audio', or 'video' for a collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashThumb(obj['path'], THUMB_BUCKET, THUMB_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashThumb(c['path'], THUMB_BUCKET, THUMB_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def media_json(self):
        ''' create and stash media.json files for a nuxeo collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashMediaJson(obj['path'], MEDIAJSON_BUCKET, MEDIAJSON_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()

        return report

def main(argv=None):
    parser = argparse.ArgumentParser(description='Deep harvest a Nuxeo collection')
    parser.add_argument('registry_id', help='UCLDC Registry ID')
    parser.add_argument('--pynuxrc', default='~/.pynuxrc', help='rc file for use by pynux')
    parser.add_argument('--replace', action='store_true', help='replace files on s3 if they already exist')
    parser.add_argument('--loglevel', default=_loglevel_)

    if argv is None:
        argv = parser.parse_args()

    registry_id = argv.registry_id
    pynuxrc = argv.pynuxrc
    replace = argv.replace
    loglevel = argv.loglevel

    # set up logging
    logfile = 'logs/stash_collection_{}'.format(registry_id)
    numeric_level = getattr(logging, loglevel, None)
    if not isinstance(numeric_level, int):
        raise ValueError('Invalid log level: %s' % loglevel)
    logging.basicConfig(level=numeric_level, filename=logfile, format='%(asctime)s (%(name)s) [%(levelname)s]: %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')
    logger = logging.getLogger(__name__)
    print "\nlogfile: {}\n".format(logfile)

    # get nuxeo path
    nxpath = s3stash.s3tools.get_nuxeo_path(registry_id)
    if nxpath is None:
        print "No record found for registry_id: {}".format(registry_id)
        sys.exit()
    info = 'nuxeo_path: {}'.format(nxpath)
    logger.info(info)
    print info, '\n'

    stash = Stash(nxpath, pynuxrc, replace)

    # stash images for use with iiif server
    print 'stashing images...'
    image_report = stash.images()
    info = 'finished stashing images'
    logger.info(info)
    print info
    reportfile = "reports/images-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(image_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash text, audio, video
    print 'stashing non-image files (text, audio, video)...'
    file_report = stash.files()
    info = 'finished stashing files'
    logger.info(info)
    print info
    reportfile = "reports/files-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(file_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash thumbnails for text, audio, video
    print 'stashing thumbnails for non-image files (text, audio, video)...'
    thumb_report = stash.thumbnails()
    info = 'finished stashing thumbnails'
    logger.info(info)
    print info
    reportfile = "reports/thumbs-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(thumb_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash media.json files
    print 'stashing media.json files for collection...'
    mediajson_report = stash.media_json()
    info = 'finished stashing media.json'
    logger.info(info)
    print info
    reportfile = "reports/mediajson-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(mediajson_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # print some information about how it went
    images_stashed = len([key for key, value in image_report.iteritems() if value['stashed']])
    files_stashed = len([key for key, value in file_report.iteritems() if value['stashed']])
    thumbs_stashed = len([key for key, value in thumb_report.iteritems() if value['stashed']])
    mediajson_stashed = len([key for key, value in mediajson_report.iteritems() if value['stashed']])

    print "SUMMARY:"
    print "objects processed:              {}".format(len(stash.objects))
    print "replaced existing files on s3:  {}".format(stash.replace)
    print "images stashed:                 {}".format(images_stashed)
    print "files stashed:                  {}".format(files_stashed)
    print "thumbnails stashed:             {}".format(thumbs_stashed)
    print "media.json files stashed:       {}".format(mediajson_stashed)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Refactor main function for use from other scripts<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys, os
import argparse
import logging
import json
from deepharvest.deepharvest_nuxeo import DeepHarvestNuxeo
import s3stash.s3tools
from s3stash.nxstashref_image import NuxeoStashImage
from s3stash.nxstashref_file import NuxeoStashFile
from s3stash.nxstash_thumb import NuxeoStashThumb
from s3stash.nxstash_mediajson import NuxeoStashMediaJson

IMAGE_BUCKET = 'ucldc-private-files/jp2000'
IMAGE_REGION = 'us-west-2'
FILE_BUCKET = 'ucldc-nuxeo-ref-media'
FILE_REGION = 'us-west-2'
THUMB_BUCKET = 'static.ucldc.cdlib.org/ucldc-nuxeo-thumb-media'
THUMB_REGION = 'us-east-1'
MEDIAJSON_BUCKET = 'static.ucldc.cdlib.org/media_json'
MEDIAJSON_REGION = 'us-east-1'

_loglevel_ = 'INFO'

class Stash(object):
    '''
        stash various files on s3 for a Nuxeo collection
        in preparation for harvesting into Calisphere
    '''
    def __init__(self, path, pynuxrc, replace=False, loglevel=_loglevel_):
        self.logger = logging.getLogger(__name__)

        self.path = path
        self.pynuxrc = pynuxrc
        self.replace = replace

        self.dh = DeepHarvestNuxeo(self.path, '', pynuxrc=self.pynuxrc)

        self.objects = self.dh.fetch_objects()

    def images(self):
        ''' stash Nuxeo image files on s3 '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashImage(obj['path'], IMAGE_BUCKET, IMAGE_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashImage(c['path'], IMAGE_BUCKET, IMAGE_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def files(self):
        ''' stash Nuxeo files of type 'file', 'audio', or 'video' for a collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashFile(obj['path'], FILE_BUCKET, FILE_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashFile(c['path'], FILE_BUCKET, FILE_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def thumbnails(self):
        ''' stash thumbnail images for Nuxeo files of type 'file', 'audio', or 'video' for a collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashThumb(obj['path'], THUMB_BUCKET, THUMB_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()
            for c in self.dh.fetch_components(obj):
                nxstash = NuxeoStashThumb(c['path'], THUMB_BUCKET, THUMB_REGION, self.pynuxrc, self.replace)
                report[nxstash.uid] = nxstash.nxstashref()

        return report

    def media_json(self):
        ''' create and stash media.json files for a nuxeo collection '''
        report = {}
        for obj in self.objects:
            nxstash = NuxeoStashMediaJson(obj['path'], MEDIAJSON_BUCKET, MEDIAJSON_REGION, self.pynuxrc, self.replace)
            report[nxstash.uid] = nxstash.nxstashref()

        return report

def main(registry_id, pynuxrc="~/.pynuxrc", replace=True, loglevel=_loglevel_):
    # set up logging
    logfile = 'logs/stash_collection_{}'.format(registry_id)
    numeric_level = getattr(logging, loglevel, None)
    if not isinstance(numeric_level, int):
        raise ValueError('Invalid log level: %s' % loglevel)
    logging.basicConfig(level=numeric_level, filename=logfile, format='%(asctime)s (%(name)s) [%(levelname)s]: %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')
    logger = logging.getLogger(__name__)
    print "\nlogfile: {}\n".format(logfile)

    # get nuxeo path
    nxpath = s3stash.s3tools.get_nuxeo_path(registry_id)
    if nxpath is None:
        print "No record found for registry_id: {}".format(registry_id)
        sys.exit()
    info = 'nuxeo_path: {}'.format(nxpath)
    logger.info(info)
    print info, '\n'

    stash = Stash(nxpath, pynuxrc, replace)

    # stash images for use with iiif server
    print 'stashing images...'
    image_report = stash.images()
    info = 'finished stashing images'
    logger.info(info)
    print info
    reportfile = "reports/images-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(image_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash text, audio, video
    print 'stashing non-image files (text, audio, video)...'
    file_report = stash.files()
    info = 'finished stashing files'
    logger.info(info)
    print info
    reportfile = "reports/files-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(file_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash thumbnails for text, audio, video
    print 'stashing thumbnails for non-image files (text, audio, video)...'
    thumb_report = stash.thumbnails()
    info = 'finished stashing thumbnails'
    logger.info(info)
    print info
    reportfile = "reports/thumbs-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(thumb_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # stash media.json files
    print 'stashing media.json files for collection...'
    mediajson_report = stash.media_json()
    info = 'finished stashing media.json'
    logger.info(info)
    print info
    reportfile = "reports/mediajson-{}.json".format(registry_id)
    with open(reportfile, 'w') as f:
        json.dump(mediajson_report, f, sort_keys=True, indent=4)
    print "report:\t{}\n".format(reportfile)

    # print some information about how it went
    images_stashed = len([key for key, value in image_report.iteritems() if value['stashed']])
    files_stashed = len([key for key, value in file_report.iteritems() if value['stashed']])
    thumbs_stashed = len([key for key, value in thumb_report.iteritems() if value['stashed']])
    mediajson_stashed = len([key for key, value in mediajson_report.iteritems() if value['stashed']])

    print "SUMMARY:"
    print "objects processed:              {}".format(len(stash.objects))
    print "replaced existing files on s3:  {}".format(stash.replace)
    print "images stashed:                 {}".format(images_stashed)
    print "files stashed:                  {}".format(files_stashed)
    print "thumbnails stashed:             {}".format(thumbs_stashed)
    print "media.json files stashed:       {}".format(mediajson_stashed)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Deep harvest a Nuxeo collection')
    parser.add_argument('registry_id', help='UCLDC Registry ID')
    parser.add_argument('--pynuxrc', default='~/.pynuxrc', help='rc file for use by pynux')
    parser.add_argument('--replace', action='store_true', help='replace files on s3 if they already exist')
    parser.add_argument('--loglevel', default=_loglevel_)

    argv = parser.parse_args()

    registry_id = argv.registry_id
    pynuxrc = argv.pynuxrc
    replace = argv.replace
    loglevel = argv.loglevel

    sys.exit(main(registry_id,
                  pytnuxrc=pynuxrc,
                  replace=replace,
                  loglevel=loglevel))
<|endoftext|>


<commit_before>from optparse import make_option
import os
from shutil import rmtree
from subprocess import Popen, PIPE

from django.core.management import call_command
from django.core.management.base import BaseCommand

from django.core.management.commands.test import Command as TestCommand

from sbo_selenium.conf import settings
from sbo_selenium.testcase import sauce_sessions
from sbo_selenium.utils import OutputMonitor

class Command(BaseCommand):
    """
    Django management command for running Selenium tests.
    """
    args = '<package or test>'
    help = 'Run Selenium tests for this application'
    requires_model_validation = True
    custom_options = (
        make_option(
            '-b',
            '--browser',
            dest='browser_name',
            default='chrome',
            help='Name of the browser to run the tests in (default is chrome)'
        ),
        make_option(
            '-n',
            type='int',
            dest='count',
            default=1,
            help='Number of times to run each test'
        ),
        make_option(
            '-p',
            '--platform',
            dest='platform',
            help='OS and version thereof for the Sauce OnDemand VM to use'
        ),
        make_option(
            '--browser-version',
            dest='browser_version',
            help='Browser version for the Sauce OnDemand VM to use'
        ),
        make_option(
            '--tunnel-identifier',
            dest='tunnel_id',
            help='Sauce Connect tunnel identifier'
        )
    )
    option_list = TestCommand.option_list + custom_options

    @staticmethod
    def clean():
        """Clear out any old screenshots"""
        screenshot_dir = settings.SELENIUM_SCREENSHOT_DIR
        if screenshot_dir and os.path.isdir(screenshot_dir):
            rmtree(screenshot_dir, ignore_errors=True)

    def handle(self, *args, **options):
        """
        Run the specified Selenium test(s) the indicated number of times in
        the specified browser.
        """
        browser_name = options['browser_name']
        count = options['count']
        if len(args) > 0:
            tests = list(args)
        else:
            tests = settings.SELENIUM_DEFAULT_TESTS

        # Kill any orphaned chromedriver processes
        process = Popen(['killall', 'chromedriver'],
                        stderr=open(os.devnull, 'w'))
        process.wait()

        # Clear any old log and screenshots
        self.clean()

        sc_process = None
        selenium_process = None
        if 'platform' in options and settings.SELENIUM_SAUCE_CONNECT_PATH:
            running, sc_process = self.verify_sauce_connect_is_running(options)
            if not running:
                return
        elif browser_name in ['opera', 'safari']:
            running, selenium_process = self.verify_selenium_server_is_running()
            if not running:
                return
        elif browser_name in ['ipad', 'iphone']:
            if not self.verify_appium_is_running():
                return

        # Ugly hack: make it so django-nose won't have nosetests choke on our
        # parameters
        BaseCommand.option_list += self.custom_options

        # Configure and run the tests
        self.update_environment(options)
        self.run_tests(tests, browser_name, count)

        # Kill Sauce Connect, if running
        if sc_process:
            sc_process.kill()

        # Kill the Selenium standalone server, if running
        if selenium_process:
            selenium_process.kill()

    def run_tests(self, tests, browser_name, count):
        """Configure and run the tests"""
        test_args = ['test'] + tests
        test_options = settings.SELENIUM_TEST_COMMAND_OPTIONS
        for i in range(count):
            msg = 'Test run %d using %s' % (i + 1, browser_name)
            self.stdout.write(msg)
            call_command(*test_args, **test_options)
            for session in sauce_sessions:
                self.stdout.write(session)
            self.stdout.flush()
            del sauce_sessions[:]

    @staticmethod
    def update_environment(options):
        """
        Populate the environment variables that need to be added for test
        execution to work correctly.  Most (but not all) of these are to match
        what the Jenkins Sauce OnDemand plugin would use for the test
        configuration that was specified on the command line.
        """
        env = os.environ
        # https://docs.djangoproject.com/en/1.6/topics/testing/tools/#liveservertestcase
        env['DJANGO_LIVE_TEST_SERVER_ADDRESS'] = settings.DJANGO_LIVE_TEST_SERVER_ADDRESS
        tunnel_id = options['tunnel_id']
        if tunnel_id:
            env['SAUCE_TUNNEL_ID'] = tunnel_id
        if 'SAUCE_API_KEY' in env:
            # Jenkins plugin has already configured the environment for us
            return
        env['SELENIUM_BROWSER'] = options['browser_name']
        platform = options['platform']
        browser_version = options['browser_version']
        if not platform or not browser_version:
            # None of the following Sauce OnDemand stuff applies
            return
        if settings.SELENIUM_SAUCE_CONNECT_PATH:
            host = 'localhost'
            port = '4445'
        else:
            host = 'ondemand.saucelabs.com'
            port = '80'
        env.update({
            'SAUCE_API_KEY': settings.SELENIUM_SAUCE_API_KEY,
            'SAUCE_USER_NAME': settings.SELENIUM_SAUCE_USERNAME,
            'SELENIUM_BROWSER': options['browser_name'],
            'SELENIUM_HOST': host,
            'SELENIUM_PLATFORM': platform,
            'SELENIUM_PORT': port,
            'SELENIUM_VERSION': browser_version,
        })

    def verify_sauce_connect_is_running(self, options):
        """
        Start Sauce Connect, if it isn't already running.  Returns a tuple of
        two elements:

        * A boolean which is True if Sauce Connect is now running
        * The Popen object representing the process so it can be terminated
          later; if it was already running, this value is "None"
        """
        sc_path = settings.SELENIUM_SAUCE_CONNECT_PATH
        if len(sc_path) < 2:
            self.stdout.write('You need to configure SELENIUM_SAUCE_CONNECT_PATH')
            return False, None
        username = settings.SELENIUM_SAUCE_USERNAME
        if not username:
            self.stdout.write('You need to configure SELENIUM_SAUCE_USERNAME')
            return False, None
        key = settings.SELENIUM_SAUCE_API_KEY
        if not key:
            self.stdout.write('You need to configure SELENIUM_SAUCE_API_KEY')
            return False, None
        # Is it already running?
        process = Popen(['ps -e | grep "%s"' % key],
                        shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        grep_command = 'grep {}'.format(key)
        lines = grep_output.split('\n')
        for line in lines:
            if 'sc' in line and username in line and grep_command not in line:
                self.stdout.write('Sauce Connect is already running')
                return True, None
        self.stdout.write('Starting Sauce Connect')
        output = OutputMonitor()
        command = [sc_path, '-u', username, '-k', key]
        tunnel_id = options['tunnel_id']
        if tunnel_id:
            command.extend(['-i', tunnel_id])
        sc_process = Popen(command,
                           stdout=output.stream.input,
                           stderr=open(os.devnull, 'w'))
        ready_log_line = 'Connection established.'
        if not output.wait_for(ready_log_line, 60):
            self.stdout.write('Timeout starting Sauce Connect:\n')
            self.stdout.write('\n'.join(output.lines))
            return False, None
        return True, sc_process

    def verify_selenium_server_is_running(self):
        """
        Start the Selenium standalone server, if it isn't already running.
        Returns a tuple of two elements:

        * A boolean which is True if the server is now running
        * The Popen object representing the process so it can be terminated
          later; if the server was already running, this value is "None"
        """
        selenium_jar = settings.SELENIUM_JAR_PATH
        if len(selenium_jar) < 5:
            self.stdout.write('You need to configure SELENIUM_JAR_PATH')
            return False, None
        _jar_dir, jar_name = os.path.split(selenium_jar)
        # Is it already running?
        process = Popen(['ps -e | grep "%s"' % jar_name[:-4]],
                        shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        lines = grep_output.split('\n')
        for line in lines:
            if jar_name in line:
                self.stdout.write('Selenium standalone server is already running')
                return True, None
        self.stdout.write('Starting the Selenium standalone server')
        output = OutputMonitor()
        selenium_process = Popen(['java', '-jar', selenium_jar],
                                 stdout=output.stream.input,
                                 stderr=open(os.devnull, 'w'))
        ready_log_line = 'Started org.openqa.jetty.jetty.Server'
        if not output.wait_for(ready_log_line, 10):
            self.stdout.write('Timeout starting the Selenium server:\n')
            self.stdout.write('\n'.join(output.lines))
            return False, None
        return True, selenium_process

    def verify_appium_is_running(self):
        """Verify that Appium is running so it can be used for local iOS tests."""
        process = Popen(['ps -e | grep "Appium"'], shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        lines = grep_output.split('\n')
        for line in lines:
            if 'Appium.app' in line:
                self.stdout.write('Appium is already running')
                return True
        self.stdout.write('Please launch and configure Appium first')
        return False
<commit_msg>Use the default browser from the Django settings<commit_after>from optparse import make_option
import os
from shutil import rmtree
from subprocess import Popen, PIPE

from django.core.management import call_command
from django.core.management.base import BaseCommand

from django.core.management.commands.test import Command as TestCommand

from sbo_selenium.conf import settings
from sbo_selenium.testcase import sauce_sessions
from sbo_selenium.utils import OutputMonitor

class Command(BaseCommand):
    """
    Django management command for running Selenium tests.
    """
    args = '<package or test>'
    help = 'Run Selenium tests for this application'
    requires_model_validation = True
    custom_options = (
        make_option(
            '-b',
            '--browser',
            dest='browser_name',
            default=settings.SELENIUM_DEFAULT_BROWSER,
            help='Name of the browser to run the tests in (default is chrome)'
        ),
        make_option(
            '-n',
            type='int',
            dest='count',
            default=1,
            help='Number of times to run each test'
        ),
        make_option(
            '-p',
            '--platform',
            dest='platform',
            help='OS and version thereof for the Sauce OnDemand VM to use'
        ),
        make_option(
            '--browser-version',
            dest='browser_version',
            help='Browser version for the Sauce OnDemand VM to use'
        ),
        make_option(
            '--tunnel-identifier',
            dest='tunnel_id',
            help='Sauce Connect tunnel identifier'
        )
    )
    option_list = TestCommand.option_list + custom_options

    @staticmethod
    def clean():
        """Clear out any old screenshots"""
        screenshot_dir = settings.SELENIUM_SCREENSHOT_DIR
        if screenshot_dir and os.path.isdir(screenshot_dir):
            rmtree(screenshot_dir, ignore_errors=True)

    def handle(self, *args, **options):
        """
        Run the specified Selenium test(s) the indicated number of times in
        the specified browser.
        """
        browser_name = options['browser_name']
        count = options['count']
        if len(args) > 0:
            tests = list(args)
        else:
            tests = settings.SELENIUM_DEFAULT_TESTS

        # Kill any orphaned chromedriver processes
        process = Popen(['killall', 'chromedriver'],
                        stderr=open(os.devnull, 'w'))
        process.wait()

        # Clear any old log and screenshots
        self.clean()

        sc_process = None
        selenium_process = None
        if 'platform' in options and settings.SELENIUM_SAUCE_CONNECT_PATH:
            running, sc_process = self.verify_sauce_connect_is_running(options)
            if not running:
                return
        elif browser_name in ['opera', 'safari']:
            running, selenium_process = self.verify_selenium_server_is_running()
            if not running:
                return
        elif browser_name in ['ipad', 'iphone']:
            if not self.verify_appium_is_running():
                return

        # Ugly hack: make it so django-nose won't have nosetests choke on our
        # parameters
        BaseCommand.option_list += self.custom_options

        # Configure and run the tests
        self.update_environment(options)
        self.run_tests(tests, browser_name, count)

        # Kill Sauce Connect, if running
        if sc_process:
            sc_process.kill()

        # Kill the Selenium standalone server, if running
        if selenium_process:
            selenium_process.kill()

    def run_tests(self, tests, browser_name, count):
        """Configure and run the tests"""
        test_args = ['test'] + tests
        test_options = settings.SELENIUM_TEST_COMMAND_OPTIONS
        for i in range(count):
            msg = 'Test run %d using %s' % (i + 1, browser_name)
            self.stdout.write(msg)
            call_command(*test_args, **test_options)
            for session in sauce_sessions:
                self.stdout.write(session)
            self.stdout.flush()
            del sauce_sessions[:]

    @staticmethod
    def update_environment(options):
        """
        Populate the environment variables that need to be added for test
        execution to work correctly.  Most (but not all) of these are to match
        what the Jenkins Sauce OnDemand plugin would use for the test
        configuration that was specified on the command line.
        """
        env = os.environ
        # https://docs.djangoproject.com/en/1.6/topics/testing/tools/#liveservertestcase
        env['DJANGO_LIVE_TEST_SERVER_ADDRESS'] = settings.DJANGO_LIVE_TEST_SERVER_ADDRESS
        tunnel_id = options['tunnel_id']
        if tunnel_id:
            env['SAUCE_TUNNEL_ID'] = tunnel_id
        if 'SAUCE_API_KEY' in env:
            # Jenkins plugin has already configured the environment for us
            return
        env['SELENIUM_BROWSER'] = options['browser_name']
        platform = options['platform']
        browser_version = options['browser_version']
        if not platform or not browser_version:
            # None of the following Sauce OnDemand stuff applies
            return
        if settings.SELENIUM_SAUCE_CONNECT_PATH:
            host = 'localhost'
            port = '4445'
        else:
            host = 'ondemand.saucelabs.com'
            port = '80'
        env.update({
            'SAUCE_API_KEY': settings.SELENIUM_SAUCE_API_KEY,
            'SAUCE_USER_NAME': settings.SELENIUM_SAUCE_USERNAME,
            'SELENIUM_BROWSER': options['browser_name'],
            'SELENIUM_HOST': host,
            'SELENIUM_PLATFORM': platform,
            'SELENIUM_PORT': port,
            'SELENIUM_VERSION': browser_version,
        })

    def verify_sauce_connect_is_running(self, options):
        """
        Start Sauce Connect, if it isn't already running.  Returns a tuple of
        two elements:

        * A boolean which is True if Sauce Connect is now running
        * The Popen object representing the process so it can be terminated
          later; if it was already running, this value is "None"
        """
        sc_path = settings.SELENIUM_SAUCE_CONNECT_PATH
        if len(sc_path) < 2:
            self.stdout.write('You need to configure SELENIUM_SAUCE_CONNECT_PATH')
            return False, None
        username = settings.SELENIUM_SAUCE_USERNAME
        if not username:
            self.stdout.write('You need to configure SELENIUM_SAUCE_USERNAME')
            return False, None
        key = settings.SELENIUM_SAUCE_API_KEY
        if not key:
            self.stdout.write('You need to configure SELENIUM_SAUCE_API_KEY')
            return False, None
        # Is it already running?
        process = Popen(['ps -e | grep "%s"' % key],
                        shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        grep_command = 'grep {}'.format(key)
        lines = grep_output.split('\n')
        for line in lines:
            if 'sc' in line and username in line and grep_command not in line:
                self.stdout.write('Sauce Connect is already running')
                return True, None
        self.stdout.write('Starting Sauce Connect')
        output = OutputMonitor()
        command = [sc_path, '-u', username, '-k', key]
        tunnel_id = options['tunnel_id']
        if tunnel_id:
            command.extend(['-i', tunnel_id])
        sc_process = Popen(command,
                           stdout=output.stream.input,
                           stderr=open(os.devnull, 'w'))
        ready_log_line = 'Connection established.'
        if not output.wait_for(ready_log_line, 60):
            self.stdout.write('Timeout starting Sauce Connect:\n')
            self.stdout.write('\n'.join(output.lines))
            return False, None
        return True, sc_process

    def verify_selenium_server_is_running(self):
        """
        Start the Selenium standalone server, if it isn't already running.
        Returns a tuple of two elements:

        * A boolean which is True if the server is now running
        * The Popen object representing the process so it can be terminated
          later; if the server was already running, this value is "None"
        """
        selenium_jar = settings.SELENIUM_JAR_PATH
        if len(selenium_jar) < 5:
            self.stdout.write('You need to configure SELENIUM_JAR_PATH')
            return False, None
        _jar_dir, jar_name = os.path.split(selenium_jar)
        # Is it already running?
        process = Popen(['ps -e | grep "%s"' % jar_name[:-4]],
                        shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        lines = grep_output.split('\n')
        for line in lines:
            if jar_name in line:
                self.stdout.write('Selenium standalone server is already running')
                return True, None
        self.stdout.write('Starting the Selenium standalone server')
        output = OutputMonitor()
        selenium_process = Popen(['java', '-jar', selenium_jar],
                                 stdout=output.stream.input,
                                 stderr=open(os.devnull, 'w'))
        ready_log_line = 'Started org.openqa.jetty.jetty.Server'
        if not output.wait_for(ready_log_line, 10):
            self.stdout.write('Timeout starting the Selenium server:\n')
            self.stdout.write('\n'.join(output.lines))
            return False, None
        return True, selenium_process

    def verify_appium_is_running(self):
        """Verify that Appium is running so it can be used for local iOS tests."""
        process = Popen(['ps -e | grep "Appium"'], shell=True, stdout=PIPE)
        (grep_output, _grep_error) = process.communicate()
        lines = grep_output.split('\n')
        for line in lines:
            if 'Appium.app' in line:
                self.stdout.write('Appium is already running')
                return True
        self.stdout.write('Please launch and configure Appium first')
        return False
<|endoftext|>


<commit_before>import copy

from scrapy.http import Request
from scrapy.spider import BaseSpider
from scrapy.conf import settings

class Rule(object):
    """
    A rule for crawling, which receives the following constructor arguments:

    link_extractor (required)
       A LinkExtractor which defines the policy for extracting links
    callback (optional)
       A function to use to process the page once it has been downloaded. If
       callback is omitted the page is not procesed, just crawled. If callback
       is a string (instead of callable) a method of the spider class with that
       name is used as the callback function
    cb_kwargs (optional)
       A dict specifying keyword arguments to pass to the callback function
    follow (optional)
       If True, links will be followed from the pages crawled by this rule.
       It defaults to True when no callback is specified or False when a
       callback is specified
    process_links (optional)
       Can be either a callable, or a string with the name of a method defined
       in the spider's class.
       This method will be called with the list of extracted links matching
       this rule (if any) and must return another list of links.
    """

    def __init__(self, link_extractor, callback=None, cb_kwargs=None, follow=None, process_links=None):
        self.link_extractor = link_extractor
        self.callback = callback
        self.cb_kwargs = cb_kwargs or {}
        self.process_links = process_links
        if follow is None:
            self.follow = False if callback else True
        else:
            self.follow = follow

class CrawlSpider(BaseSpider):
    """
    Class for spiders that crawl over web pages and extract/parse their links
    given some crawling rules.

    These crawling rules are established by setting the 'rules' class attribute,
    which is a tuple of Rule objects.
    When the spider is running, it iterates over these rules with each response
    and do what it has to (extract links if follow=True, and return items/requests if
    there's a parsing method defined in the rule).
    """
    rules = ()

    def __init__(self):
        """Constructor takes care of compiling rules"""
        super(CrawlSpider, self).__init__()
        self._compile_rules()

    def parse(self, response):
        """This function is called by the framework core for all the
        start_urls. Do not override this function, override parse_start_url
        instead."""
        return self._response_downloaded(response, self.parse_start_url, cb_kwargs={}, follow=True)

    def parse_start_url(self, response):
        """Overrideable callback function for processing start_urls. It must
        return a list of ScrapedItems and/or Requests"""
        return []

    def process_spider_output(self, results, response):
        """This overridable method is called for each result (item or request)
        returned by the spider, and it's intended to perform any last time
        processing required before returning the results to the framework core,
        for example setting the item GUIDs. It receives a list of results and
        the response which originated that results. It must return a list
        of results (Items or Requests)."""
        return results

    def _requests_to_follow(self, response):
        """
        This method iterates over each of the spider's rules, extracts the links
        matching each case, filters them (if needed), and returns a list of unique
        requests per response.
        """
        requests = []
        seen = set()
        for rule in self._rules:
            links = [l for l in rule.link_extractor.extract_links(response) if l not in seen]
            if links and rule.process_links:
                links = rule.process_links(links)
            seen = seen.union(links)
            for link in links:
                r = Request(url=link.url, link_text=link.text)
                r.append_callback(self._response_downloaded, rule.callback, cb_kwargs=rule.cb_kwargs, follow=rule.follow)
                requests.append(r)
        return requests

    def _response_downloaded(self, response, callback, cb_kwargs, follow):
        """
        This is were any response arrives, and were it's decided whether
        to extract links or not from it, and if it will be parsed or not.
        It returns a list of requests/items.
        """
        res = []

        if follow and settings.getbool('CRAWLSPIDER_FOLLOW_LINKS', True):
            res.extend(self._requests_to_follow(response))
        if callback:
            cb_res = callback(response, **cb_kwargs) or ()
            cb_res = self.process_results(cb_res, response)
            res.extend(cb_res)
        return res

    def _compile_rules(self):
        """Compile the crawling rules"""

        def get_method(method):
            if callable(method):
                return method
            elif isinstance(method, basestring):
                return getattr(self, method, None)

        self._rules = [copy.copy(r) for r in self.rules]
        for rule in self._rules:
            rule.callback = get_method(rule.callback)
            rule.process_links = get_method(rule.process_links)
<commit_msg>fix to revision 660<commit_after>import copy

from scrapy.http import Request
from scrapy.spider import BaseSpider
from scrapy.conf import settings

class Rule(object):
    """
    A rule for crawling, which receives the following constructor arguments:

    link_extractor (required)
       A LinkExtractor which defines the policy for extracting links
    callback (optional)
       A function to use to process the page once it has been downloaded. If
       callback is omitted the page is not procesed, just crawled. If callback
       is a string (instead of callable) a method of the spider class with that
       name is used as the callback function
    cb_kwargs (optional)
       A dict specifying keyword arguments to pass to the callback function
    follow (optional)
       If True, links will be followed from the pages crawled by this rule.
       It defaults to True when no callback is specified or False when a
       callback is specified
    process_links (optional)
       Can be either a callable, or a string with the name of a method defined
       in the spider's class.
       This method will be called with the list of extracted links matching
       this rule (if any) and must return another list of links.
    """

    def __init__(self, link_extractor, callback=None, cb_kwargs=None, follow=None, process_links=None):
        self.link_extractor = link_extractor
        self.callback = callback
        self.cb_kwargs = cb_kwargs or {}
        self.process_links = process_links
        if follow is None:
            self.follow = False if callback else True
        else:
            self.follow = follow

class CrawlSpider(BaseSpider):
    """
    Class for spiders that crawl over web pages and extract/parse their links
    given some crawling rules.

    These crawling rules are established by setting the 'rules' class attribute,
    which is a tuple of Rule objects.
    When the spider is running, it iterates over these rules with each response
    and do what it has to (extract links if follow=True, and return items/requests if
    there's a parsing method defined in the rule).
    """
    rules = ()

    def __init__(self):
        """Constructor takes care of compiling rules"""
        super(CrawlSpider, self).__init__()
        self._compile_rules()

    def parse(self, response):
        """This function is called by the framework core for all the
        start_urls. Do not override this function, override parse_start_url
        instead."""
        return self._response_downloaded(response, self.parse_start_url, cb_kwargs={}, follow=True)

    def parse_start_url(self, response):
        """Overrideable callback function for processing start_urls. It must
        return a list of ScrapedItems and/or Requests"""
        return []

    def process_spider_output(self, results, response):
        """This overridable method is called for each result (item or request)
        returned by the spider, and it's intended to perform any last time
        processing required before returning the results to the framework core,
        for example setting the item GUIDs. It receives a list of results and
        the response which originated that results. It must return a list
        of results (Items or Requests)."""
        return results

    def _requests_to_follow(self, response):
        """
        This method iterates over each of the spider's rules, extracts the links
        matching each case, filters them (if needed), and returns a list of unique
        requests per response.
        """
        requests = []
        seen = set()
        for rule in self._rules:
            links = [l for l in rule.link_extractor.extract_links(response) if l not in seen]
            if links and rule.process_links:
                links = rule.process_links(links)
            seen = seen.union(links)
            for link in links:
                r = Request(url=link.url, link_text=link.text)
                r.append_callback(self._response_downloaded, rule.callback, cb_kwargs=rule.cb_kwargs, follow=rule.follow)
                requests.append(r)
        return requests

    def _response_downloaded(self, response, callback, cb_kwargs, follow):
        """
        This is were any response arrives, and were it's decided whether
        to extract links or not from it, and if it will be parsed or not.
        It returns a list of requests/items.
        """
        res = []

        if follow and settings.getbool('CRAWLSPIDER_FOLLOW_LINKS', True):
            res.extend(self._requests_to_follow(response))
        if callback:
            cb_res = callback(response, **cb_kwargs) or ()
            cb_res = self.process_spider_output(cb_res, response)
            res.extend(cb_res)
        return res

    def _compile_rules(self):
        """Compile the crawling rules"""

        def get_method(method):
            if callable(method):
                return method
            elif isinstance(method, basestring):
                return getattr(self, method, None)

        self._rules = [copy.copy(r) for r in self.rules]
        for rule in self._rules:
            rule.callback = get_method(rule.callback)
            rule.process_links = get_method(rule.process_links)
<|endoftext|>