<commit_before># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import six
from jsontableschema.model import SchemaModel
from sqlalchemy import (
        Table, Column, MetaData,
        Text, Integer, Float, Boolean,
        PrimaryKeyConstraint, ForeignKeyConstraint)

# Module API

class Storage(object):
    """SQL Tabular Storage.

    Parameters
    ----------
    engine: object
        SQLAlchemy engine.
    dbschema: str
        Database schema name.
    prefix: str
        Prefix for all tables.

    """

    # Public

    def __init__(self, engine, dbschema=None, prefix=''):

        # Set attributes
        self.__engine = engine
        self.__dbschema = dbschema
        self.__prefix = prefix

        # Create metadata
        self.__metadata = MetaData(
                bind=self.__engine,
                schema=self.__dbschema,
                reflect=True)

    def __repr__(self):

        # Template and format
        template = 'Storage <{engine}/{dbschema}>'
        text = template.format(
                engine=self.__engine,
                dbschema=self.__dbschema)

        return text

    @property
    def tables(self):
        """Return list of storage's table names.
        """

        # Collect
        tables = []
        for dbtable in self.__metadata.sorted_tables:
            table = dbtable.name
            table = _restore_table(self.__prefix, table)
            if table is not None:
                tables.append(table)

        return tables

    def check(self, table):
        """Return if table exists.
        """

        # Check existence
        existence = table in self.tables

        return existence

    def create(self, table, schema):
        """Create table by schema.

        Parameters
        ----------
        table: str/list
            Table name or list of table names.
        schema: dict/list
            JSONTableSchema schema or list of schemas.

        Raises
        ------
        RuntimeError
            If table already exists.

        """

        # Make lists
        tables = table
        if isinstance(table, six.string_types):
            tables = [table]
        schemas = schema
        if isinstance(schema, dict):
            schemas = [schema]

        # Iterate over tables/schemas
        for table, schema in zip(tables, schemas):

            # Check not existent
            if self.check(table):
                message = 'Table "%s" already exists.' % table
                raise RuntimeError(message)

            # Define table
            table = _convert_table(self.__prefix, table)
            columns, constraints = _convert_schema(
                    self.__prefix, table, schema)
            Table(table, self.__metadata, *(columns+constraints))

        # Create tables, update metadata
        self.__metadata.create_all()
        # Metadata reflect is auto

    def delete(self, table):
        """Delete table.

        Parameters
        ----------
        table: str/list
            Table name or list of table names.

        Raises
        ------
        RuntimeError
            If table doesn't exist.

        """

        # Make lists
        tables = table
        if isinstance(table, six.string_types):
            tables = [table]

        # Iterate over tables
        dbtables = []
        for table in tables:

            # Check existent
            if not self.check(table):
                message = 'Table "%s" doesn\'t exist.' % self
                raise RuntimeError(message)

            # Add table to dbtables
            dbtable = self.__get_dbtable(table)
            dbtables.append(dbtable)

        # Drop tables, update metadata
        self.__metadata.drop_all(tables=dbtables)
        self.__metadata.clear()
        self.__metadata.reflect()

    def describe(self, table):
        """Return table's JSONTableSchema schema.

        Parameters
        ----------
        table: str
            Table name.

        Returns
        -------
        dict
            JSONTableSchema schema.

        """

        # Get schema
        dbtable = self.__get_dbtable(table)
        table = _convert_table(self.__prefix, table)
        schema = _restore_schema(
                self.__prefix, table, dbtable.columns, dbtable.constraints)

        return schema

    def read(self, table):
        """Read data from table.

        Parameters
        ----------
        table: str
            Table name.

        Returns
        -------
        generator
            Data tuples generator.

        """

        # Get result
        dbtable = self.__get_dbtable(table)
        result = dbtable.select().execute()

        # Yield data
        schema = self.describe(table)
        model = SchemaModel(schema)
        for row in result:
            row = tuple(model.convert_row(*row))
            yield row

    def write(self, table, data):
        """Write data to table.

        Parameters
        ----------
        table: str
            Table name.
        data: list
            List of data tuples.

        """

        # Process data
        schema = self.describe(table)
        model = SchemaModel(schema)
        cdata = []
        for row in data:
            rdata = {}
            row = tuple(model.convert_row(*row))
            for index, field in enumerate(model.fields):
                rdata[field['name']] = row[index]
            cdata.append(rdata)

        # Insert data
        dbtable = self.__get_dbtable(table)
        dbtable.insert().execute(cdata)

    # Private

    def __get_dbtable(self, table):
        """Return dbtable instance from metadata.
        """

        # Prepare dict key
        key = _convert_table(self.__prefix, table)
        if self.__dbschema:
            key = '.'.join(self.__dbschema, key)

        return self.__metadata.tables[key]

# Internal

def _convert_table(prefix, table):
    """Convert high-level table name to database name.
    """
    return prefix + table

def _restore_table(prefix, table):
    """Restore database table name to high-level name.
    """
    if table.startswith(prefix):
        return table.replace(prefix, '', 1)
    return None

def _convert_schema(prefix, table, schema):  # noqa
    """Convert JSONTableSchema schema to SQLAlchemy columns and constraints.
    """

    # Init
    columns = []
    constraints = []

    # Mapping
    mapping = {
        'string': Text(),
        'integer': Integer(),
        'number': Float(),
        'boolean': Boolean(),
    }

    # Fields
    for field in schema['fields']:
        try:
            column_type = mapping[field['type']]
        except KeyError:
            message = 'Type %s is not supported' % field['type']
            raise TypeError(message)
        nullable = not field.get('constraints', {}).get('required', True)
        column = Column(field['name'], column_type, nullable=nullable)
        columns.append(column)

    # Primary key
    pk = schema.get('primaryKey', None)
    if pk is not None:
        if isinstance(pk, six.string_types):
            pk = [pk]
        constraint = PrimaryKeyConstraint(*pk)
        constraints.append(constraint)

    # Foreign keys
    fks = schema.get('foreignKeys', [])
    for fk in fks:
        fields = fk['fields']
        if isinstance(fields, six.string_types):
            fields = [fields]
        resource = fk['reference']['resource']
        if resource == 'self':
            resource = table
        elif resource == '<table>':
            resource = _convert_table(prefix, fk['reference']['table'])
        else:
            message = 'Supported only "self" and "<table>" references.'
            raise ValueError(message)
        references = fk['reference']['fields']
        if isinstance(references, six.string_types):
            references = [references]
        joiner = lambda reference: '.'.join([resource, reference])  # noqa
        references = list(map(joiner, references))
        constraint = ForeignKeyConstraint(fields, references)
        constraints.append(constraint)

    return (columns, constraints)

def _restore_schema(prefix, table, columns, constraints):  # noqa
    """Convert SQLAlchemy columns and constraints to JSONTableSchema schema.
    """

    # Init
    schema = {}

    # Mapping
    mapping = {
        Text: 'string',
        Integer: 'integer',
        Float: 'number',
        Boolean: 'boolean',
    }

    # Fields
    fields = []
    for column in columns:
        try:
            field_type = mapping[column.type.__class__]
        except KeyError:
            message = 'Type %s is not supported' % column.type
            raise TypeError(message)
        field = {'name': column.name, 'type': field_type}
        if column.nullable:
            field['constraints'] = {'required': False}
        fields.append(field)
    schema['fields'] = fields

    # Primary key
    pk = []
    for constraint in constraints:
        if isinstance(constraint, PrimaryKeyConstraint):
            for column in constraint.columns:
                pk.append(column.name)
    if len(pk) > 0:
        if len(pk) == 1:
            pk = pk.pop()
        schema['primaryKey'] = pk

    # Foreign keys
    fks = []
    for constraint in constraints:
        if isinstance(constraint, ForeignKeyConstraint):
            fields = []
            reftable = None
            resource = None
            references = []
            for element in constraint.elements:
                fields.append(element.parent.name)
                references.append(element.column.name)
                if element.column.table.name == table:
                    resource = 'self'
                else:
                    reftable = _restore_table(
                            prefix, element.column.table.name)
            if len(fields) == len(references) == 1:
                fields = fields.pop()
                references = references.pop()
            fk = {
                'fields': fields,
                'reference': {
                    'fields': references,
                }
            }
            if resource is not None:
                fk['reference']['resource'] = resource
            if reftable is not None:
                fk['reference']['resource'] = '<table>'
                fk['reference']['table'] = reftable
            fks.append(fk)
    if len(fks) > 0:
        schema['foreignKeys'] = fks

    return schema
<commit_msg>fixed Storage.create<commit_after># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import six
from jsontableschema.model import SchemaModel
from sqlalchemy import (
        Table, Column, MetaData,
        Text, Integer, Float, Boolean,
        PrimaryKeyConstraint, ForeignKeyConstraint)

# Module API

class Storage(object):
    """SQL Tabular Storage.

    Parameters
    ----------
    engine: object
        SQLAlchemy engine.
    dbschema: str
        Database schema name.
    prefix: str
        Prefix for all tables.

    """

    # Public

    def __init__(self, engine, dbschema=None, prefix=''):

        # Set attributes
        self.__engine = engine
        self.__dbschema = dbschema
        self.__prefix = prefix

        # Create metadata
        self.__metadata = MetaData(
                bind=self.__engine,
                schema=self.__dbschema,
                reflect=True)

    def __repr__(self):

        # Template and format
        template = 'Storage <{engine}/{dbschema}>'
        text = template.format(
                engine=self.__engine,
                dbschema=self.__dbschema)

        return text

    @property
    def tables(self):
        """Return list of storage's table names.
        """

        # Collect
        tables = []
        for dbtable in self.__metadata.sorted_tables:
            table = dbtable.name
            table = _restore_table(self.__prefix, table)
            if table is not None:
                tables.append(table)

        return tables

    def check(self, table):
        """Return if table exists.
        """

        # Check existence
        existence = table in self.tables

        return existence

    def create(self, table, schema):
        """Create table by schema.

        Parameters
        ----------
        table: str/list
            Table name or list of table names.
        schema: dict/list
            JSONTableSchema schema or list of schemas.

        Raises
        ------
        RuntimeError
            If table already exists.

        """

        # Make lists
        tables = table
        if isinstance(table, six.string_types):
            tables = [table]
        schemas = schema
        if isinstance(schema, dict):
            schemas = [schema]

        # Check tables for existence
        for table in tables:
            if self.check(table):
                message = 'Table "%s" already exists.' % table
                raise RuntimeError(message)

        # Define tables
        for table, schema in zip(tables, schemas):
            table = _convert_table(self.__prefix, table)
            columns, constraints = _convert_schema(
                    self.__prefix, table, schema)
            Table(table, self.__metadata, *(columns+constraints))

        # Create tables, update metadata
        self.__metadata.create_all()
        # Metadata reflect is auto

    def delete(self, table):
        """Delete table.

        Parameters
        ----------
        table: str/list
            Table name or list of table names.

        Raises
        ------
        RuntimeError
            If table doesn't exist.

        """

        # Make lists
        tables = table
        if isinstance(table, six.string_types):
            tables = [table]

        # Iterate over tables
        dbtables = []
        for table in tables:

            # Check existent
            if not self.check(table):
                message = 'Table "%s" doesn\'t exist.' % self
                raise RuntimeError(message)

            # Add table to dbtables
            dbtable = self.__get_dbtable(table)
            dbtables.append(dbtable)

        # Drop tables, update metadata
        self.__metadata.drop_all(tables=dbtables)
        self.__metadata.clear()
        self.__metadata.reflect()

    def describe(self, table):
        """Return table's JSONTableSchema schema.

        Parameters
        ----------
        table: str
            Table name.

        Returns
        -------
        dict
            JSONTableSchema schema.

        """

        # Get schema
        dbtable = self.__get_dbtable(table)
        table = _convert_table(self.__prefix, table)
        schema = _restore_schema(
                self.__prefix, table, dbtable.columns, dbtable.constraints)

        return schema

    def read(self, table):
        """Read data from table.

        Parameters
        ----------
        table: str
            Table name.

        Returns
        -------
        generator
            Data tuples generator.

        """

        # Get result
        dbtable = self.__get_dbtable(table)
        result = dbtable.select().execute()

        # Yield data
        schema = self.describe(table)
        model = SchemaModel(schema)
        for row in result:
            row = tuple(model.convert_row(*row))
            yield row

    def write(self, table, data):
        """Write data to table.

        Parameters
        ----------
        table: str
            Table name.
        data: list
            List of data tuples.

        """

        # Process data
        schema = self.describe(table)
        model = SchemaModel(schema)
        cdata = []
        for row in data:
            rdata = {}
            row = tuple(model.convert_row(*row))
            for index, field in enumerate(model.fields):
                rdata[field['name']] = row[index]
            cdata.append(rdata)

        # Insert data
        dbtable = self.__get_dbtable(table)
        dbtable.insert().execute(cdata)

    # Private

    def __get_dbtable(self, table):
        """Return dbtable instance from metadata.
        """

        # Prepare dict key
        key = _convert_table(self.__prefix, table)
        if self.__dbschema:
            key = '.'.join(self.__dbschema, key)

        return self.__metadata.tables[key]

# Internal

def _convert_table(prefix, table):
    """Convert high-level table name to database name.
    """
    return prefix + table

def _restore_table(prefix, table):
    """Restore database table name to high-level name.
    """
    if table.startswith(prefix):
        return table.replace(prefix, '', 1)
    return None

def _convert_schema(prefix, table, schema):  # noqa
    """Convert JSONTableSchema schema to SQLAlchemy columns and constraints.
    """

    # Init
    columns = []
    constraints = []

    # Mapping
    mapping = {
        'string': Text(),
        'integer': Integer(),
        'number': Float(),
        'boolean': Boolean(),
    }

    # Fields
    for field in schema['fields']:
        try:
            column_type = mapping[field['type']]
        except KeyError:
            message = 'Type %s is not supported' % field['type']
            raise TypeError(message)
        nullable = not field.get('constraints', {}).get('required', True)
        column = Column(field['name'], column_type, nullable=nullable)
        columns.append(column)

    # Primary key
    pk = schema.get('primaryKey', None)
    if pk is not None:
        if isinstance(pk, six.string_types):
            pk = [pk]
        constraint = PrimaryKeyConstraint(*pk)
        constraints.append(constraint)

    # Foreign keys
    fks = schema.get('foreignKeys', [])
    for fk in fks:
        fields = fk['fields']
        if isinstance(fields, six.string_types):
            fields = [fields]
        resource = fk['reference']['resource']
        if resource == 'self':
            resource = table
        elif resource == '<table>':
            resource = _convert_table(prefix, fk['reference']['table'])
        else:
            message = 'Supported only "self" and "<table>" references.'
            raise ValueError(message)
        references = fk['reference']['fields']
        if isinstance(references, six.string_types):
            references = [references]
        joiner = lambda reference: '.'.join([resource, reference])  # noqa
        references = list(map(joiner, references))
        constraint = ForeignKeyConstraint(fields, references)
        constraints.append(constraint)

    return (columns, constraints)

def _restore_schema(prefix, table, columns, constraints):  # noqa
    """Convert SQLAlchemy columns and constraints to JSONTableSchema schema.
    """

    # Init
    schema = {}

    # Mapping
    mapping = {
        Text: 'string',
        Integer: 'integer',
        Float: 'number',
        Boolean: 'boolean',
    }

    # Fields
    fields = []
    for column in columns:
        try:
            field_type = mapping[column.type.__class__]
        except KeyError:
            message = 'Type %s is not supported' % column.type
            raise TypeError(message)
        field = {'name': column.name, 'type': field_type}
        if column.nullable:
            field['constraints'] = {'required': False}
        fields.append(field)
    schema['fields'] = fields

    # Primary key
    pk = []
    for constraint in constraints:
        if isinstance(constraint, PrimaryKeyConstraint):
            for column in constraint.columns:
                pk.append(column.name)
    if len(pk) > 0:
        if len(pk) == 1:
            pk = pk.pop()
        schema['primaryKey'] = pk

    # Foreign keys
    fks = []
    for constraint in constraints:
        if isinstance(constraint, ForeignKeyConstraint):
            fields = []
            reftable = None
            resource = None
            references = []
            for element in constraint.elements:
                fields.append(element.parent.name)
                references.append(element.column.name)
                if element.column.table.name == table:
                    resource = 'self'
                else:
                    reftable = _restore_table(
                            prefix, element.column.table.name)
            if len(fields) == len(references) == 1:
                fields = fields.pop()
                references = references.pop()
            fk = {
                'fields': fields,
                'reference': {
                    'fields': references,
                }
            }
            if resource is not None:
                fk['reference']['resource'] = resource
            if reftable is not None:
                fk['reference']['resource'] = '<table>'
                fk['reference']['table'] = reftable
            fks.append(fk)
    if len(fks) > 0:
        schema['foreignKeys'] = fks

    return schema
<|endoftext|>


<commit_before>
#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

import operator
import itertools
import models # damn!
import pprint
import sys

COL_INDICATOR = 0
COL_TEST = 1
COL_RESULTS_DATA = 2
COL_RESULTS_NUM = 3
COL_HIERARCHY = 4
COL_PACKAGE = 5

class NoRelevantResults(Exception): pass

def reform_dict(d):
    def inner(k1):
        matches_first = lambda i: i[0] == k1
        return dict([ (k2, d[(k1, k2)]) for k2 in
                      map(lambda i: i[1], filter(matches_first, d.keys())) ])

    return dict([ (k1, inner(k1))
                   for k1 in set( k1 for k1, k2 in d.keys() ) ])

def remove_empty_dicts(h):
    has_keys = lambda kvp: len(kvp[1])
    return dict([
            (K, dict(filter(has_keys, V.items())))
            for K, V in h.items()
            ])

class TestInfo(object):
    def __init__(self, test_id, results_pct, results_num, sampling_ok):
        test = models.Test.query.filter(models.Test.id == test_id).first()
        self.test_id = test_id
        self.test_name = test.name
        self.test_description = test.description
        self.test_group = test.test_group
        self.test_level = test.test_level
        self.results_raw_score = results_pct
        self.results_num = results_num
        self.sampling_ok = sampling_ok

        if sampling_ok:
            self.results_pct = self.results_raw_score
        else:
            self.results_pct = 0.0

    def as_dict(self):
        return {
            "test": {
                "id": self.test_id,
                "name": self.test_name,
                "description": self.test_description,
                "test_group": self.test_group,
                "test_level": self.test_level
                },
            "results_pct": self.results_pct,
            "results_num": self.results_num,
            "sampling_ok": self.sampling_ok,
            "results_raw_score": self.results_raw_score
            }

class IndicatorInfo(object):
    def __init__(self, indicator_id):
        self.indicator_id = indicator_id
        self.ind = models.Indicator.query.filter(
            models.Indicator.id == self.indicator_id).first()

    def as_dict(self):
        return self.ind.as_dict()

    def as_dict_minus_group(self):
        tmp = self.ind.as_dict()
        del(tmp['indicatorgroup_id'])
        return tmp

def publisher_indicators(indicators, indicators_tests, simple_out):
    # get all tests which belong to a specific indicator
    # average the results for all tests in that indicator
    def per_indicator(indicator):
        indicator_test_data = []
        results_pct = 0.0
        results_num = 0.0
        results_weighted_pct_average_numerator = 0.0

        relevant = lambda test: (indicator, test) in indicators_tests

        for test in filter(relevant, simple_out.keys()):
            indic_info = simple_out[test]
            results_pct += indic_info["results_pct"]
            results_num += indic_info["results_num"]
            results_weighted_pct_average_numerator += (
                indic_info["results_pct"] *
                indic_info["results_num"]
                )
            indicator_test_data.append(indic_info)

        return {
            "indicator": IndicatorInfo(indicator).as_dict_minus_group(),
            "tests": indicator_test_data,
            "results_pct": (results_weighted_pct_average_numerator/results_num),
            "results_num": results_num
            }

    return dict([ (i, per_indicator(i)) for i in indicators ])

def make_summary(test_id, results_pct, results_num, sampling_ok):
    t = TestInfo(test_id, results_pct, results_num, sampling_ok)
    return t.as_dict()

def publisher_simple(out, cdtns):
    hierarchies = set(out)
    tests = set()
    for h in hierarchies:
        tests.update(set(out[h]))

    def per_test(t):
        results_pct = 0.0
        results_num = 0.0
        results_weighted_pct_average_numerator = 0.0

        def relevant(hierarchy):
            key = (t,'activity hierarchy', str(hierarchy))

            if t not in out[hierarchy]:
                return False

            if not cdtns:
                return True

            if key not in cdtns:
                return True

            if cdtns[key][0] == 0:
                return False

            return True

        # This makes sure information about a test is returned.
        def get_okhierarchy(out, t):
            for h in out:
                if (t in out[h] and 'test' in out[h][t]):
                    return h
            raise NoRelevantResults(
                "Summary could not be generated for test %d" % t
                )

        for hierarchy in filter(relevant, hierarchies):
            test_info = out[hierarchy][t]

            results_pct += test_info["results_pct"]
            results_num += test_info["results_num"]
            results_weighted_pct_average_numerator += (
                test_info["results_pct"] *
                test_info["results_num"]
                )

        if results_num == 0:
            raise NoRelevantResults("Results_num == 0 for test: %d" % t)

        # Result aggregation throws away hierarchies if there are 0 results.
        # This means some hierarchies won't have the 'test' dict. But, at
        # least one must, because we wouldn't be here otherwise.
        okhierarchy = get_okhierarchy(out, t)

        tmp = make_summary(
            out[okhierarchy][t]['test']["id"],
            (results_weighted_pct_average_numerator / results_num),
            results_num,
            True
            )
        tmp["indicator"] = out[okhierarchy][t]['indicator']
        return tmp

    return dict([ (t, per_test(t)) for t in tests ])

def sum_for_publishers(packages, d, h, t):
    # aggregate data across multiple packages for a single publisher ;
    # for each package, add percentage for each ;
    # need below to only include packages that are in this hierarchy

    relevant = lambda p: (h, t, p) in d
    relevant_data = map(lambda p: d[(h, t, p)], filter(relevant, packages))

    ## FIXME
    # this is an appalling hack
    def indicator_for_test(test_id):
        return relevant_data[-1][0]

    pct = lambda i: i[COL_RESULTS_DATA]
    activities = lambda i: i[COL_RESULTS_NUM]

    total_pct        = reduce(operator.add, map(pct,        relevant_data), 0)
    total_activities = reduce(operator.add, map(activities, relevant_data), 0)

    packages_in_hierarchy = len(relevant_data)

    if total_activities <= 0:
        return {}

    indicator_id = indicator_for_test(t)

    tmp = make_summary(
        t,
        float(total_pct/packages_in_hierarchy),
        total_activities,
        True
        )
    tmp["indicator"] = IndicatorInfo(indicator_id).as_dict()
    return tmp

class Summary(object):
    def __init__(self, data, conditions=None, manual=False):
        self.data = data
        self.conditions = conditions
        self.manual = manual
        self._summary = self.calculate()

    def calculate(self):
        if self.manual:
            return self.aggregate()

        def replace_first(tupl, newval):
            return tuple([newval] + list(tupl)[COL_TEST:])
        switch_first = lambda t: replace_first(t, t[COL_INDICATOR])
        self.data = map(switch_first, self.data)

        return self.aggregate()

    def summary(self):
        return self._summary

    def gen_hierarchies(self):
        for i in set(map(lambda x: (x[COL_HIERARCHY]), self.data)):
            yield i

    def gen_tests(self):
        for i in set(map(lambda x: (x[COL_TEST]), self.data)):
            yield i

    def get_conditions(self):
        if not self.conditions:
            return None

        return dict(map(lambda x: (
                    (x.test_id, x.condition, x.condition_value),
                    (x.operation, x.description)
                    ), self.conditions))

    def restructure_data(self):
        return lambda x: ((x[COL_HIERARCHY], x[COL_TEST], x[COL_PACKAGE]), (x))

    def setmap(self, lam):
        return set(map(lam, self.data))

    def get_indicator_data(self):
        ind_f = lambda x: x[COL_INDICATOR]
        return ind_f

    def aggregate(self):
        hierarchies = self.gen_hierarchies()
        tests = self.gen_tests()
        cdtns = self.get_conditions()

        d_f = self.restructure_data()

        d = dict(map(d_f, self.data))

        indicators = self.setmap(self.get_indicator_data())

        ind_test_f = lambda x: (x[COL_INDICATOR], x[COL_TEST])
        indicators_tests = list(self.setmap(ind_test_f))

        pkg_f = lambda x: x[COL_PACKAGE]
        packages = self.setmap(pkg_f)

        summary = lambda h, t: sum_for_publishers(packages, d, h, t)

        return self.summarise_results(hierarchies,
                                 tests, cdtns, indicators,
                                 indicators_tests, packages, d, summary)

    def summarise_results(self, hierarchies,
                      tests, cdtns, indicators,
                      indicators_tests, packages, d, summary):

        def summaries(summary_f):
            for h, t in itertools.product(hierarchies, tests):
                yield h, t, summary_f(h, t)

        def add_condition(i):
            h, t, tdata = i
            if self.conditions:
                key = (t, 'activity hierarchy', str(h))
                if key in cdtns:
                    tdata["condition"] = cdtns[key]
            return h, t, tdata

        summaries = (add_condition(i) for i in summaries(summary))

        even_more_tmp_out = [ ((h, t), tdata) for h, t, tdata in summaries ]
        for i in even_more_tmp_out:
            pass # print >>sys.stderr, i
        tmp_out = dict(even_more_tmp_out)

        out = reform_dict(tmp_out)

        out = remove_empty_dicts(out)
        return self.add_indicator_info(out, cdtns, indicators,
                                       indicators_tests)

    def add_indicator_info(self, out, cdtns, indicators,
                           indicators_tests):
        return out

class PublisherSummary(Summary):
    pass

class PublisherIndicatorsSummary(PublisherSummary):
    def add_indicator_info(self, out, cdtns, indicators,
                           indicators_tests):

        simple_out = publisher_simple(out, cdtns)
        return publisher_indicators(indicators, indicators_tests, simple_out)

class SummaryCreator(object):
    @property
    def summary(self):
        return self._summary

    @property
    def aggregate_results(self):
        return self._aggregate_results

from models import *

class PublisherSummaryCreator(SummaryCreator):
    def __init__(self, organisation, aggregation_type):
        aggregate_results = db.session.query(
            Indicator.id,
            Test.id,
            AggregateResult.results_data,
            AggregateResult.results_num,
            AggregateResult.result_hierarchy,
            AggregateResult.package_id
            ).filter(
            AggregateResult.organisation_id==organisation.id
            ).filter(
                AggregateResult.aggregateresulttype_id == aggregation_type
            )

        aggregate_results2 = aggregate_results.group_by(
            Indicator.id,
            AggregateResult.result_hierarchy,
            Test.id,
            AggregateResult.package_id,
            AggregateResult.results_data,
            AggregateResult.results_num
        ).join(IndicatorTest
        ).join(Test
        ).join(AggregateResult
        ).join(Package
        ).join(OrganisationPackage
        ).join(Organisation
        ).all()

        pconditions = OrganisationCondition.query.filter_by(
            organisation_id=organisation.id
            ).all()

        self._aggregate_results = aggregate_results2

        self._summary = PublisherSummary(
            self._aggregate_results, conditions=pconditions)

class PackageSummaryCreator(SummaryCreator):
    def __init__(self, package, latest_runtime, aggregation_type):
        p = package

        self._aggregate_results = db.session.query(
            Indicator.id,
            Test.id,
            AggregateResult.results_data,
            AggregateResult.results_num,
            AggregateResult.result_hierarchy,
            AggregateResult.package_id
        ).filter(
            AggregateResult.package_id==p[0].id,
            AggregateResult.aggregateresulttype_id==aggregation_type
        ).group_by(
            AggregateResult.result_hierarchy,
            Test.id,
            AggregateResult.package_id,
            Indicator.id,
            AggregateResult.results_data,
            AggregateResult.results_num
            ).join(IndicatorTest
            ).join(Test
            ).join(AggregateResult
            ).all()

        self._summary = PublisherSummary(self._aggregate_results, {}).summary()

class PublisherIndicatorsSummaryCreator(SummaryCreator):
    def __init__(self, organisation, aggregation_type):
        aggregate_results = db.session.query(Indicator.id,
                                     Test.id,
                                     AggregateResult.results_data,
                                     AggregateResult.results_num,
                                     AggregateResult.result_hierarchy,
                                     AggregateResult.package_id
        ).filter(Organisation.organisation_code==organisation.organisation_code
        ).filter(AggregateResult.aggregateresulttype_id == aggregation_type
        ).filter(AggregateResult.organisation_id == organisation.id
        ).group_by(AggregateResult.result_hierarchy,
                   Test.id,
                   AggregateResult.package_id,
                   Indicator.id,
                   AggregateResult.results_data,
                   AggregateResult.results_num
        ).join(IndicatorTest
        ).join(Test
        ).join(AggregateResult
        ).join(Package
        ).join(OrganisationPackage
        ).join(Organisation
        ).order_by(Indicator.indicator_type,
                   Indicator.indicator_category_name,
                   Indicator.indicator_subcategory_name
        ).all()

        pconditions = OrganisationCondition.query.filter_by(
            organisation_id=organisation.id
            ).all()

        self._aggregate_results = aggregate_results

        self._summary = PublisherIndicatorsSummary(
            aggregate_results,
            conditions=pconditions)

<commit_msg>factor publisher conditions lookups together<commit_after>
#  IATI Data Quality, tools for Data QA on IATI-formatted  publications
#  by Mark Brough, Martin Keegan, Ben Webb and Jennifer Smith
#
#  Copyright (C) 2013  Publish What You Fund
#
#  This programme is free software; you may redistribute and/or modify
#  it under the terms of the GNU Affero General Public License v3.0

import operator
import itertools
import models # damn!
import pprint
import sys

COL_INDICATOR = 0
COL_TEST = 1
COL_RESULTS_DATA = 2
COL_RESULTS_NUM = 3
COL_HIERARCHY = 4
COL_PACKAGE = 5

class NoRelevantResults(Exception): pass

def reform_dict(d):
    def inner(k1):
        matches_first = lambda i: i[0] == k1
        return dict([ (k2, d[(k1, k2)]) for k2 in
                      map(lambda i: i[1], filter(matches_first, d.keys())) ])

    return dict([ (k1, inner(k1))
                   for k1 in set( k1 for k1, k2 in d.keys() ) ])

def remove_empty_dicts(h):
    has_keys = lambda kvp: len(kvp[1])
    return dict([
            (K, dict(filter(has_keys, V.items())))
            for K, V in h.items()
            ])

class TestInfo(object):
    def __init__(self, test_id, results_pct, results_num, sampling_ok):
        test = models.Test.query.filter(models.Test.id == test_id).first()
        self.test_id = test_id
        self.test_name = test.name
        self.test_description = test.description
        self.test_group = test.test_group
        self.test_level = test.test_level
        self.results_raw_score = results_pct
        self.results_num = results_num
        self.sampling_ok = sampling_ok

        if sampling_ok:
            self.results_pct = self.results_raw_score
        else:
            self.results_pct = 0.0

    def as_dict(self):
        return {
            "test": {
                "id": self.test_id,
                "name": self.test_name,
                "description": self.test_description,
                "test_group": self.test_group,
                "test_level": self.test_level
                },
            "results_pct": self.results_pct,
            "results_num": self.results_num,
            "sampling_ok": self.sampling_ok,
            "results_raw_score": self.results_raw_score
            }

class IndicatorInfo(object):
    def __init__(self, indicator_id):
        self.indicator_id = indicator_id
        self.ind = models.Indicator.query.filter(
            models.Indicator.id == self.indicator_id).first()

    def as_dict(self):
        return self.ind.as_dict()

    def as_dict_minus_group(self):
        tmp = self.ind.as_dict()
        del(tmp['indicatorgroup_id'])
        return tmp

def publisher_indicators(indicators, indicators_tests, simple_out):
    # get all tests which belong to a specific indicator
    # average the results for all tests in that indicator
    def per_indicator(indicator):
        indicator_test_data = []
        results_pct = 0.0
        results_num = 0.0
        results_weighted_pct_average_numerator = 0.0

        relevant = lambda test: (indicator, test) in indicators_tests

        for test in filter(relevant, simple_out.keys()):
            indic_info = simple_out[test]
            results_pct += indic_info["results_pct"]
            results_num += indic_info["results_num"]
            results_weighted_pct_average_numerator += (
                indic_info["results_pct"] *
                indic_info["results_num"]
                )
            indicator_test_data.append(indic_info)

        return {
            "indicator": IndicatorInfo(indicator).as_dict_minus_group(),
            "tests": indicator_test_data,
            "results_pct": (results_weighted_pct_average_numerator/results_num),
            "results_num": results_num
            }

    return dict([ (i, per_indicator(i)) for i in indicators ])

def make_summary(test_id, results_pct, results_num, sampling_ok):
    t = TestInfo(test_id, results_pct, results_num, sampling_ok)
    return t.as_dict()

def publisher_simple(out, cdtns):
    hierarchies = set(out)
    tests = set()
    for h in hierarchies:
        tests.update(set(out[h]))

    def per_test(t):
        results_pct = 0.0
        results_num = 0.0
        results_weighted_pct_average_numerator = 0.0

        def relevant(hierarchy):
            key = (t,'activity hierarchy', str(hierarchy))

            if t not in out[hierarchy]:
                return False

            if not cdtns:
                return True

            if key not in cdtns:
                return True

            if cdtns[key][0] == 0:
                return False

            return True

        # This makes sure information about a test is returned.
        def get_okhierarchy(out, t):
            for h in out:
                if (t in out[h] and 'test' in out[h][t]):
                    return h
            raise NoRelevantResults(
                "Summary could not be generated for test %d" % t
                )

        for hierarchy in filter(relevant, hierarchies):
            test_info = out[hierarchy][t]

            results_pct += test_info["results_pct"]
            results_num += test_info["results_num"]
            results_weighted_pct_average_numerator += (
                test_info["results_pct"] *
                test_info["results_num"]
                )

        if results_num == 0:
            raise NoRelevantResults("Results_num == 0 for test: %d" % t)

        # Result aggregation throws away hierarchies if there are 0 results.
        # This means some hierarchies won't have the 'test' dict. But, at
        # least one must, because we wouldn't be here otherwise.
        okhierarchy = get_okhierarchy(out, t)

        tmp = make_summary(
            out[okhierarchy][t]['test']["id"],
            (results_weighted_pct_average_numerator / results_num),
            results_num,
            True
            )
        tmp["indicator"] = out[okhierarchy][t]['indicator']
        return tmp

    return dict([ (t, per_test(t)) for t in tests ])

def sum_for_publishers(packages, d, h, t):
    # aggregate data across multiple packages for a single publisher ;
    # for each package, add percentage for each ;
    # need below to only include packages that are in this hierarchy

    relevant = lambda p: (h, t, p) in d
    relevant_data = map(lambda p: d[(h, t, p)], filter(relevant, packages))

    ## FIXME
    # this is an appalling hack
    def indicator_for_test(test_id):
        return relevant_data[-1][0]

    pct = lambda i: i[COL_RESULTS_DATA]
    activities = lambda i: i[COL_RESULTS_NUM]

    total_pct        = reduce(operator.add, map(pct,        relevant_data), 0)
    total_activities = reduce(operator.add, map(activities, relevant_data), 0)

    packages_in_hierarchy = len(relevant_data)

    if total_activities <= 0:
        return {}

    indicator_id = indicator_for_test(t)

    tmp = make_summary(
        t,
        float(total_pct/packages_in_hierarchy),
        total_activities,
        True
        )
    tmp["indicator"] = IndicatorInfo(indicator_id).as_dict()
    return tmp

class Summary(object):
    def __init__(self, data, conditions=None, manual=False):
        self.data = data
        self.conditions = conditions
        self.manual = manual
        self._summary = self.calculate()

    def calculate(self):
        if self.manual:
            return self.aggregate()

        def replace_first(tupl, newval):
            return tuple([newval] + list(tupl)[COL_TEST:])
        switch_first = lambda t: replace_first(t, t[COL_INDICATOR])
        self.data = map(switch_first, self.data)

        return self.aggregate()

    def summary(self):
        return self._summary

    def gen_hierarchies(self):
        for i in set(map(lambda x: (x[COL_HIERARCHY]), self.data)):
            yield i

    def gen_tests(self):
        for i in set(map(lambda x: (x[COL_TEST]), self.data)):
            yield i

    def get_conditions(self):
        if not self.conditions:
            return None

        return dict(map(lambda x: (
                    (x.test_id, x.condition, x.condition_value),
                    (x.operation, x.description)
                    ), self.conditions))

    def restructure_data(self):
        return lambda x: ((x[COL_HIERARCHY], x[COL_TEST], x[COL_PACKAGE]), (x))

    def setmap(self, lam):
        return set(map(lam, self.data))

    def get_indicator_data(self):
        ind_f = lambda x: x[COL_INDICATOR]
        return ind_f

    def aggregate(self):
        hierarchies = self.gen_hierarchies()
        tests = self.gen_tests()
        cdtns = self.get_conditions()

        d_f = self.restructure_data()

        d = dict(map(d_f, self.data))

        indicators = self.setmap(self.get_indicator_data())

        ind_test_f = lambda x: (x[COL_INDICATOR], x[COL_TEST])
        indicators_tests = list(self.setmap(ind_test_f))

        pkg_f = lambda x: x[COL_PACKAGE]
        packages = self.setmap(pkg_f)

        summary = lambda h, t: sum_for_publishers(packages, d, h, t)

        return self.summarise_results(hierarchies,
                                 tests, cdtns, indicators,
                                 indicators_tests, packages, d, summary)

    def summarise_results(self, hierarchies,
                      tests, cdtns, indicators,
                      indicators_tests, packages, d, summary):

        def summaries(summary_f):
            for h, t in itertools.product(hierarchies, tests):
                yield h, t, summary_f(h, t)

        def add_condition(i):
            h, t, tdata = i
            if self.conditions:
                key = (t, 'activity hierarchy', str(h))
                if key in cdtns:
                    tdata["condition"] = cdtns[key]
            return h, t, tdata

        summaries = (add_condition(i) for i in summaries(summary))

        even_more_tmp_out = [ ((h, t), tdata) for h, t, tdata in summaries ]
        for i in even_more_tmp_out:
            pass # print >>sys.stderr, i
        tmp_out = dict(even_more_tmp_out)

        out = reform_dict(tmp_out)

        out = remove_empty_dicts(out)
        return self.add_indicator_info(out, cdtns, indicators,
                                       indicators_tests)

    def add_indicator_info(self, out, cdtns, indicators,
                           indicators_tests):
        return out

class PublisherSummary(Summary):
    pass

class PublisherIndicatorsSummary(PublisherSummary):
    def add_indicator_info(self, out, cdtns, indicators,
                           indicators_tests):

        simple_out = publisher_simple(out, cdtns)
        return publisher_indicators(indicators, indicators_tests, simple_out)

from models import *

class SummaryCreator(object):
    @property
    def summary(self):
        return self._summary

    @property
    def aggregate_results(self):
        return self._aggregate_results

    def conditions_for_organisation(self, organisation_id):
        return OrganisationCondition.query.filter_by(
            organisation_id=organisation_id
            ).all()

class PublisherSummaryCreator(SummaryCreator):
    def __init__(self, organisation, aggregation_type):
        aggregate_results = db.session.query(
            Indicator.id,
            Test.id,
            AggregateResult.results_data,
            AggregateResult.results_num,
            AggregateResult.result_hierarchy,
            AggregateResult.package_id
            ).filter(
            AggregateResult.organisation_id==organisation.id
            ).filter(
                AggregateResult.aggregateresulttype_id == aggregation_type
            )

        aggregate_results2 = aggregate_results.group_by(
            Indicator.id,
            AggregateResult.result_hierarchy,
            Test.id,
            AggregateResult.package_id,
            AggregateResult.results_data,
            AggregateResult.results_num
        ).join(IndicatorTest
        ).join(Test
        ).join(AggregateResult
        ).join(Package
        ).join(OrganisationPackage
        ).join(Organisation
        ).all()

        pconditions = self.conditions_for_organisation(organisation.id)

        self._aggregate_results = aggregate_results2

        self._summary = PublisherSummary(
            self._aggregate_results, conditions=pconditions)

class PackageSummaryCreator(SummaryCreator):
    def __init__(self, package, latest_runtime, aggregation_type):
        p = package

        self._aggregate_results = db.session.query(
            Indicator.id,
            Test.id,
            AggregateResult.results_data,
            AggregateResult.results_num,
            AggregateResult.result_hierarchy,
            AggregateResult.package_id
        ).filter(
            AggregateResult.package_id==p[0].id,
            AggregateResult.aggregateresulttype_id==aggregation_type
        ).group_by(
            AggregateResult.result_hierarchy,
            Test.id,
            AggregateResult.package_id,
            Indicator.id,
            AggregateResult.results_data,
            AggregateResult.results_num
            ).join(IndicatorTest
            ).join(Test
            ).join(AggregateResult
            ).all()

        self._summary = PublisherSummary(self._aggregate_results, {}).summary()

class PublisherIndicatorsSummaryCreator(SummaryCreator):
    def __init__(self, organisation, aggregation_type):
        aggregate_results = db.session.query(Indicator.id,
                                     Test.id,
                                     AggregateResult.results_data,
                                     AggregateResult.results_num,
                                     AggregateResult.result_hierarchy,
                                     AggregateResult.package_id
        ).filter(Organisation.organisation_code==organisation.organisation_code
        ).filter(AggregateResult.aggregateresulttype_id == aggregation_type
        ).filter(AggregateResult.organisation_id == organisation.id
        ).group_by(AggregateResult.result_hierarchy,
                   Test.id,
                   AggregateResult.package_id,
                   Indicator.id,
                   AggregateResult.results_data,
                   AggregateResult.results_num
        ).join(IndicatorTest
        ).join(Test
        ).join(AggregateResult
        ).join(Package
        ).join(OrganisationPackage
        ).join(Organisation
        ).order_by(Indicator.indicator_type,
                   Indicator.indicator_category_name,
                   Indicator.indicator_subcategory_name
        ).all()

        pconditions = self.conditions_for_organisation(organisation.id)

        self._aggregate_results = aggregate_results

        self._summary = PublisherIndicatorsSummary(
            aggregate_results,
            conditions=pconditions)

<|endoftext|>


<commit_before>import argparse
import logging
import os
import re

import simplejson as json

from Bio.PDB.PDBParser import PDBParser
from Bio.PDB import NeighborSearch

from config import PDB_RESIDUE_TYPES_BY_RESIDUE

LEVEL_MAP = {
    'A': 'atom',
    'R': 'residue',
    'C': 'chain'
}

# MATCHES PDB ACCESSIONS
# E.G. 1XKK
PDB_REGEX = r'([0-9]{1}[a-zA-Z]{3})'
# MATCHES PDB-BIOMOLID
# E.G. 1XKK-1
PDB_BIOMOL_REGEX = PDB_REGEX + r'-([0-9]+)'
# MATCHES RSCB STYLE ASSEMBLIES
# E.G. 1XKK.pdb1
PDB_RCSB_ASM_REGEX = PDB_REGEX + r'\.pdb([0-9]+)'

PDB_REGEX = re.compile(PDB_REGEX)
PDB_BIOMOL_REGEX = re.compile(PDB_BIOMOL_REGEX)
PDB_RCSB_ASM_REGEX = re.compile(PDB_RCSB_ASM_REGEX)

def cap(string, char=','):
    '''Pad the beginning and end of string with a character.'''
    return '{}{}{}'.format(char, string, char)

if __name__ == '__main__':

    # ARGUMENT PARSING
    parser = argparse.ArgumentParser(description='''
# Determine Interactions

Calculate all atom-atom, residue-residue, and chain-chain interactions
in PDB files.

This program assumes that the PDB file has already been cleaned and only takes
into account the first model.
''', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        'inputfile', type=str, help='Path to the PDB file to be analysed.')
    parser.add_argument('-i', '--interacting', type=float,
                        default=5.0, help='Distance cutoff for interactions.')
    parser.add_argument('-o', '--outputs', type=str,
                        default='cr', help='Granularity of output: string'
                        'including letters "c" for chain level, "r" for '
                        'residue level, "a" for atom level.')
    parser.add_argument('-v', '--verbose',
                        action='store_true', help='Be chatty.')

    args = parser.parse_args()

    # SET ARGS TO CONSTANTS
    INPUT_FILE = args.inputfile
    INPUT_FILE_SPLITEXT = os.path.splitext(INPUT_FILE)[0]
    INPUT_FILENAME = os.path.split(INPUT_FILE)[1]
    INTERACTION_THRESHOLD = args.interacting
    OUTPUTS = args.outputs.upper()

    # LOGGING
    if args.verbose:
        logging.basicConfig(
            level=logging.INFO,
            format='%(levelname)s//%(asctime)s.%(msecs).03d//%(message)s',
            datefmt='%H:%M:%S')
    else:
        logging.basicConfig(
            level=logging.WARN,
            format='%(levelname)s//%(asctime)s.%(msecs).03d//%(message)s',
            datefmt='%H:%M:%S')

    logging.info('Program begin.')

    # DETECT PDB ACCESSION FROM FILENAME IF POSSIBLE
    PDB = ''
    BIOMOL_ID = ''

    pdb_match = re.search(PDB_REGEX, INPUT_FILENAME)
    if pdb_match:
        PDB = pdb_match.group(1)

    pdb_biomol_match = re.search(PDB_BIOMOL_REGEX, INPUT_FILENAME)
    if pdb_biomol_match:
        PDB = pdb_biomol_match.group(1)
        BIOMOL_ID = pdb_biomol_match.group(2)

    pdb_rcsb_asm_match = re.search(PDB_RCSB_ASM_REGEX, INPUT_FILENAME)
    if pdb_rcsb_asm_match:
        PDB = pdb_rcsb_asm_match.group(1)
        BIOMOL_ID = pdb_rcsb_asm_match.group(2)

    # LOAD STRUCTURE
    structure = PDBParser().get_structure('structure', INPUT_FILE)
    structure_atoms = list(structure.get_atoms())

    logging.info('Loaded PDB structure (BioPython).')

    # CONSTRUCT KDTREE
    neighborsearch = NeighborSearch(structure_atoms)

    logging.info('Constructured NeighborSearch.')

    # GET INTERACTIONS
    logging.info('Calculating interactions...')
    for interaction_level in 'ARC':

        if interaction_level in OUTPUTS:

            logging.info('Calculating interactions for {}s...'.format(LEVEL_MAP[interaction_level]))

            pairs = neighborsearch.search_all(INTERACTION_THRESHOLD,
                                              level=interaction_level)

            logging.info('Search complete for {}s.'.format(LEVEL_MAP[interaction_level]))

            logging.info('Organising interactions for {}s...'.format(LEVEL_MAP[interaction_level]))
            interactions = {}

            for entities in pairs:

                entity1, entity2 = entities

                # NO SELFIES
                if (entity1 is entity2) or (entity1 == entity2):
                    continue

                id1 = entity1.get_full_id()
                id2 = entity2.get_full_id()

                if interaction_level == 'A':
                    res1 = entity1.get_parent()
                    res2 = entity2.get_parent()

                    res1 = res1.resname.strip()
                    res2 = res2.resname.strip()

                    entity1 = cap(','.join([id1[2], str(id1[3][1]) + id1[3][2].strip() + '`' + res1, entity1.name]))
                    entity2 = cap(','.join([id2[2], str(id2[3][1]) + id2[3][2].strip() + '`' + res2, entity2.name]))

                elif interaction_level == 'R':
                    entity1 = cap(','.join([id1[2], str(id1[3][1]) + id1[3][2].strip() + '`' + entity1.resname.strip()]))
                    entity2 = cap(','.join([id2[2], str(id2[3][1]) + id2[3][2].strip() + '`' + entity2.resname.strip()]))

                elif interaction_level == 'C':
                    entity1 = cap(entity1.id)
                    entity2 = cap(entity2.id)

                # ADD INTERACTING ENTITY TO LIST OF INTERACTORS
                if entity1 not in interactions:
                    interactions[entity1] = []

                if entity2 not in interactions:
                    interactions[entity2] = []

                if entity2 not in interactions[entity1]:
                    interactions[entity1].append(entity2)

                if entity1 not in interactions[entity2]:
                    interactions[entity2].append(entity1)

            for entity in interactions:
                interactions[entity] = sorted(interactions)

            logging.info('Organisation complete for {}s.'.format(LEVEL_MAP[interaction_level]))

            logging.info('Constructing JSON for {}s...'.format(LEVEL_MAP[interaction_level]))

            json_output = {
                'input': INPUT_FILE,
                'pdb': PDB,
                'biomol_id': BIOMOL_ID,
                'level': LEVEL_MAP[interaction_level],
                'interactions': interactions
            }

            # TYPE RESIDUES IF POSSIBLE
            if interaction_level in 'AR' and PDB_RESIDUE_TYPES_BY_RESIDUE:

                logging.info('Typing residues for {} output...'.format(LEVEL_MAP[interaction_level]))

                json_output['residue_types'] = {}

                for entity in json_output['interactions']:

                    resname = None

                    if interaction_level == 'A':
                        resname = entity.split(',')[-3].split('`')[1]
                    if interaction_level == 'R':
                        resname = entity.split(',')[-2].split('`')[1]

                    if resname:
                        json_output['residue_types'][entity] = PDB_RESIDUE_TYPES_BY_RESIDUE[resname]

            # WRITE OUT JSON OUTPUT
            logging.info('Writing JSON for {}s...'.format(LEVEL_MAP[interaction_level]))

            with open('.'.join([INPUT_FILE_SPLITEXT, LEVEL_MAP[interaction_level], 'interactions', 'json']), 'wb') as fo:
                json.dump(json_output, fo)

            logging.info('JSON output written for {}s.'.format(LEVEL_MAP[interaction_level]))

    # FINISH UP
    logging.info('Program end.')
<commit_msg>Typo.<commit_after>import argparse
import logging
import os
import re

import simplejson as json

from Bio.PDB.PDBParser import PDBParser
from Bio.PDB import NeighborSearch

from config import PDB_RESIDUE_TYPES_BY_RESIDUE

LEVEL_MAP = {
    'A': 'atom',
    'R': 'residue',
    'C': 'chain'
}

# MATCHES PDB ACCESSIONS
# E.G. 1XKK
PDB_REGEX = r'([0-9]{1}[a-zA-Z]{3})'
# MATCHES PDB-BIOMOLID
# E.G. 1XKK-1
PDB_BIOMOL_REGEX = PDB_REGEX + r'-([0-9]+)'
# MATCHES RSCB STYLE ASSEMBLIES
# E.G. 1XKK.pdb1
PDB_RCSB_ASM_REGEX = PDB_REGEX + r'\.pdb([0-9]+)'

PDB_REGEX = re.compile(PDB_REGEX)
PDB_BIOMOL_REGEX = re.compile(PDB_BIOMOL_REGEX)
PDB_RCSB_ASM_REGEX = re.compile(PDB_RCSB_ASM_REGEX)

def cap(string, char=','):
    '''Pad the beginning and end of string with a character.'''
    return '{}{}{}'.format(char, string, char)

if __name__ == '__main__':

    # ARGUMENT PARSING
    parser = argparse.ArgumentParser(description='''
# Determine Interactions

Calculate all atom-atom, residue-residue, and chain-chain interactions
in PDB files.

This program assumes that the PDB file has already been cleaned and only takes
into account the first model.
''', formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument(
        'inputfile', type=str, help='Path to the PDB file to be analysed.')
    parser.add_argument('-i', '--interacting', type=float,
                        default=5.0, help='Distance cutoff for interactions.')
    parser.add_argument('-o', '--outputs', type=str,
                        default='cr', help='Granularity of output: string '
                        'including letters "c" for chain level, "r" for '
                        'residue level, "a" for atom level.')
    parser.add_argument('-v', '--verbose',
                        action='store_true', help='Be chatty.')

    args = parser.parse_args()

    # SET ARGS TO CONSTANTS
    INPUT_FILE = args.inputfile
    INPUT_FILE_SPLITEXT = os.path.splitext(INPUT_FILE)[0]
    INPUT_FILENAME = os.path.split(INPUT_FILE)[1]
    INTERACTION_THRESHOLD = args.interacting
    OUTPUTS = args.outputs.upper()

    # LOGGING
    if args.verbose:
        logging.basicConfig(
            level=logging.INFO,
            format='%(levelname)s//%(asctime)s.%(msecs).03d//%(message)s',
            datefmt='%H:%M:%S')
    else:
        logging.basicConfig(
            level=logging.WARN,
            format='%(levelname)s//%(asctime)s.%(msecs).03d//%(message)s',
            datefmt='%H:%M:%S')

    logging.info('Program begin.')

    # DETECT PDB ACCESSION FROM FILENAME IF POSSIBLE
    PDB = ''
    BIOMOL_ID = ''

    pdb_match = re.search(PDB_REGEX, INPUT_FILENAME)
    if pdb_match:
        PDB = pdb_match.group(1)

    pdb_biomol_match = re.search(PDB_BIOMOL_REGEX, INPUT_FILENAME)
    if pdb_biomol_match:
        PDB = pdb_biomol_match.group(1)
        BIOMOL_ID = pdb_biomol_match.group(2)

    pdb_rcsb_asm_match = re.search(PDB_RCSB_ASM_REGEX, INPUT_FILENAME)
    if pdb_rcsb_asm_match:
        PDB = pdb_rcsb_asm_match.group(1)
        BIOMOL_ID = pdb_rcsb_asm_match.group(2)

    # LOAD STRUCTURE
    structure = PDBParser().get_structure('structure', INPUT_FILE)
    structure_atoms = list(structure.get_atoms())

    logging.info('Loaded PDB structure (BioPython).')

    # CONSTRUCT KDTREE
    neighborsearch = NeighborSearch(structure_atoms)

    logging.info('Constructured NeighborSearch.')

    # GET INTERACTIONS
    logging.info('Calculating interactions...')
    for interaction_level in 'ARC':

        if interaction_level in OUTPUTS:

            logging.info('Calculating interactions for {}s...'.format(LEVEL_MAP[interaction_level]))

            pairs = neighborsearch.search_all(INTERACTION_THRESHOLD,
                                              level=interaction_level)

            logging.info('Search complete for {}s.'.format(LEVEL_MAP[interaction_level]))

            logging.info('Organising interactions for {}s...'.format(LEVEL_MAP[interaction_level]))
            interactions = {}

            for entities in pairs:

                entity1, entity2 = entities

                # NO SELFIES
                if (entity1 is entity2) or (entity1 == entity2):
                    continue

                id1 = entity1.get_full_id()
                id2 = entity2.get_full_id()

                if interaction_level == 'A':
                    res1 = entity1.get_parent()
                    res2 = entity2.get_parent()

                    res1 = res1.resname.strip()
                    res2 = res2.resname.strip()

                    entity1 = cap(','.join([id1[2], str(id1[3][1]) + id1[3][2].strip() + '`' + res1, entity1.name]))
                    entity2 = cap(','.join([id2[2], str(id2[3][1]) + id2[3][2].strip() + '`' + res2, entity2.name]))

                elif interaction_level == 'R':
                    entity1 = cap(','.join([id1[2], str(id1[3][1]) + id1[3][2].strip() + '`' + entity1.resname.strip()]))
                    entity2 = cap(','.join([id2[2], str(id2[3][1]) + id2[3][2].strip() + '`' + entity2.resname.strip()]))

                elif interaction_level == 'C':
                    entity1 = cap(entity1.id)
                    entity2 = cap(entity2.id)

                # ADD INTERACTING ENTITY TO LIST OF INTERACTORS
                if entity1 not in interactions:
                    interactions[entity1] = []

                if entity2 not in interactions:
                    interactions[entity2] = []

                if entity2 not in interactions[entity1]:
                    interactions[entity1].append(entity2)

                if entity1 not in interactions[entity2]:
                    interactions[entity2].append(entity1)

            for entity in interactions:
                interactions[entity] = sorted(interactions)

            logging.info('Organisation complete for {}s.'.format(LEVEL_MAP[interaction_level]))

            logging.info('Constructing JSON for {}s...'.format(LEVEL_MAP[interaction_level]))

            json_output = {
                'input': INPUT_FILE,
                'pdb': PDB,
                'biomol_id': BIOMOL_ID,
                'level': LEVEL_MAP[interaction_level],
                'interactions': interactions
            }

            # TYPE RESIDUES IF POSSIBLE
            if interaction_level in 'AR' and PDB_RESIDUE_TYPES_BY_RESIDUE:

                logging.info('Typing residues for {} output...'.format(LEVEL_MAP[interaction_level]))

                json_output['residue_types'] = {}

                for entity in json_output['interactions']:

                    resname = None

                    if interaction_level == 'A':
                        resname = entity.split(',')[-3].split('`')[1]
                    if interaction_level == 'R':
                        resname = entity.split(',')[-2].split('`')[1]

                    if resname:
                        json_output['residue_types'][entity] = PDB_RESIDUE_TYPES_BY_RESIDUE[resname]

            # WRITE OUT JSON OUTPUT
            logging.info('Writing JSON for {}s...'.format(LEVEL_MAP[interaction_level]))

            with open('.'.join([INPUT_FILE_SPLITEXT, LEVEL_MAP[interaction_level], 'interactions', 'json']), 'wb') as fo:
                json.dump(json_output, fo)

            logging.info('JSON output written for {}s.'.format(LEVEL_MAP[interaction_level]))

    # FINISH UP
    logging.info('Program end.')
<|endoftext|>


<commit_before>#!/usr/bin/python3.4
# -*-coding:Utf-8 -*
'''module to manage antialiasing settings'''
import xml.etree.ElementTree as xmlMod

class OSA:
	'''class to manage antialiasing settings'''

	def __init__(self, xml= None):
		'''initialize OSA settings with default value or values extracted from an xml object'''
		if xml is None:
			self.defaultInit()
		else:
			self.fromXml(xml)

	def defaultInit(self):
		'''initialize antialiasing settings with default value'''
		self.enabled = True
		self.samples = 8
		self.fullSample = False

	def fromXml(self, xml):
		'''initialize antialiasing settings with values extracted from an xml object'''
		boolDict = { 'True':True, 'False':False }
		self.enabled = boolDict[xml.get('enabled')]
		self.samples = int(xml.get('samples'))
		self.fullSample = boolDict[xml.get('fullSample')]

	def toXml(self):
		'''export antialiasing settings into xml syntaxed string'''
		txt = '<OSA enabled="'+str(self.enabled)+'" samples="'\
				+str(self.samples)+'" fullSample="'+str(self.fullSample)+'"/>\n'
		return txt

	def menu(self, log):
		'''menu to explore and edit antialiasing settings'''
		change = False
		log.menuIn('OSA Settings')

		while True:

			log.print()

			self.print()

			print('''

(Antialiasing setting only work with blender internal engine and not with Cycles)

        Menu :
1- Switch To Enable Or Disabled OSA
2- Edit OSA Samples
0- Quit

''')

			choice = input('Action?').strip().lower()

			if choice in ['0', 'q', 'quit', 'cancel']:
				log.menuOut()
				return change
			elif choice == '1':
				change = True
				self.enabled = not self.enabled
			elif choice == '2':
				change = (self.editSamples() or change)
			else:
				log.error('Unvalid menu choice', False)

	def print(self):
		'''a method to print preset'''
		enable = { True:'enabled', False:'disabled' }
		print('Antialiasing :          '\
					+enable[self.enabled] )

		if self.enabled:
			print('OSA samples :           '+str(self.samples))
			print('Full Sample :           '+enable[self.fullSample])

	def editSamples(self):
		'''menu to edit OSA sampes setting'''
		log.menuIn('Edit Samples Setting')
		options = [ 5, 8, 11, 16 ]

		while True:
			log.print()

			print('\n\nCurrent settings : '+str(self.samples)+'samples per pixel.\n\n       available settings choice :\n0- Quit whithout change' )

			for opt in options:
				print(str(opt)+'- '+str(opt)+' samples per pixel')
			choice = input().strip().lower()

			try:
				choice = int(choice)
			except ValueError:
				log.error('expect integer value.')
				continue

			if choice == 0:
				return False
			if choice in options
				self.samples = choice
				return True
			log.error('unexpected «'+str(choice)+'» integer value. expected 0, 5, 8, 11 or 16.')

<commit_msg>implement OSA Full Sample option switcher<commit_after>#!/usr/bin/python3.4
# -*-coding:Utf-8 -*
'''module to manage antialiasing settings'''
import xml.etree.ElementTree as xmlMod

class OSA:
	'''class to manage antialiasing settings'''

	def __init__(self, xml= None):
		'''initialize OSA settings with default value or values extracted from an xml object'''
		if xml is None:
			self.defaultInit()
		else:
			self.fromXml(xml)

	def defaultInit(self):
		'''initialize antialiasing settings with default value'''
		self.enabled = True
		self.samples = 8
		self.fullSample = False

	def fromXml(self, xml):
		'''initialize antialiasing settings with values extracted from an xml object'''
		boolDict = { 'True':True, 'False':False }
		self.enabled = boolDict[xml.get('enabled')]
		self.samples = int(xml.get('samples'))
		self.fullSample = boolDict[xml.get('fullSample')]

	def toXml(self):
		'''export antialiasing settings into xml syntaxed string'''
		txt = '<OSA enabled="'+str(self.enabled)+'" samples="'\
				+str(self.samples)+'" fullSample="'+str(self.fullSample)+'"/>\n'
		return txt

	def menu(self, log):
		'''menu to explore and edit antialiasing settings'''
		change = False
		log.menuIn('OSA Settings')

		while True:

			log.print()

			self.print()

			print('''

(Antialiasing setting only work with blender internal engine and not with Cycles)

        Menu :
1- Switch To Enable Or Disabled OSA
2- Edit OSA Samples
3- Switch To Enable Or Disabled OSA Full Sample Option
0- Quit

''')

			choice = input('Action?').strip().lower()

			if choice in ['0', 'q', 'quit', 'cancel']:
				log.menuOut()
				return change
			elif choice == '1':
				change = True
				self.enabled = not self.enabled
			elif choice == '2':
				change = (self.editSamples() or change)
			elif choice == '3':
				change = True
				self.fullSample = not self.fullSample
			else:
				log.error('Unvalid menu choice', False)

	def print(self):
		'''a method to print preset'''
		enable = { True:'enabled', False:'disabled' }
		print('Antialiasing :          '\
					+enable[self.enabled] )

		if self.enabled:
			print('OSA samples :           '+str(self.samples))
			print('Full Sample :           '+enable[self.fullSample])

	def editSamples(self):
		'''menu to edit OSA sampes setting'''
		log.menuIn('Edit Samples Setting')
		options = [ 5, 8, 11, 16 ]

		while True:
			log.print()

			print('\n\nCurrent settings : '+str(self.samples)+'samples per pixel.\n\n       available settings choice :\n0- Quit whithout change' )

			for opt in options:
				print(str(opt)+'- '+str(opt)+' samples per pixel')
			choice = input().strip().lower()

			try:
				choice = int(choice)
			except ValueError:
				log.error('expect integer value.')
				continue

			if choice == 0:
				return False
			if choice in options
				self.samples = choice
				return True
			log.error('unexpected «'+str(choice)+'» integer value. expected 0, 5, 8, 11 or 16.')

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Generated by Django 1.10.7 on 2017-04-22 17:18
from __future__ import unicode_literals

from django.db import migrations

from corehq.sql_db.operations import HqRunSQL

class Migration(migrations.Migration):

    dependencies = [
        ('accounting', '0005_automatic_downgrade_adjustment_method'),
    ]

    operations = [
        HqRunSQL(
            """
            CREATE UNIQUE INDEX accounting_subscription_active_subscriber
            ON accounting_subscription(subscriber_id) WHERE (is_active = TRUE and is_hidden_to_ops = FALSE);
            """,
            reverse_sql=
            """
            DROP INDEX accounting_subscription_active_subscriber;
            """,
        )
    ]
<commit_msg>add if not exists to create<commit_after># -*- coding: utf-8 -*-
# Generated by Django 1.10.7 on 2017-04-22 17:18
from __future__ import unicode_literals

from django.db import migrations

from corehq.sql_db.operations import HqRunSQL
from corehq.util.django_migrations import add_if_not_exists

class Migration(migrations.Migration):

    dependencies = [
        ('accounting', '0005_automatic_downgrade_adjustment_method'),
    ]

    operations = [
        HqRunSQL(
            add_if_not_exists(
                """
                CREATE UNIQUE INDEX accounting_subscription_active_subscriber
                ON accounting_subscription(subscriber_id) WHERE (is_active = TRUE and is_hidden_to_ops = FALSE)
                """
            ),
            reverse_sql=
            """
            DROP INDEX accounting_subscription_active_subscriber;
            """,
        )
    ]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import re
import time

from django.core.urlresolvers import reverse
from django.http import QueryDict
from django.utils.translation import ugettext as _

from desktop.lib.exceptions_renderable import PopupException
from desktop.models import Document2

from notebook.connectors.base import Api, QueryError

LOG = logging.getLogger(__name__)

try:
  from oozie.models2 import Workflow, WorkflowBuilder
  from oozie.views.api import get_log as get_workflow_logs
  from oozie.views.dashboard import check_job_access_permission, check_job_edition_permission
  from oozie.views.editor2 import _submit_workflow
except Exception, e:
  LOG.exception('Oozie application is not enabled: %s' % e)

class OozieApi(Api):

  LOG_START_PATTERN = '(>>> Invoking Main class now >>>.+)'
  LOG_END_PATTERN = '<<< Invocation of Main class completed <<<'
  RESULTS_PATTERN = "(?P<results>>>> Invoking Beeline command line now >>>.+<<< Invocation of Beeline command completed <<<)"
  RESULTS_PATTERN_GENERIC = "(?P<results>>>> Invoking Main class now >>>.+<<< Invocation of Main class completed <<<)"

  def __init__(self, *args, **kwargs):
    Api.__init__(self, *args, **kwargs)

    self.fs = self.request.fs
    self.jt = self.request.jt

  def execute(self, notebook, snippet):
    # Get document from notebook
    if not notebook.get('uuid', ''):
      raise PopupException(_('Notebook is missing a uuid, please save the notebook before executing as a batch job.'))

    notebook_doc = Document2.objects.get_by_uuid(user=self.user, uuid=notebook['uuid'], perm_type='read')

    if notebook_doc.type == 'notebook':
      pass
    else:
      # Create a managed workflow from the notebook doc
      workflow_doc = WorkflowBuilder().create_workflow(document=notebook_doc, user=self.user, managed=True, name=_("Batch job for %s") % (notebook_doc.name or notebook_doc.type))
      workflow = Workflow(document=workflow_doc, user=self.user)

    # Submit workflow
    job_id = _submit_workflow(user=self.user, fs=self.fs, jt=self.jt, workflow=workflow, mapping=None)

    return {
      'id': job_id,
      'has_result_set': True,
    }

  def check_status(self, notebook, snippet):
    response = {'status': 'running'}

    job_id = snippet['result']['handle']['id']
    oozie_job = check_job_access_permission(self.request, job_id)

    if oozie_job.is_running():
      return response
    elif oozie_job.status in ('KILLED', 'FAILED'):
      raise QueryError(_('Job was %s') % oozie_job.status)
    else:
      # Check if job results are actually available, since YARN takes a while to move logs to JHS,
      log_output = self.get_log(notebook, snippet)
      if log_output:
        results = self._get_results(log_output, snippet['type'])
        if results:
          response['status'] = 'available'
      else:
        response['status'] = 'failed'

    return response

  def fetch_result(self, notebook, snippet, rows, start_over):
    log_output = self.get_log(notebook, snippet)
    results = self._get_results(log_output, snippet['type'])

    return {
        'data':  [[line] for line in results.split('\n')],  # hdfs_link()
        'meta': [{'name': 'Header', 'type': 'STRING_TYPE', 'comment': ''}],
        'type': 'table',
        'has_more': False,
    }

  def cancel(self, notebook, snippet):
    job_id = snippet['result']['handle']['id']

    job = check_job_access_permission(self, job_id)
    oozie_job = check_job_edition_permission(job, self.user)

    oozie_job.kill()

    return {'status': 0}

  def get_log(self, notebook, snippet, startFrom=0, size=None):
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    return self._get_log_output(oozie_job)

  def progress(self, snippet, logs):
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    return oozie_job.get_progress(),

  def get_jobs(self, notebook, snippet, logs):
    jobs = []
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    actions = oozie_job.get_working_actions()
    for action in actions:
      if action.externalId is not None:
        jobs.append({
          'name': action.externalId,
          'url': reverse('jobbrowser.views.single_job', kwargs={'job': action.externalId}),
          'started': action.startTime is not None,
          'finished': action.endTime is not None
        })
    return jobs

  def close_statement(self, snippet):
    pass

  def close_session(self, session):
    pass

  def _get_log_output(self, oozie_workflow):
    log_output = ''
    q = QueryDict(self.request.GET, mutable=True)
    q['format'] = 'python'  # Hack for triggering the good section in single_task_attempt_logs
    self.request.GET = q

    attempts = 0
    max_attempts = 10
    logs_found = False
    while not logs_found and attempts < max_attempts:
      logs, workflow_actions, is_really_done = get_workflow_logs(self.request, oozie_workflow, make_links=False,
                                                                 log_start_pattern=self.LOG_START_PATTERN,
                                                                 log_end_pattern=self.LOG_END_PATTERN)
      if len(logs) > 0:
        log_output = logs.values()[0]
        if log_output.startswith('Unable to locate'):
          LOG.debug('Failed to get job attempt logs, possibly due to YARN archiving job to JHS. Will sleep and try again.')
          time.sleep(2.0)
        else:
          logs_found = True

    return log_output

  def _get_results(self, log_output, action_type):
    results = ''

    pattern = self.RESULTS_PATTERN if action_type == 'hive' else self.RESULTS_PATTERN_GENERIC

    re_results = re.compile(pattern, re.M | re.DOTALL)
    if re_results.search(log_output):
      results = re.search(re_results, log_output).group('results').strip()
    return results
<commit_msg>HUE-4554 [editor] Status checking can go in infinite loop in batch mode<commit_after>#!/usr/bin/env python
# Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import re
import time

from django.core.urlresolvers import reverse
from django.http import QueryDict
from django.utils.translation import ugettext as _

from desktop.lib.exceptions_renderable import PopupException
from desktop.models import Document2

from notebook.connectors.base import Api, QueryError

LOG = logging.getLogger(__name__)

try:
  from oozie.models2 import Workflow, WorkflowBuilder
  from oozie.views.api import get_log as get_workflow_logs
  from oozie.views.dashboard import check_job_access_permission, check_job_edition_permission
  from oozie.views.editor2 import _submit_workflow
except Exception, e:
  LOG.exception('Oozie application is not enabled: %s' % e)

class OozieApi(Api):

  LOG_START_PATTERN = '(>>> Invoking Main class now >>>.+)'
  LOG_END_PATTERN = '<<< Invocation of Main class completed <<<'
  RESULTS_PATTERN = "(?P<results>>>> Invoking Beeline command line now >>>.+<<< Invocation of Beeline command completed <<<)"
  RESULTS_PATTERN_GENERIC = "(?P<results>>>> Invoking Main class now >>>.+<<< Invocation of Main class completed <<<)"

  def __init__(self, *args, **kwargs):
    Api.__init__(self, *args, **kwargs)

    self.fs = self.request.fs
    self.jt = self.request.jt

  def execute(self, notebook, snippet):
    # Get document from notebook
    if not notebook.get('uuid', ''):
      raise PopupException(_('Notebook is missing a uuid, please save the notebook before executing as a batch job.'))

    notebook_doc = Document2.objects.get_by_uuid(user=self.user, uuid=notebook['uuid'], perm_type='read')

    if notebook_doc.type == 'notebook':
      pass
    else:
      # Create a managed workflow from the notebook doc
      workflow_doc = WorkflowBuilder().create_workflow(document=notebook_doc, user=self.user, managed=True, name=_("Batch job for %s") % (notebook_doc.name or notebook_doc.type))
      workflow = Workflow(document=workflow_doc, user=self.user)

    # Submit workflow
    job_id = _submit_workflow(user=self.user, fs=self.fs, jt=self.jt, workflow=workflow, mapping=None)

    return {
      'id': job_id,
      'has_result_set': True,
    }

  def check_status(self, notebook, snippet):
    response = {'status': 'running'}

    job_id = snippet['result']['handle']['id']
    oozie_job = check_job_access_permission(self.request, job_id)

    if oozie_job.is_running():
      return response
    elif oozie_job.status in ('KILLED', 'FAILED'):
      raise QueryError(_('Job was %s') % oozie_job.status)
    else:
      # Check if job results are actually available, since YARN takes a while to move logs to JHS,
      log_output = self.get_log(notebook, snippet)
      if log_output:
        results = self._get_results(log_output, snippet['type'])
        if results:
          response['status'] = 'available'
      else:
        response['status'] = 'failed'

    return response

  def fetch_result(self, notebook, snippet, rows, start_over):
    log_output = self.get_log(notebook, snippet)
    results = self._get_results(log_output, snippet['type'])

    return {
        'data':  [[line] for line in results.split('\n')],  # hdfs_link()
        'meta': [{'name': 'Header', 'type': 'STRING_TYPE', 'comment': ''}],
        'type': 'table',
        'has_more': False,
    }

  def cancel(self, notebook, snippet):
    job_id = snippet['result']['handle']['id']

    job = check_job_access_permission(self, job_id)
    oozie_job = check_job_edition_permission(job, self.user)

    oozie_job.kill()

    return {'status': 0}

  def get_log(self, notebook, snippet, startFrom=0, size=None):
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    return self._get_log_output(oozie_job)

  def progress(self, snippet, logs):
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    return oozie_job.get_progress(),

  def get_jobs(self, notebook, snippet, logs):
    jobs = []
    job_id = snippet['result']['handle']['id']

    oozie_job = check_job_access_permission(self.request, job_id)
    actions = oozie_job.get_working_actions()
    for action in actions:
      if action.externalId is not None:
        jobs.append({
          'name': action.externalId,
          'url': reverse('jobbrowser.views.single_job', kwargs={'job': action.externalId}),
          'started': action.startTime is not None,
          'finished': action.endTime is not None
        })
    return jobs

  def close_statement(self, snippet):
    pass

  def close_session(self, session):
    pass

  def _get_log_output(self, oozie_workflow):
    log_output = ''
    q = QueryDict(self.request.GET, mutable=True)
    q['format'] = 'python'  # Hack for triggering the good section in single_task_attempt_logs
    self.request.GET = q

    attempts = 0
    max_attempts = 10
    logs_found = False
    while not logs_found and attempts < max_attempts:
      logs, workflow_actions, is_really_done = get_workflow_logs(self.request, oozie_workflow, make_links=False,
                                                                 log_start_pattern=self.LOG_START_PATTERN,
                                                                 log_end_pattern=self.LOG_END_PATTERN)
      if logs:
        log_output = logs.values()[0]
        if log_output.startswith('Unable to locate'):
          LOG.debug('Failed to get job attempt logs, possibly due to YARN archiving job to JHS. Will sleep and try again.')
          time.sleep(2.0)
        else:
          logs_found = True

      attempts += 1
    return log_output

  def _get_results(self, log_output, action_type):
    results = ''

    pattern = self.RESULTS_PATTERN if action_type == 'hive' else self.RESULTS_PATTERN_GENERIC

    re_results = re.compile(pattern, re.M | re.DOTALL)
    if re_results.search(log_output):
      results = re.search(re_results, log_output).group('results').strip()
    return results
<|endoftext|>