<commit_before>import logging
import fingerprints
from pprint import pprint, pformat  # noqa
from banal import ensure_list, first
from followthemoney import model
from followthemoney.types import registry
from elasticsearch.helpers import scan

from aleph.core import es, cache
from aleph.model import Entity
from aleph.index.indexes import entities_write_index, entities_read_index
from aleph.index.util import unpack_result, delete_safe
from aleph.index.util import authz_query, bulk_actions
from aleph.index.util import MAX_PAGE, NUMERIC_TYPES
from aleph.index.util import MAX_REQUEST_TIMEOUT, MAX_TIMEOUT

log = logging.getLogger(__name__)
PROXY_INCLUDES = [
    "schema",
    "properties",
    "collection_id",
    "role_id",
    "mutable",
    "created_at",
    "updated_at",
]

def _source_spec(includes, excludes):
    includes = ensure_list(includes)
    excludes = ensure_list(excludes)
    return {"includes": includes, "excludes": excludes}

def _entities_query(filters, authz, collection_id, schemata):
    filters = filters or []
    if authz is not None:
        filters.append(authz_query(authz))
    if collection_id is not None:
        filters.append({"term": {"collection_id": collection_id}})
    if ensure_list(schemata):
        filters.append({"terms": {"schemata": ensure_list(schemata)}})
    return {"bool": {"filter": filters}}

def get_field_type(field):
    field = field.split(".")[-1]
    if field in registry.groups:
        return registry.groups[field]
    for prop in model.properties:
        if prop.name == field:
            return prop.type
    return registry.string

def iter_entities(
    authz=None,
    collection_id=None,
    schemata=None,
    includes=PROXY_INCLUDES,
    excludes=None,
    filters=None,
    sort=None,
):
    """Scan all entities matching the given criteria."""
    query = {
        "query": _entities_query(filters, authz, collection_id, schemata),
        "_source": _source_spec(includes, excludes),
    }
    preserve_order = False
    if sort is not None:
        query["sort"] = ensure_list(sort)
        preserve_order = True
    index = entities_read_index(schema=schemata)
    for res in scan(
        es,
        index=index,
        query=query,
        timeout=MAX_TIMEOUT,
        request_timeout=MAX_REQUEST_TIMEOUT,
        preserve_order=preserve_order,
    ):
        entity = unpack_result(res)
        if entity is not None:
            yield entity

def iter_proxies(**kw):
    for data in iter_entities(**kw):
        schema = model.get(data.get("schema"))
        if schema is None:
            continue
        yield model.get_proxy(data)

def iter_adjacent(entity):
    """Used for recursively deleting entities and their linked associations."""
    query = {"term": {"entities": entity.get("id")}}
    yield from iter_entities(
        includes=["collection_id"],
        collection_id=entity.get("collection_id"),
        filters=[query],
    )

def entities_by_ids(
    ids, schemata=None, cached=False, includes=PROXY_INCLUDES, excludes=None
):
    """Iterate over unpacked entities based on a search for the given
    entity IDs."""
    ids = ensure_list(ids)
    if not len(ids):
        return
    cached = cached and excludes is None and includes == PROXY_INCLUDES
    entities = {}
    if cached:
        keys = [cache.object_key(Entity, i) for i in ids]
        for _, entity in cache.get_many_complex(keys):
            if entity is not None:
                entities[entity.get("id")] = entity

    missing = [i for i in ids if entities.get(id) is None]
    index = entities_read_index(schema=schemata)
    query = {
        "query": {"ids": {"values": missing}},
        "_source": _source_spec(includes, excludes),
        "size": MAX_PAGE,
    }
    result = es.search(index=index, body=query)
    for doc in result.get("hits", {}).get("hits", []):
        entity = unpack_result(doc)
        if entity is not None:
            entity_id = entity.get("id")
            entities[entity_id] = entity
            if cached:
                key = cache.object_key(Entity, entity_id)
                cache.set_complex(key, entity, expires=60 * 60 * 2)

    for i in ids:
        entity = entities.get(i)
        if entity is not None:
            yield entity

def get_entity(entity_id, **kwargs):
    """Fetch an entity from the index."""
    for entity in entities_by_ids(entity_id, cached=True, **kwargs):
        return entity

def index_entity(entity, sync=False):
    """Index an entity."""
    return index_proxy(entity.collection, entity.to_proxy(), sync=sync)

def index_proxy(collection, proxy, sync=False):
    delete_entity(proxy.id, exclude=proxy.schema, sync=False)
    return index_bulk(collection, [proxy], sync=sync)

def index_bulk(collection, entities, sync=False):
    """Index a set of entities."""
    entities = (format_proxy(p, collection) for p in entities)
    bulk_actions(entities, sync=sync)

def _numeric_values(type_, values):
    values = [type_.to_number(v) for v in ensure_list(values)]
    return [v for v in values if v is not None]

def format_proxy(proxy, collection):
    """Apply final denormalisations to the index."""
    data = proxy.to_full_dict()
    data["schemata"] = list(proxy.schema.names)
    data["caption"] = proxy.caption

    names = data.get("names", [])
    fps = set([fingerprints.generate(name) for name in names])
    fps.update(names)
    data["fingerprints"] = [fp for fp in fps if fp is not None]

    # Slight hack: a magic property in followthemoney that gets taken out
    # of the properties and added straight to the index text.
    properties = data.get("properties")
    data["text"] = properties.pop("indexText", [])

    # integer casting
    numeric = {}
    for prop in proxy.iterprops():
        if prop.type in NUMERIC_TYPES:
            values = proxy.get(prop)
            numeric[prop.name] = _numeric_values(prop.type, values)
    # also cast group field for dates
    numeric["dates"] = _numeric_values(registry.date, data.get("dates"))
    data["numeric"] = numeric

    # Context data - from aleph system, not followthemoney.
    # FIXME: Can there ever really be multiple role_ids?
    data["role_id"] = first(data.get("role_id"))
    data["mutable"] = max(ensure_list(data.get("mutable")), default=False)
    data["origin"] = ensure_list(data.get("origin"))
    created_at = data.get("created_at")
    if created_at:
        data["updated_at"] = data.get("updated_at", created_at)
    data["collection_id"] = collection.id
    # log.info("%s", pformat(data))
    entity_id = data.pop("id")
    return {
        "_id": entity_id,
        "_index": entities_write_index(data.get("schema")),
        "_source": data,
    }

def delete_entity(entity_id, exclude=None, sync=False):
    """Delete an entity from the index."""
    if exclude is not None:
        exclude = entities_write_index(exclude)
    for entity in entities_by_ids(entity_id, excludes="*"):
        index = entity.get("_index")
        if index == exclude:
            continue
        delete_safe(index, entity_id)
<commit_msg>Pass proxy object<commit_after>import logging
import fingerprints
from pprint import pprint, pformat  # noqa
from banal import ensure_list, first
from followthemoney import model
from followthemoney.types import registry
from elasticsearch.helpers import scan

from aleph.core import es, cache
from aleph.model import Entity
from aleph.index.indexes import entities_write_index, entities_read_index
from aleph.index.util import unpack_result, delete_safe
from aleph.index.util import authz_query, bulk_actions
from aleph.index.util import MAX_PAGE, NUMERIC_TYPES
from aleph.index.util import MAX_REQUEST_TIMEOUT, MAX_TIMEOUT

log = logging.getLogger(__name__)
PROXY_INCLUDES = [
    "schema",
    "properties",
    "collection_id",
    "role_id",
    "mutable",
    "created_at",
    "updated_at",
]

def _source_spec(includes, excludes):
    includes = ensure_list(includes)
    excludes = ensure_list(excludes)
    return {"includes": includes, "excludes": excludes}

def _entities_query(filters, authz, collection_id, schemata):
    filters = filters or []
    if authz is not None:
        filters.append(authz_query(authz))
    if collection_id is not None:
        filters.append({"term": {"collection_id": collection_id}})
    if ensure_list(schemata):
        filters.append({"terms": {"schemata": ensure_list(schemata)}})
    return {"bool": {"filter": filters}}

def get_field_type(field):
    field = field.split(".")[-1]
    if field in registry.groups:
        return registry.groups[field]
    for prop in model.properties:
        if prop.name == field:
            return prop.type
    return registry.string

def iter_entities(
    authz=None,
    collection_id=None,
    schemata=None,
    includes=PROXY_INCLUDES,
    excludes=None,
    filters=None,
    sort=None,
):
    """Scan all entities matching the given criteria."""
    query = {
        "query": _entities_query(filters, authz, collection_id, schemata),
        "_source": _source_spec(includes, excludes),
    }
    preserve_order = False
    if sort is not None:
        query["sort"] = ensure_list(sort)
        preserve_order = True
    index = entities_read_index(schema=schemata)
    for res in scan(
        es,
        index=index,
        query=query,
        timeout=MAX_TIMEOUT,
        request_timeout=MAX_REQUEST_TIMEOUT,
        preserve_order=preserve_order,
    ):
        entity = unpack_result(res)
        if entity is not None:
            yield entity

def iter_proxies(**kw):
    for data in iter_entities(**kw):
        schema = model.get(data.get("schema"))
        if schema is None:
            continue
        yield model.get_proxy(data)

def iter_adjacent(entity):
    """Used for recursively deleting entities and their linked associations."""
    query = {"term": {"entities": entity.get("id")}}
    yield from iter_entities(
        includes=["collection_id"],
        collection_id=entity.get("collection_id"),
        filters=[query],
    )

def entities_by_ids(
    ids, schemata=None, cached=False, includes=PROXY_INCLUDES, excludes=None
):
    """Iterate over unpacked entities based on a search for the given
    entity IDs."""
    ids = ensure_list(ids)
    if not len(ids):
        return
    cached = cached and excludes is None and includes == PROXY_INCLUDES
    entities = {}
    if cached:
        keys = [cache.object_key(Entity, i) for i in ids]
        for _, entity in cache.get_many_complex(keys):
            if entity is not None:
                entities[entity.get("id")] = entity

    missing = [i for i in ids if entities.get(id) is None]
    index = entities_read_index(schema=schemata)
    query = {
        "query": {"ids": {"values": missing}},
        "_source": _source_spec(includes, excludes),
        "size": MAX_PAGE,
    }
    result = es.search(index=index, body=query)
    for doc in result.get("hits", {}).get("hits", []):
        entity = unpack_result(doc)
        if entity is not None:
            entity_id = entity.get("id")
            entities[entity_id] = entity
            if cached:
                key = cache.object_key(Entity, entity_id)
                cache.set_complex(key, entity, expires=60 * 60 * 2)

    for i in ids:
        entity = entities.get(i)
        if entity is not None:
            yield entity

def get_entity(entity_id, **kwargs):
    """Fetch an entity from the index."""
    for entity in entities_by_ids(entity_id, cached=True, **kwargs):
        return entity

def index_entity(entity, sync=False):
    """Index an entity."""
    return index_proxy(entity.collection, entity.to_proxy(), sync=sync)

def index_proxy(collection, proxy, sync=False):
    delete_entity(proxy.id, exclude=proxy.schema, sync=False)
    return index_bulk(collection, [proxy], sync=sync)

def index_bulk(collection, entities, sync=False):
    """Index a set of entities."""
    entities = (format_proxy(p, collection) for p in entities)
    bulk_actions(entities, sync=sync)

def _numeric_values(type_, values):
    values = [type_.to_number(v) for v in ensure_list(values)]
    return [v for v in values if v is not None]

def format_proxy(proxy, collection):
    """Apply final denormalisations to the index."""
    data = proxy.to_full_dict()
    data["schemata"] = list(proxy.schema.names)
    data["caption"] = proxy.caption

    names = data.get("names", [])
    fps = set([fingerprints.generate(name) for name in names])
    fps.update(names)
    data["fingerprints"] = [fp for fp in fps if fp is not None]

    # Slight hack: a magic property in followthemoney that gets taken out
    # of the properties and added straight to the index text.
    properties = data.get("properties")
    data["text"] = properties.pop("indexText", [])

    # integer casting
    numeric = {}
    for prop in proxy.iterprops():
        if prop.type in NUMERIC_TYPES:
            values = proxy.get(prop)
            numeric[prop.name] = _numeric_values(prop.type, values)
    # also cast group field for dates
    numeric["dates"] = _numeric_values(registry.date, data.get("dates"))
    data["numeric"] = numeric

    # Context data - from aleph system, not followthemoney.
    # FIXME: Can there ever really be multiple role_ids?
    data["role_id"] = first(data.get("role_id"))
    data["mutable"] = max(ensure_list(data.get("mutable")), default=False)
    data["origin"] = ensure_list(data.get("origin"))
    created_at = data.get("created_at")
    if created_at:
        data["updated_at"] = data.get("updated_at", created_at)
    data["collection_id"] = collection.id
    # log.info("%s", pformat(data))
    entity_id = data.pop("id")
    return {
        "_id": entity_id,
        "_index": entities_write_index(proxy.schema),
        "_source": data,
    }

def delete_entity(entity_id, exclude=None, sync=False):
    """Delete an entity from the index."""
    if exclude is not None:
        exclude = entities_write_index(exclude)
    for entity in entities_by_ids(entity_id, excludes="*"):
        index = entity.get("_index")
        if index == exclude:
            continue
        delete_safe(index, entity_id)
<|endoftext|>


<commit_before>"""Convert parsed trees into SQLAlchemy objects """
from lark import Lark, Transformer, v_args
from sqlalchemy import func, distinct, case, and_, or_, not_, cast, Float
from datetime import date
from dateutil.relativedelta import relativedelta

from .field_grammar import (
    field_parser,
    noag_field_parser,
    noag_full_condition_parser,
    full_condition_parser,
)
from .utils import aggregations, find_column, ingredient_class_for_name, convert_value, calc_date_range
from recipe.exceptions import BadIngredient

@v_args(inline=True)  # Affects the signatures of the methods
class TransformToSQLAlchemyExpression(Transformer):
    """Converts a field to a SQLAlchemy expression """

    from operator import add, sub, mul, neg

    def __init__(self, selectable, require_aggregation=False):
        self.selectable = selectable
        self.require_aggregation = require_aggregation

    def number(self, value):
        try:
            return int(value)
        except ValueError:
            return float(value)

    def true(self, value):
        return True

    def false(self, value):
        return False

    def IN(self, value):
        return "IN"

    def NOTIN(self, value):
        return "NOTIN"

    def IS(self, value):
        return "IS"

    def NULL(self, value):
        return None

    def BETWEEN(self, value):
        return "BETWEEN"

    def aggregate(self, name):
        return aggregations.get(name.lower())

    def div(self, num, denom):
        """SQL safe division"""
        return case([(denom == 0, None)], else_=cast(num, Float) / cast(denom, Float))

    def column(self, name):
        return find_column(self.selectable, name)

    def agex(self, aggr, val):
        if val == "*":
            return func.count()
        else:
            return aggr(val)

    def expr(self, expr):
        return expr

    def case(self, *args):
        args = list(args)

        # If there's an odd number of args, pop the last one to use as the else
        if len(args) % 2:
            else_expr = args.pop()
        else:
            else_expr = None

        # collect the other args into pairs
        # ['a','b','c','d'] --> [('a',b'), ('c','d')]
        pairs = zip(args[::2], args[1::2])
        return case(pairs, else_=else_expr)

    def is_comparison(self, *args):
        """Things that can be compared with IS

        will be either IS NULL or an intelligent date
        like "IS prior year" or "IS next week".

        If intelligent dates are used, determine the relevant dates and
        return a tuple of start_date, end_date
        """
        if len(args) == 1 and args[0] is None:
            return None
        else:
            # INTELLIGENT_DATE_OFFSET: /prior/i | /current/i | /next/i
            # INTELLIGENT_DATE_UNITS: /ytd/i | /year/i | /qtr/i | /month/i | /week/i | /day/i
            offset, units = str(args[0]).lower(), str(args[1]).lower()
            return calc_date_range(offset, units, date.today())

    def relation_expr(self, left, rel, right):
        rel = rel.lower()
        comparators = {
            "=": "__eq__",
            ">": "__gt__",
            ">=": "__gte__",
            "!=": "__ne__",
            "<": "__lt__",
            "<=": "__lte__",
        }
        if right is None:
            return
        # Convert the right into a type compatible with the left
        right = convert_value(left, right)
        return getattr(left, comparators[rel])(right)

    def relation_expr_using_is(self, left, rel, right):
        """A relation expression like age is null or
        birth_date is last month"""
        # TODO: Why is this not handled by the tokenization
        if str(right).upper() == 'NULL':
            return left.is_(None)
        else:
            return left.between(*right)

    def array(self, *args):
        # TODO  check these are all the same type
        # And match the type of the column!
        return args

    def vector_relation_expr(self, left, rel, right):
        comparators = {"IN": "in_", "NOTIN": "notin_"}
        return getattr(left, comparators[rel.upper()])(right)

    def between_relation_expr(self, col, between, low, _, high):
        # TODO: check data types and convert dates.
        return col.between(convert_value(col, low), convert_value(col, high))

    def bool_expr(self, *exprs):
        if len(exprs) > 1:
            return or_(*exprs)
        else:
            return exprs[0]

    def string_literal(self, value):
        # Strip the quotes off of this string
        return value.value[1:-1]

    def bool_term(self, *exprs):
        if len(exprs) > 1:
            return and_(*exprs)
        else:
            return exprs[0]

    def not_bool_factor(self, notval, expr):
        return not_(expr)

def create_ingredient_from_parsed(ingr_dict, selectable):
    """ Create an ingredient from config version 2 object . """
    kind = ingr_dict.pop("kind", "metric")
    IngredientClass = ingredient_class_for_name(kind.title())
    if IngredientClass is None:
        raise BadIngredient("Unknown ingredient kind")

    if kind in ("metric", "dimension"):
        parser = field_parser if kind == "metric" else noag_field_parser
        fld_defn = ingr_dict.pop("field", None)
        tree = parser.parse(fld_defn)
        # Create a sqlalchemy expression from 'field' and pass it as the first arg
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                parser.parse(fld_defn)
            )
        ]

        # Convert quickselects to a kwarg with sqlalchemy expressions
        parsed_quickselects = []
        for qf in ingr_dict.pop("quickselects", []):
            parsed_quickselects.append(
                {
                    "name": qf["name"],
                    "condition": TransformToSQLAlchemyExpression(
                        selectable=selectable
                    ).transform(noag_full_condition_parser.parse(qf.get("condition"))),
                }
            )
        ingr_dict["quickselects"] = parsed_quickselects

        # Convert extra fields to sqlalchemy expressions and add them directly to
        # the kwargs
        for extra in ingr_dict.pop("extra_fields", []):
            ingr_dict[extra.get("name")] = TransformToSQLAlchemyExpression(
                selectable=selectable
            ).transform(parser.parse(extra.get("field")))

    elif kind == "filter":
        # Create a sqlalchemy expression from 'condition' and pass it as the first arg
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                noag_full_condition_parser.parse(ingr_dict.pop("condition", None))
            )
        ]

    elif kind == "having":
        # Create a sqlalchemy expression from 'condition' and pass it as the first arg
        # TODO: Force this to be an aggregate
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                full_condition_parser.parse(ingr_dict.pop("condition", None))
            )
        ]

    return IngredientClass(*args, **ingr_dict)
<commit_msg>Improve and, or and divide<commit_after>"""Convert parsed trees into SQLAlchemy objects """
from lark import Lark, Transformer, v_args
from sqlalchemy import func, distinct, case, and_, or_, not_, cast, Float
from datetime import date
from dateutil.relativedelta import relativedelta

from .field_grammar import (
    field_parser,
    noag_field_parser,
    noag_full_condition_parser,
    full_condition_parser,
)
from .utils import aggregations, find_column, ingredient_class_for_name, convert_value, calc_date_range
from recipe.exceptions import BadIngredient

@v_args(inline=True)  # Affects the signatures of the methods
class TransformToSQLAlchemyExpression(Transformer):
    """Converts a field to a SQLAlchemy expression """

    from operator import add, sub, mul, neg

    def __init__(self, selectable, require_aggregation=False):
        self.selectable = selectable
        self.require_aggregation = require_aggregation

    def number(self, value):
        try:
            return int(value)
        except ValueError:
            return float(value)

    def true(self, value):
        return True

    def false(self, value):
        return False

    def IN(self, value):
        return "IN"

    def NOTIN(self, value):
        return "NOTIN"

    def IS(self, value):
        return "IS"

    def NULL(self, value):
        return None

    def BETWEEN(self, value):
        return "BETWEEN"

    def aggregate(self, name):
        return aggregations.get(name.lower())

    def div(self, num, denom):
        """SQL safe division"""
        if isinstance(denom, (int, float)):
            if denom == 0:
                raise ValueError("Denominator can not be zero")
            elif isinstance(num, (int, float)):
                return num / denom
            else:
                return cast(num, Float) / denom
        else:
            if isinstance(num, (int, float)):
                return case([(denom == 0, None)], else_=num / cast(denom, Float))
            else:
                return case([(denom == 0, None)], else_=cast(num, Float) / cast(denom, Float))

    def column(self, name):
        return find_column(self.selectable, name)

    def agex(self, aggr, val):
        if val == "*":
            return func.count()
        else:
            return aggr(val)

    def expr(self, expr):
        return expr

    def case(self, *args):
        args = list(args)

        # If there's an odd number of args, pop the last one to use as the else
        if len(args) % 2:
            else_expr = args.pop()
        else:
            else_expr = None

        # collect the other args into pairs
        # ['a','b','c','d'] --> [('a',b'), ('c','d')]
        pairs = zip(args[::2], args[1::2])
        return case(pairs, else_=else_expr)

    def is_comparison(self, *args):
        """Things that can be compared with IS

        will be either IS NULL or an intelligent date
        like "IS prior year" or "IS next week".

        If intelligent dates are used, determine the relevant dates and
        return a tuple of start_date, end_date
        """
        if len(args) == 1 and args[0] is None:
            return None
        else:
            # INTELLIGENT_DATE_OFFSET: /prior/i | /current/i | /next/i
            # INTELLIGENT_DATE_UNITS: /ytd/i | /year/i | /qtr/i | /month/i | /week/i | /day/i
            offset, units = str(args[0]).lower(), str(args[1]).lower()
            return calc_date_range(offset, units, date.today())

    def relation_expr(self, left, rel, right):
        rel = rel.lower()
        comparators = {
            "=": "__eq__",
            ">": "__gt__",
            ">=": "__gte__",
            "!=": "__ne__",
            "<": "__lt__",
            "<=": "__lte__",
        }
        if right is None:
            return
        # Convert the right into a type compatible with the left
        right = convert_value(left, right)
        return getattr(left, comparators[rel])(right)

    def relation_expr_using_is(self, left, rel, right):
        """A relation expression like age is null or
        birth_date is last month"""
        # TODO: Why is this not handled by the tokenization
        if str(right).upper() == 'NULL':
            return left.is_(None)
        else:
            return left.between(*right)

    def array(self, *args):
        # TODO  check these are all the same type
        # And match the type of the column!
        return args

    def vector_relation_expr(self, left, rel, right):
        comparators = {"IN": "in_", "NOTIN": "notin_"}
        return getattr(left, comparators[rel.upper()])(right)

    def between_relation_expr(self, col, between, low, _, high):
        # TODO: check data types and convert dates.
        return col.between(convert_value(col, low), convert_value(col, high))

    def bool_expr(self, *exprs):
        if len(exprs) > 1:
            left, _, right = exprs
            return or_(left, right)
        else:
            return exprs[0]

    def string_literal(self, value):
        # Strip the quotes off of this string
        return value.value[1:-1]

    def bool_term(self, *exprs):
        if len(exprs) > 1:
            left, _, right = exprs
            return and_(left, right)
        else:
            return exprs[0]

    def not_bool_factor(self, notval, expr):
        return not_(expr)

def create_ingredient_from_parsed(ingr_dict, selectable):
    """ Create an ingredient from config version 2 object . """
    kind = ingr_dict.pop("kind", "metric")
    IngredientClass = ingredient_class_for_name(kind.title())
    if IngredientClass is None:
        raise BadIngredient("Unknown ingredient kind")

    if kind in ("metric", "dimension"):
        parser = field_parser if kind == "metric" else noag_field_parser
        fld_defn = ingr_dict.pop("field", None)
        tree = parser.parse(fld_defn)
        # Create a sqlalchemy expression from 'field' and pass it as the first arg
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                parser.parse(fld_defn)
            )
        ]

        # Convert quickselects to a kwarg with sqlalchemy expressions
        parsed_quickselects = []
        for qf in ingr_dict.pop("quickselects", []):
            parsed_quickselects.append(
                {
                    "name": qf["name"],
                    "condition": TransformToSQLAlchemyExpression(
                        selectable=selectable
                    ).transform(noag_full_condition_parser.parse(qf.get("condition"))),
                }
            )
        ingr_dict["quickselects"] = parsed_quickselects

        # Convert extra fields to sqlalchemy expressions and add them directly to
        # the kwargs
        for extra in ingr_dict.pop("extra_fields", []):
            ingr_dict[extra.get("name")] = TransformToSQLAlchemyExpression(
                selectable=selectable
            ).transform(parser.parse(extra.get("field")))

    elif kind == "filter":
        # Create a sqlalchemy expression from 'condition' and pass it as the first arg
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                noag_full_condition_parser.parse(ingr_dict.pop("condition", None))
            )
        ]

    elif kind == "having":
        # Create a sqlalchemy expression from 'condition' and pass it as the first arg
        # TODO: Force this to be an aggregate
        args = [
            TransformToSQLAlchemyExpression(selectable=selectable).transform(
                full_condition_parser.parse(ingr_dict.pop("condition", None))
            )
        ]

    return IngredientClass(*args, **ingr_dict)
<|endoftext|>


<commit_before>import logging
from flask_babel import gettext
from flask import Blueprint, request, send_file
from werkzeug.exceptions import BadRequest

from aleph.model import Linkage
from aleph.search import XrefQuery
from aleph.index.xref import get_xref
from aleph.logic.xref import export_matches
from aleph.logic.linkages import decide_xref
from aleph.views.serializers import XrefSerializer
from aleph.queues import queue_task, OP_XREF
from aleph.views.util import get_db_collection, get_index_collection
from aleph.views.util import parse_request, require, jsonify, obj_or_404

XLSX_MIME = 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'  # noqa
blueprint = Blueprint('xref_api', __name__)
log = logging.getLogger(__name__)

@blueprint.route('/api/2/collections/<int:collection_id>/xref', methods=['GET'])  # noqa
def index(collection_id):
    """
    ---
    get:
      summary: Fetch cross-reference results
      description: >-
        Fetch cross-reference matches for entities in the collection
        with id `collection_id`
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '200':
          description: OK
          content:
            application/json:
              schema:
                type: object
                allOf:
                - $ref: '#/components/schemas/QueryResponse'
                properties:
                  results:
                    type: array
                    items:
                      $ref: '#/components/schemas/XrefResponse'
      tags:
      - Xref
      - Collection
    """
    get_index_collection(collection_id)
    result = XrefQuery.handle(request, collection_id=collection_id)
    context_id = result.parser.getint('context_id', request.authz.id)
    require(request.authz.can_read_role(context_id))
    pairs = [(x.get('entity_id'), x.get('match_id')) for x in result.results]
    decisions = Linkage.decisions(pairs, context_id)
    for xref in result.results:
        key = (xref.get('entity_id'), xref.get('match_id'))
        xref['decision'] = decisions.get(key)
    return XrefSerializer.jsonify_result(result)

@blueprint.route('/api/2/collections/<int:collection_id>/xref', methods=['POST'])  # noqa
def generate(collection_id):
    """
    ---
    post:
      summary: Generate cross-reference matches
      description: >
        Generate cross-reference matches for entities in a collection.
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '202':
          content:
            application/json:
              schema:
                properties:
                  status:
                    description: accepted
                    type: string
                type: object
          description: Accepted
      tags:
      - Xref
      - Collection
    """
    collection = get_db_collection(collection_id, request.authz.WRITE)
    queue_task(collection, OP_XREF)
    return jsonify({'status': 'accepted'}, status=202)

@blueprint.route('/api/2/collections/<int:collection_id>/xref.xlsx')
def export(collection_id):
    """
    ---
    get:
      summary: Download cross-reference results
      description: Download results of cross-referencing as an Excel file
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '200':
          description: OK
          content:
            application/vnd.openxmlformats-officedocument.spreadsheetml.sheet:
              schema:
                type: object
      tags:
      - Xref
      - Collection
    """
    collection = get_db_collection(collection_id, request.authz.READ)
    buffer = export_matches(collection, request.authz)
    file_name = '%s - Crossreference.xlsx' % collection.label
    return send_file(buffer,
                     mimetype=XLSX_MIME,
                     as_attachment=True,
                     attachment_filename=file_name)

@blueprint.route('/api/2/collections/<int:collection_id>/xref/<xref_id>', methods=['POST'])  # noqa
def decide(collection_id, xref_id):
    """
    ---
    post:
      summary: Give feedback about the veracity of an xref match.
      description: >
        This lets a user decide if they think a given xref match is a true or
        false match, and what group of users (context) should be privy to this
        insight.
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      - in: path
        name: xref_id
        required: true
        schema:
          type: string
      requestBody:
        content:
          application/json:
            schema:
              $ref: '#/components/schemas/XrefDecide'
      responses:
        '202':
          content:
            application/json:
              schema:
                properties:
                  status:
                    description: accepted
                    type: string
                type: object
          description: Accepted
      tags:
      - Xref
      - Linkage
    """
    require(request.authz.logged_in)
    data = parse_request('XrefDecide')
    xref = obj_or_404(get_xref(xref_id, collection_id=collection_id))
    try:
        context_id = int(data.get('context_id', request.authz.id))
    except (ValueError, TypeError):
        raise BadRequest(gettext("Invalid context."))
    require(context_id in request.authz.private_roles)
    decide_xref(xref,
                decision=data.get('decision'),
                context_id=context_id,
                decider_id=request.authz.id)
    return jsonify({'status': 'ok'}, status=204)
<commit_msg>fix test<commit_after>import logging
from flask_babel import gettext
from flask import Blueprint, request, send_file
from werkzeug.exceptions import BadRequest

from aleph.model import Linkage
from aleph.search import XrefQuery
from aleph.index.xref import get_xref
from aleph.logic.xref import export_matches
from aleph.logic.linkages import decide_xref
from aleph.views.serializers import XrefSerializer
from aleph.queues import queue_task, OP_XREF
from aleph.views.util import get_db_collection, get_index_collection
from aleph.views.util import parse_request, require, jsonify, obj_or_404

XLSX_MIME = 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'  # noqa
blueprint = Blueprint('xref_api', __name__)
log = logging.getLogger(__name__)

@blueprint.route('/api/2/collections/<int:collection_id>/xref', methods=['GET'])  # noqa
def index(collection_id):
    """
    ---
    get:
      summary: Fetch cross-reference results
      description: >-
        Fetch cross-reference matches for entities in the collection
        with id `collection_id`
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '200':
          description: OK
          content:
            application/json:
              schema:
                type: object
                allOf:
                - $ref: '#/components/schemas/QueryResponse'
                properties:
                  results:
                    type: array
                    items:
                      $ref: '#/components/schemas/XrefResponse'
      tags:
      - Xref
      - Collection
    """
    get_index_collection(collection_id)
    result = XrefQuery.handle(request, collection_id=collection_id)
    context_id = result.parser.getint('context_id')
    context_id = context_id or request.authz.id
    require(request.authz.can_read_role(context_id))
    pairs = [(x.get('entity_id'), x.get('match_id')) for x in result.results]
    decisions = Linkage.decisions(pairs, context_id)
    for xref in result.results:
        key = (xref.get('entity_id'), xref.get('match_id'))
        xref['decision'] = decisions.get(key)
    return XrefSerializer.jsonify_result(result)

@blueprint.route('/api/2/collections/<int:collection_id>/xref', methods=['POST'])  # noqa
def generate(collection_id):
    """
    ---
    post:
      summary: Generate cross-reference matches
      description: >
        Generate cross-reference matches for entities in a collection.
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '202':
          content:
            application/json:
              schema:
                properties:
                  status:
                    description: accepted
                    type: string
                type: object
          description: Accepted
      tags:
      - Xref
      - Collection
    """
    collection = get_db_collection(collection_id, request.authz.WRITE)
    queue_task(collection, OP_XREF)
    return jsonify({'status': 'accepted'}, status=202)

@blueprint.route('/api/2/collections/<int:collection_id>/xref.xlsx')
def export(collection_id):
    """
    ---
    get:
      summary: Download cross-reference results
      description: Download results of cross-referencing as an Excel file
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      responses:
        '200':
          description: OK
          content:
            application/vnd.openxmlformats-officedocument.spreadsheetml.sheet:
              schema:
                type: object
      tags:
      - Xref
      - Collection
    """
    collection = get_db_collection(collection_id, request.authz.READ)
    buffer = export_matches(collection, request.authz)
    file_name = '%s - Crossreference.xlsx' % collection.label
    return send_file(buffer,
                     mimetype=XLSX_MIME,
                     as_attachment=True,
                     attachment_filename=file_name)

@blueprint.route('/api/2/collections/<int:collection_id>/xref/<xref_id>', methods=['POST'])  # noqa
def decide(collection_id, xref_id):
    """
    ---
    post:
      summary: Give feedback about the veracity of an xref match.
      description: >
        This lets a user decide if they think a given xref match is a true or
        false match, and what group of users (context) should be privy to this
        insight.
      parameters:
      - in: path
        name: collection_id
        required: true
        schema:
          type: integer
      - in: path
        name: xref_id
        required: true
        schema:
          type: string
      requestBody:
        content:
          application/json:
            schema:
              $ref: '#/components/schemas/XrefDecide'
      responses:
        '202':
          content:
            application/json:
              schema:
                properties:
                  status:
                    description: accepted
                    type: string
                type: object
          description: Accepted
      tags:
      - Xref
      - Linkage
    """
    require(request.authz.logged_in)
    data = parse_request('XrefDecide')
    xref = obj_or_404(get_xref(xref_id, collection_id=collection_id))
    try:
        context_id = int(data.get('context_id', request.authz.id))
    except (ValueError, TypeError):
        raise BadRequest(gettext("Invalid context."))
    require(context_id in request.authz.private_roles)
    decide_xref(xref,
                decision=data.get('decision'),
                context_id=context_id,
                decider_id=request.authz.id)
    return jsonify({'status': 'ok'}, status=204)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Licensed 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
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2017

# This module contains functions that are used with the get_parameters() function defined in the information module.

# WARNING: IN GENERAL, NEEDS TO USE PLUG-IN MANAGER

from pilot.info import infosys

import logging
logger = logging.getLogger(__name__)

def get_maximum_input_sizes(queuedata):
    """
    This function returns the maximum allowed size for all input files. The sum of all input file sizes should not
    exceed this value.

    :param queuedata: the queuedata dictionary from schedconfig.
    :return: maxinputsizes (integer value in GB).
    """

    try:
        _maxinputsizes = infosys.queuedata.maxwdir  # normally 14336+2000 MB
    except TypeError, e:
        from pilot.util.config import config
        _maxinputsizes = config.Pilot.maximum_input_file_sizes  # MB
        logger.warning('could not convert schedconfig value for maxwdir: %s (will use default value instead - %d GB)' %
                       (e, _maxinputsizes))

    return _maxinputsizes

def convert_to_int(parameter, default=None):
    """
    Try to convert a given parameter to an integer value.
    The default parameter can be used to force the function to always return a given value in case the integer
    conversion, int(parameter), fails.

    :param parameter: parameter (any type).
    :param default: None by default (if set, always return an integer; the given value will be returned if
    conversion to integer fails).
    :return: converted integer.
    """

    try:
        value = int(parameter)
    except Exception:  # can be ValueError or TypeValue (for None)
        value = default

    return value
<commit_msg>Removed useless queuedata argument from get_maximum_input_sizes()<commit_after>#!/usr/bin/env python
# Licensed 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
#
# Authors:
# - Paul Nilsson, paul.nilsson@cern.ch, 2017

# This module contains functions that are used with the get_parameters() function defined in the information module.

# WARNING: IN GENERAL, NEEDS TO USE PLUG-IN MANAGER

from pilot.info import infosys

import logging
logger = logging.getLogger(__name__)

def get_maximum_input_sizes():
    """
    This function returns the maximum allowed size for all input files. The sum of all input file sizes should not
    exceed this value.

    :return: maxinputsizes (integer value in GB).
    """

    try:
        _maxinputsizes = infosys.queuedata.maxwdir  # normally 14336+2000 MB
    except TypeError, e:
        from pilot.util.config import config
        _maxinputsizes = config.Pilot.maximum_input_file_sizes  # MB
        logger.warning('could not convert schedconfig value for maxwdir: %s (will use default value instead - %d GB)' %
                       (e, _maxinputsizes))

    return _maxinputsizes

def convert_to_int(parameter, default=None):
    """
    Try to convert a given parameter to an integer value.
    The default parameter can be used to force the function to always return a given value in case the integer
    conversion, int(parameter), fails.

    :param parameter: parameter (any type).
    :param default: None by default (if set, always return an integer; the given value will be returned if
    conversion to integer fails).
    :return: converted integer.
    """

    try:
        value = int(parameter)
    except Exception:  # can be ValueError or TypeValue (for None)
        value = default

    return value
<|endoftext|>


<commit_before><commit_msg>QChem: inherit Inputfile in QCInput<commit_after><|endoftext|>


<commit_before><commit_msg>Modify check for ISPIN = 2 and separate_spins=True<commit_after><|endoftext|>


<commit_before>from django.contrib import admin
from django.urls import reverse
from django.utils.html import format_html

from .models import Referral, ReferralResponse

@admin.register(Referral)
class ReferralAdmin(admin.ModelAdmin):
    list_display = [
        "user",
        "code",
        "label",
        "redirect_to",
        "target_content_type",
        "target_object_id",
        "url",
    ]
    readonly_fields = ["code", "created_at"]
    list_filter = ["target_content_type", "created_at"]
    search_fields = ["user__first_name", "user__last_name", "user__email", "user__username", "code"]
    autocomplete_fields = ["user"]

@admin.register(ReferralResponse)
class ReferralResponseAdmin(admin.ModelAdmin):
    list_display = [
        "referral",
        "session_key",
        "user",
        "ip_address",
        "action",
        "target_object_link",
    ]
    readonly_fields = ["referral", "session_key", "user", "ip_address", "action", "target_object_link"]
    list_filter = ["action", "created_at"]
    search_fields = [
        "referral__code",
        "referral__user__email",
        "referral__user__username",
        "referral__user__first_name",
        "referral__user__last_name",
        "user__email",
        "user__username",
        "user__first_name",
        "user__last_name",
        "ip_address",
    ]

    def target_object_link(self, obj):
        if obj.pk and obj.target:
            admin_link = reverse("admin:%s_%s_change" % (obj.target_content_type.app_label, obj.target_content_type.model), args=(obj.target.pk,))
            return format_html('{}', admin_link, obj.target.__str__())
<commit_msg>Update admin.py<commit_after>from django.contrib import admin
from django.urls import reverse
from django.utils.html import format_html

from .models import Referral, ReferralResponse

@admin.register(Referral)
class ReferralAdmin(admin.ModelAdmin):
    list_display = [
        "user",
        "code",
        "label",
        "redirect_to",
        "target_content_type",
        "target_object_id",
        "url",
    ]
    readonly_fields = ["code", "created_at"]
    list_filter = ["target_content_type"]
    search_fields = ["user__username", "user__email", "code"]
    autocomplete_fields = ["user"]

@admin.register(ReferralResponse)
class ReferralResponseAdmin(admin.ModelAdmin):
    list_display = [
        "referral",
        "session_key",
        "user",
        "ip_address",
        "action",
        "target_object_link",
    ]
    readonly_fields = ["referral", "session_key", "user", "ip_address", "action", "target_object_link"]
    list_filter = ["action"]
    search_fields = [
        "referral__code",
        "referral__user__email",
        "referral__user__username",
        "user__email",
        "user__username",
        "ip_address",
    ]

    def target_object_link(self, obj):
        if obj.pk and obj.target:
            admin_link = reverse("admin:%s_%s_change" % (obj.target_content_type.app_label, obj.target_content_type.model), args=(obj.target.pk,))
            return format_html('{}', admin_link, obj.target.__str__())
<|endoftext|>


<commit_before>import discord

from redbot.core import commands, checks, Config, bot
from redbot.core.utils.chat_formatting import box, humanize_list
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS

from io import BufferedIOBase, IOBase, BytesIO
from datetime import datetime, timezone

# https://red-discordbot.readthedocs.io/en/latest/framework_utils.html
# https://github.com/Cog-Creators/Red-DiscordBot/blob/V3/develop/redbot/cogs/economy/economy.py

# https://discordpy.readthedocs.io/en/latest/api.html#discord.TextChannel.send
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Message.attachments
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Attachment

# commands
class PinHistory(commands.Cog):
    """My custom counter cog"""
    def __init__(self):
        self.config = Config.get_conf(self, identifier=108586980896678, force_registration=True)
        default_guild = {
            "pin_limit": 50, # Maximum number of pinned messages
            "manage_pins": True, # Whether the bot should manage pinned messages in the monitored channels (Remove old pins etc)
            "pin_history": [], # List of ID's for all pinned messages that have been archived
            "monitored_channels": [],  # Channels to monitor for changes in pins
            "archive_channels": [], # Channels to archive the pinned messages to
        }

        self.config.register_guild(**default_guild)

    # Commands
    @checks.admin() # Checks if they have admin role - https://red-discordbot.readthedocs.io/en/latest/framework_checks.html
    @commands.group(invoke_without_command=False)
    async def pinhistory(self, ctx): #Recount group
        """
        Base postset group
        """
        pass

    @checks.admin()
    @pinhistory.group(name="settings", invoke_without_command=True)
    async def pinhistory_settings(self, ctx):
        """
        Display current settings
        """
        # Should display current settings
        settings = await self.config.guild(ctx.channel.guild).get_raw()
        await ctx.send("```{}```".format(str(settings)))

    @checks.admin()
    @pinhistory.group(name="pinlimit", invoke_without_command=True)
    async def pinhistory_pinlimit(self, ctx, new_pin_limit : int):
        """
        Set the number of pins allowed in a monitored channel
        """
        async with self.config.guild(ctx.channel.guild).pin_limit() as pin_limit:
            pin_limit = new_pin_limit
            await ctx.send("Set pin limit to {}".format(pin_limit))

    @checks.admin()
    @pinhistory.group(name="managepins", invoke_without_command=True)
    async def pinhistory_managepins(self, ctx):
        """
        Toggle if pinned messages should be managed. EG, deleted once they go over pin limit
        """
        async with self.config.guild(ctx.channel.guild).manage_pins() as manage_pins:
            manage_pins = not manage_pins
            await ctx.send("Set pin management to {}".format(new_value))

    @checks.admin()
    @pinhistory.group(name="toggle", invoke_without_command=False)
    async def pinhistory_toggle(self, ctx):
        pass

    @checks.admin()
    @pinhistory_toggle.group(name="monitor", invoke_without_command=True)
    async def pinhistory_toggle_monitor(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).monitored_channels() as monitored_channels:
            if channel.id not in monitored_channels:
                monitored_channels.append(channel.id)
                await ctx.send("Enabled monitoring for {}".format(channel.name))
            else:
                monitored_channels.remove(channel.id)
                await ctx.send("Disabled monitoring for {}".format(channel.name))

    @checks.admin()
    @pinhistory_toggle.group(name="archive", invoke_without_command=True)
    async def pinhistory_toggle_archive(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).archive_channels() as archive_channels:
            if channel.id not in archive_channels:
                archive_channels.append(channel.id)
                await ctx.send("Enabled archiving to {}".format(channel.name))
            else:
                archive_channels.remove(channel.id)
                await ctx.send("Disabled archiving for {}".format(channel.name))

    # https://leovoel.github.io/embed-visualizer/
    # https://cog-creators.github.io/discord-embed-sandbox/

    def create_embed(self, message):
        "Create the correctly formatted embed for pinned messages"
        # Create embed using information from message
        embed_message = discord.Embed(description=message.content, timestamp=datetime.replace(message.created_at, tzinfo=timezone.utc))
        embed_message.set_author(name=message.author.display_name, url="https://discord.com/users/{}".format(message.author.id), icon_url=message.author.avatar_url)
        embed_message.set_thumbnail(url=message.author.avatar_url)
        embed_message.add_field(name="Sources", value="{} | [Message]({})".format(message.channel.mention, message.jump_url), inline=False)
        return embed_message

    async def return_attachments(self, message):
        "Returns attachments in a file list"
        files = []
        for attachment in message.attachments:
            attachment_bytes = BytesIO(await attachment.read())
            files.append(discord.File(fp=attachment_bytes))
        return files

    # https://discordpy.readthedocs.io/en/latest/api.html#discord.on_guild_channel_pins_update
    # Monitored Events
    @commands.Cog.listener('on_guild_channel_pins_update') # Executes the below command when a channels pinned messages changes
    async def on_pin_update(self, channel, last_pin):
        monitored_channels = await self.config.guild(channel.guild).monitored_channels()
        archive_channels = await self.config.guild(channel.guild).archive_channels()
        if (channel.id in monitored_channels) and (last_pin != None):
            # Check if this pin hasn't already been archived
            last_pinned_message = (await channel.pins())[0]
            async with self.config.guild(channel.guild).pin_history() as pin_history:
                if last_pinned_message.id not in pin_history:
                    # Get list of attachments from message
                    # Save list of attachment objects to a list of file like objects
                    # Send embed with provided attachments
                    embed_message = self.create_embed(last_pinned_message)
                    files = await self.return_attachments(last_pinned_message)
                    for channel_id in archive_channels:
                        channel = channel.guild.get_channel(channel_id)
                        await channel.send(embed=embed_message, files=files)
                    pin_history.append(last_pinned_message.id)
<commit_msg>Included filename in file upload<commit_after>import discord

from redbot.core import commands, checks, Config, bot
from redbot.core.utils.chat_formatting import box, humanize_list
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS

from io import BufferedIOBase, IOBase, BytesIO
from datetime import datetime, timezone

# https://red-discordbot.readthedocs.io/en/latest/framework_utils.html
# https://github.com/Cog-Creators/Red-DiscordBot/blob/V3/develop/redbot/cogs/economy/economy.py

# https://discordpy.readthedocs.io/en/latest/api.html#discord.TextChannel.send
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Message.attachments
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Attachment

# commands
class PinHistory(commands.Cog):
    """My custom counter cog"""
    def __init__(self):
        self.config = Config.get_conf(self, identifier=108586980896678, force_registration=True)
        default_guild = {
            "pin_limit": 50, # Maximum number of pinned messages
            "manage_pins": True, # Whether the bot should manage pinned messages in the monitored channels (Remove old pins etc)
            "pin_history": [], # List of ID's for all pinned messages that have been archived
            "monitored_channels": [],  # Channels to monitor for changes in pins
            "archive_channels": [], # Channels to archive the pinned messages to
        }

        self.config.register_guild(**default_guild)

    # Commands
    @checks.admin() # Checks if they have admin role - https://red-discordbot.readthedocs.io/en/latest/framework_checks.html
    @commands.group(invoke_without_command=False)
    async def pinhistory(self, ctx): #Recount group
        """
        Base postset group
        """
        pass

    @checks.admin()
    @pinhistory.group(name="settings", invoke_without_command=True)
    async def pinhistory_settings(self, ctx):
        """
        Display current settings
        """
        # Should display current settings
        settings = await self.config.guild(ctx.channel.guild).get_raw()
        await ctx.send("```{}```".format(str(settings)))

    @checks.admin()
    @pinhistory.group(name="pinlimit", invoke_without_command=True)
    async def pinhistory_pinlimit(self, ctx, new_pin_limit : int):
        """
        Set the number of pins allowed in a monitored channel
        """
        async with self.config.guild(ctx.channel.guild).pin_limit() as pin_limit:
            pin_limit = new_pin_limit
            await ctx.send("Set pin limit to {}".format(pin_limit))

    @checks.admin()
    @pinhistory.group(name="managepins", invoke_without_command=True)
    async def pinhistory_managepins(self, ctx):
        """
        Toggle if pinned messages should be managed. EG, deleted once they go over pin limit
        """
        async with self.config.guild(ctx.channel.guild).manage_pins() as manage_pins:
            manage_pins = not manage_pins
            await ctx.send("Set pin management to {}".format(new_value))

    @checks.admin()
    @pinhistory.group(name="toggle", invoke_without_command=False)
    async def pinhistory_toggle(self, ctx):
        pass

    @checks.admin()
    @pinhistory_toggle.group(name="monitor", invoke_without_command=True)
    async def pinhistory_toggle_monitor(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).monitored_channels() as monitored_channels:
            if channel.id not in monitored_channels:
                monitored_channels.append(channel.id)
                await ctx.send("Enabled monitoring for {}".format(channel.name))
            else:
                monitored_channels.remove(channel.id)
                await ctx.send("Disabled monitoring for {}".format(channel.name))

    @checks.admin()
    @pinhistory_toggle.group(name="archive", invoke_without_command=True)
    async def pinhistory_toggle_archive(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).archive_channels() as archive_channels:
            if channel.id not in archive_channels:
                archive_channels.append(channel.id)
                await ctx.send("Enabled archiving to {}".format(channel.name))
            else:
                archive_channels.remove(channel.id)
                await ctx.send("Disabled archiving for {}".format(channel.name))

    # https://leovoel.github.io/embed-visualizer/
    # https://cog-creators.github.io/discord-embed-sandbox/

    def create_embed(self, message):
        "Create the correctly formatted embed for pinned messages"
        # Create embed using information from message
        embed_message = discord.Embed(description=message.content, timestamp=datetime.replace(message.created_at, tzinfo=timezone.utc))
        embed_message.set_author(name=message.author.display_name, url="https://discord.com/users/{}".format(message.author.id), icon_url=message.author.avatar_url)
        embed_message.set_thumbnail(url=message.author.avatar_url)
        embed_message.add_field(name="Sources", value="{} | [Message]({})".format(message.channel.mention, message.jump_url), inline=False)
        return embed_message

    async def return_attachments(self, message):
        "Returns attachments in a file list"
        files = []
        for attachment in message.attachments:
            attachment_bytes = BytesIO(await attachment.read())
            files.append(discord.File(fp=attachment_bytes, filename=attachment.filename))
        return files

    # https://discordpy.readthedocs.io/en/latest/api.html#discord.on_guild_channel_pins_update
    # Monitored Events
    @commands.Cog.listener('on_guild_channel_pins_update') # Executes the below command when a channels pinned messages changes
    async def on_pin_update(self, channel, last_pin):
        monitored_channels = await self.config.guild(channel.guild).monitored_channels()
        archive_channels = await self.config.guild(channel.guild).archive_channels()
        if (channel.id in monitored_channels) and (last_pin != None):
            # Check if this pin hasn't already been archived
            last_pinned_message = (await channel.pins())[0]
            async with self.config.guild(channel.guild).pin_history() as pin_history:
                if last_pinned_message.id not in pin_history:
                    # Get list of attachments from message
                    # Save list of attachment objects to a list of file like objects
                    # Send embed with provided attachments
                    embed_message = self.create_embed(last_pinned_message)
                    files = await self.return_attachments(last_pinned_message)
                    for channel_id in archive_channels:
                        channel = channel.guild.get_channel(channel_id)
                        await channel.send(embed=embed_message, files=files)
                    pin_history.append(last_pinned_message.id)
<|endoftext|>


<commit_before>import discord

from redbot.core import commands, checks, Config, bot
from redbot.core.utils.chat_formatting import box, humanize_list
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS

from io import BytesIO
from datetime import datetime, timezone

# https://red-discordbot.readthedocs.io/en/latest/framework_utils.html
# https://github.com/Cog-Creators/Red-DiscordBot/blob/V3/develop/redbot/cogs/economy/economy.py

# https://discordpy.readthedocs.io/en/latest/api.html#discord.TextChannel.send
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Message.attachments
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Attachment

# commands
class PinHistory(commands.Cog):
    """My custom counter cog"""
    def __init__(self):
        self.config = Config.get_conf(self, identifier=108586980896678, force_registration=True)
        default_guild = {
            "pin_limit": 50, # Maximum number of pinned messages
            "manage_pins": True, # Whether the bot should manage pinned messages in the monitored channels (Remove old pins etc)
            "pin_history": [], # List of ID's for all pinned messages that have been archived
            "monitored_channels": [],  # Channels to monitor for changes in pins
            "archive_channels": [], # Channels to archive the pinned messages to
        }

        self.config.register_guild(**default_guild)

    # Commands
    @checks.admin() # Checks if they have admin role - https://red-discordbot.readthedocs.io/en/latest/framework_checks.html
    @commands.group(invoke_without_command=False)
    async def pinhistory(self, ctx): #Recount group
        """
        Base pinhistory group
        """
        pass

    @checks.admin()
    @pinhistory.group(name="clear", invoke_without_command=False)
    async def pinhistory_clear(self, ctx):
        """
        Base clear group
        """
        pass

    @checks.admin()
    @pinhistory_clear.group(name="history", invoke_without_command=True)
    async def pinhistory_clear_history(self, ctx):
        """
        Clears all pin history
        """
        async with self.config.guild(channel.guild).pin_history() as pin_history:
            pin_history = []

    @checks.admin()
    @pinhistory.group(name="update", invoke_without_command=False)
    async def pinhistory_update(self, ctx):
        """
        Base update group. Configure where to update archived pins
        """
        pass

    @checks.admin()
    @pinhistory_update.group(name="all", invoke_without_command=True)
    async def pinhistory_update_all(self, ctx):
        """
        Updates all pin history
        """
        monitored_channels = await self.config.guild(channel.guild).monitored_channels()
        for monitored_channel_id in monitored_channels:
            monitored_channel = ctx.guild.get_channel(monitored_channel_id)
            for pin in await monitored_channel.pins():
                await self.archive_pin(pin)

    @checks.admin()
    @pinhistory.group(name="settings", invoke_without_command=True)
    async def pinhistory_settings(self, ctx):
        """
        Display current settings
        """
        # Should display current settings
        settings = await self.config.guild(ctx.channel.guild).get_raw()
        await ctx.send("```{}```".format(str(settings)))

    @checks.admin()
    @pinhistory.group(name="pinlimit", invoke_without_command=True)
    async def pinhistory_pinlimit(self, ctx, new_pin_limit : int):
        """
        Set the number of pins allowed in a monitored channel
        """
        async with self.config.guild(ctx.channel.guild).pin_limit() as pin_limit:
            pin_limit = new_pin_limit
            await ctx.send("Set pin limit to {}".format(pin_limit))

    @checks.admin()
    @pinhistory.group(name="managepins", invoke_without_command=True)
    async def pinhistory_managepins(self, ctx):
        """
        Toggle if pinned messages should be managed. EG, deleted once they go over pin limit
        """
        async with self.config.guild(ctx.channel.guild).manage_pins() as manage_pins:
            manage_pins = not manage_pins
            await ctx.send("Set pin management to {}".format(new_value))

    @checks.admin()
    @pinhistory.group(name="toggle", invoke_without_command=False)
    async def pinhistory_toggle(self, ctx):
        pass

    @checks.admin()
    @pinhistory_toggle.group(name="monitor", invoke_without_command=True)
    async def pinhistory_toggle_monitor(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).monitored_channels() as monitored_channels:
            if channel.id not in monitored_channels:
                monitored_channels.append(channel.id)
                await ctx.send("Enabled monitoring for {}".format(channel.name))
            else:
                monitored_channels.remove(channel.id)
                await ctx.send("Disabled monitoring for {}".format(channel.name))

    @checks.admin()
    @pinhistory_toggle.group(name="archive", invoke_without_command=True)
    async def pinhistory_toggle_archive(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).archive_channels() as archive_channels:
            if channel.id not in archive_channels:
                archive_channels.append(channel.id)
                await ctx.send("Enabled archiving to {}".format(channel.name))
            else:
                archive_channels.remove(channel.id)
                await ctx.send("Disabled archiving for {}".format(channel.name))

    # https://leovoel.github.io/embed-visualizer/
    # https://cog-creators.github.io/discord-embed-sandbox/

    async def create_embed(self, message):
        "Create the correctly formatted embed for pinned messages"
        # Create embed using information from message
        embed_message = discord.Embed()
        embed_message.title = "{} posted".format(message.author.display_name)
        embed_message.description = message.content
        embed_message.timestamp = datetime.replace(message.created_at, tzinfo=timezone.utc)
        embed_message.set_author(name="{}#{}".format(message.author.name, message.author.discriminator), url="https://discord.com/users/{}".format(message.author.id), icon_url=message.author.avatar_url)
        embed_message.set_thumbnail(url=message.author.avatar_url)
        embed_message.add_field(name="Sources", value="{} | [Message]({})".format(message.channel.mention, message.jump_url), inline=False)
        attachment_count = len(message.attachments)
        if attachment_count != 0:
            if attachment_count == 1:
                attachment_text = "attachment"
            else:
                attachment_text = "attachments"
            embed_message.set_footer(text="{} {}".format(attachment_count, attachment_text))
        return embed_message

    async def return_attachments(self, message):
        "Returns attachments in a file list"
        files = []
        for attachment in message.attachments:
            attachment_bytes = BytesIO(await attachment.read())
            files.append(discord.File(fp=attachment_bytes, filename=attachment.filename))
        return files

    async def archive_pin(self, channel, message):
        """
        Archive pinned messages
        """
        archive_channels = await self.config.guild(channel.guild).archive_channels()
        async with self.config.guild(channel.guild).pin_history() as pin_history:
            if message.id not in pin_history:
                # Get list of attachments from message
                # Save list of attachment objects to a list of file like objects
                # Send embed with provided attachments
                embed_message = await self.create_embed(message)
                files = await self.return_attachments(message)
                for channel_id in archive_channels:
                    archive_channel = channel.guild.get_channel(channel_id)
                    await archive_channel.send(embed=embed_message, files=files)
                pin_history.append(message.id)
    # https://discordpy.readthedocs.io/en/latest/api.html#discord.on_guild_channel_pins_update
    # Monitored Events
    @commands.Cog.listener('on_guild_channel_pins_update') # Executes the below command when a channels pinned messages changes
    async def on_pin_update(self, channel, last_pin):
        monitored_channels = await self.config.guild(channel.guild).monitored_channels()
        if (channel.id in monitored_channels) and (last_pin != None):
            # Check if this pin hasn't already been archived
            last_pinned_message = (await channel.pins())[0]
            await self.archive_pin(channel, last_pinned_message)
<commit_msg>Corrected ctx references<commit_after>import discord

from redbot.core import commands, checks, Config, bot
from redbot.core.utils.chat_formatting import box, humanize_list
from redbot.core.utils.menus import menu, DEFAULT_CONTROLS

from io import BytesIO
from datetime import datetime, timezone

# https://red-discordbot.readthedocs.io/en/latest/framework_utils.html
# https://github.com/Cog-Creators/Red-DiscordBot/blob/V3/develop/redbot/cogs/economy/economy.py

# https://discordpy.readthedocs.io/en/latest/api.html#discord.TextChannel.send
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Message.attachments
# https://discordpy.readthedocs.io/en/latest/api.html#discord.Attachment

# commands
class PinHistory(commands.Cog):
    """My custom counter cog"""
    def __init__(self):
        self.config = Config.get_conf(self, identifier=108586980896678, force_registration=True)
        default_guild = {
            "pin_limit": 50, # Maximum number of pinned messages
            "manage_pins": True, # Whether the bot should manage pinned messages in the monitored channels (Remove old pins etc)
            "pin_history": [], # List of ID's for all pinned messages that have been archived
            "monitored_channels": [],  # Channels to monitor for changes in pins
            "archive_channels": [], # Channels to archive the pinned messages to
        }

        self.config.register_guild(**default_guild)

    # Commands
    @checks.admin() # Checks if they have admin role - https://red-discordbot.readthedocs.io/en/latest/framework_checks.html
    @commands.group(invoke_without_command=False)
    async def pinhistory(self, ctx): #Recount group
        """
        Base pinhistory group
        """
        pass

    @checks.admin()
    @pinhistory.group(name="clear", invoke_without_command=False)
    async def pinhistory_clear(self, ctx):
        """
        Base clear group
        """
        pass

    @checks.admin()
    @pinhistory_clear.group(name="history", invoke_without_command=True)
    async def pinhistory_clear_history(self, ctx):
        """
        Clears all pin history
        """
        async with self.config.guild(ctx.guild).pin_history() as pin_history:
            pin_history = []

    @checks.admin()
    @pinhistory.group(name="update", invoke_without_command=False)
    async def pinhistory_update(self, ctx):
        """
        Base update group. Configure where to update archived pins
        """
        pass

    @checks.admin()
    @pinhistory_update.group(name="all", invoke_without_command=True)
    async def pinhistory_update_all(self, ctx):
        """
        Updates all pin history
        """
        monitored_channels = await self.config.guild(ctx.guild).monitored_channels()
        for monitored_channel_id in monitored_channels:
            monitored_channel = ctx.guild.get_channel(monitored_channel_id)
            for pin in await monitored_channel.pins():
                await self.archive_pin(pin)

    @checks.admin()
    @pinhistory.group(name="settings", invoke_without_command=True)
    async def pinhistory_settings(self, ctx):
        """
        Display current settings
        """
        # Should display current settings
        settings = await self.config.guild(ctx.channel.guild).get_raw()
        await ctx.send("```{}```".format(str(settings)))

    @checks.admin()
    @pinhistory.group(name="pinlimit", invoke_without_command=True)
    async def pinhistory_pinlimit(self, ctx, new_pin_limit : int):
        """
        Set the number of pins allowed in a monitored channel
        """
        async with self.config.guild(ctx.channel.guild).pin_limit() as pin_limit:
            pin_limit = new_pin_limit
            await ctx.send("Set pin limit to {}".format(pin_limit))

    @checks.admin()
    @pinhistory.group(name="managepins", invoke_without_command=True)
    async def pinhistory_managepins(self, ctx):
        """
        Toggle if pinned messages should be managed. EG, deleted once they go over pin limit
        """
        async with self.config.guild(ctx.channel.guild).manage_pins() as manage_pins:
            manage_pins = not manage_pins
            await ctx.send("Set pin management to {}".format(new_value))

    @checks.admin()
    @pinhistory.group(name="toggle", invoke_without_command=False)
    async def pinhistory_toggle(self, ctx):
        pass

    @checks.admin()
    @pinhistory_toggle.group(name="monitor", invoke_without_command=True)
    async def pinhistory_toggle_monitor(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).monitored_channels() as monitored_channels:
            if channel.id not in monitored_channels:
                monitored_channels.append(channel.id)
                await ctx.send("Enabled monitoring for {}".format(channel.name))
            else:
                monitored_channels.remove(channel.id)
                await ctx.send("Disabled monitoring for {}".format(channel.name))

    @checks.admin()
    @pinhistory_toggle.group(name="archive", invoke_without_command=True)
    async def pinhistory_toggle_archive(self, ctx, channel=None):
        """
        Monitors channel, if none is given it'll use the one in context. If one is mentioned, it'll use that one
        """
        if channel == None:
            channel = ctx.channel
        async with self.config.guild(channel.guild).archive_channels() as archive_channels:
            if channel.id not in archive_channels:
                archive_channels.append(channel.id)
                await ctx.send("Enabled archiving to {}".format(channel.name))
            else:
                archive_channels.remove(channel.id)
                await ctx.send("Disabled archiving for {}".format(channel.name))

    # https://leovoel.github.io/embed-visualizer/
    # https://cog-creators.github.io/discord-embed-sandbox/

    async def create_embed(self, message):
        "Create the correctly formatted embed for pinned messages"
        # Create embed using information from message
        embed_message = discord.Embed()
        embed_message.title = "{} posted".format(message.author.display_name)
        embed_message.description = message.content
        embed_message.timestamp = datetime.replace(message.created_at, tzinfo=timezone.utc)
        embed_message.set_author(name="{}#{}".format(message.author.name, message.author.discriminator), url="https://discord.com/users/{}".format(message.author.id), icon_url=message.author.avatar_url)
        embed_message.set_thumbnail(url=message.author.avatar_url)
        embed_message.add_field(name="Sources", value="{} | [Message]({})".format(message.channel.mention, message.jump_url), inline=False)
        attachment_count = len(message.attachments)
        if attachment_count != 0:
            if attachment_count == 1:
                attachment_text = "attachment"
            else:
                attachment_text = "attachments"
            embed_message.set_footer(text="{} {}".format(attachment_count, attachment_text))
        return embed_message

    async def return_attachments(self, message):
        "Returns attachments in a file list"
        files = []
        for attachment in message.attachments:
            attachment_bytes = BytesIO(await attachment.read())
            files.append(discord.File(fp=attachment_bytes, filename=attachment.filename))
        return files

    async def archive_pin(self, channel, message):
        """
        Archive pinned messages
        """
        archive_channels = await self.config.guild(channel.guild).archive_channels()
        async with self.config.guild(channel.guild).pin_history() as pin_history:
            if message.id not in pin_history:
                # Get list of attachments from message
                # Save list of attachment objects to a list of file like objects
                # Send embed with provided attachments
                embed_message = await self.create_embed(message)
                files = await self.return_attachments(message)
                for channel_id in archive_channels:
                    archive_channel = channel.guild.get_channel(channel_id)
                    await archive_channel.send(embed=embed_message, files=files)
                pin_history.append(message.id)
    # https://discordpy.readthedocs.io/en/latest/api.html#discord.on_guild_channel_pins_update
    # Monitored Events
    @commands.Cog.listener('on_guild_channel_pins_update') # Executes the below command when a channels pinned messages changes
    async def on_pin_update(self, channel, last_pin):
        monitored_channels = await self.config.guild(channel.guild).monitored_channels()
        if (channel.id in monitored_channels) and (last_pin != None):
            # Check if this pin hasn't already been archived
            last_pinned_message = (await channel.pins())[0]
            await self.archive_pin(channel, last_pinned_message)
<|endoftext|>


<commit_before>
import click
import numpy as np
import rasterio as rio
from scipy import sparse

import pyamg
from pyamg.relaxation.smoothing import change_smoothers

from geoblend import blend, matrix_from_mask_numba
from geoblend.utilities import get_mask
from geoblend.solver import create_multilevel_solver, load_multilevel_solver

@click.group()
def geoblend():
    pass

@click.command('poisson')
@click.argument('srcpath', type=click.Path(exists=True))
@click.argument('refpath', type=click.Path(exists=True))
@click.argument('dstpath', type=click.Path(exists=False))
@click.option("--matrix")
def poisson(srcpath, refpath, dstpath, matrix):
    """
    Poisson blend the source image against the reference image.
    """

    mask = get_mask(srcpath)
    indices = np.nonzero(mask)

    with rio.drivers():
        with rio.open(srcpath) as src, rio.open(refpath) as ref:

            profile = src.profile

            if matrix:
                levels = load_multilevel_solver(matrix)
                ml = pyamg.multilevel.multilevel_solver(levels, coarse_solver='pinv2')
                change_smoothers(ml, 'gauss_seidel', 'gauss_seidel')
            else:
                mat = matrix_from_mask_numba(mask)

                v = np.ones((mat.shape[0], 1))
                ml = pyamg.smoothed_aggregation_solver(mat, v, max_coarse=10)

            shape = src.shape

            with rio.open(dstpath, 'w', **profile) as dst:

                for bidx in range(1, 4):

                    source = src.read(bidx)
                    reference = ref.read(bidx)

                    arr = blend(source, reference, mask, ml)
                    dst.write_band(bidx, arr)

@click.command('create-solver')
@click.argument('srcpath', type=click.Path(exists=True))
@click.argument('dstpath', type=click.Path(exists=False))
def create_solver(srcpath, dstpath):
    """
    Create a multi-level solver and save to disk.
    """

    mask = get_mask(srcpath)
    indices = np.nonzero(mask)

    with rio.drivers():
        with rio.open(srcpath) as src:
            mat = matrix_from_mask_numba(mask)

    print mat.shape
    create_multilevel_solver(dstpath, mat)

geoblend.add_command(poisson)
geoblend.add_command(create_solver)

if __name__ == '__main__':
    geoblend()
<commit_msg>switch to cython impl<commit_after>
import click
import numpy as np
import rasterio as rio
from scipy import sparse

import pyamg
from pyamg.relaxation.smoothing import change_smoothers

from geoblend import blend
from geoblend.coefficients import matrix_from_mask
from geoblend.utilities import get_mask
from geoblend.solver import create_multilevel_solver, load_multilevel_solver

@click.group()
def geoblend():
    pass

@click.command('poisson')
@click.argument('srcpath', type=click.Path(exists=True))
@click.argument('refpath', type=click.Path(exists=True))
@click.argument('dstpath', type=click.Path(exists=False))
@click.option("--matrix")
def poisson(srcpath, refpath, dstpath, matrix):
    """
    Poisson blend the source image against the reference image.
    """

    mask = get_mask(srcpath)
    indices = np.nonzero(mask)

    with rio.drivers():
        with rio.open(srcpath) as src, rio.open(refpath) as ref:

            profile = src.profile

            if matrix:
                levels = load_multilevel_solver(matrix)
                ml = pyamg.multilevel.multilevel_solver(levels, coarse_solver='pinv2')
                change_smoothers(ml, 'gauss_seidel', 'gauss_seidel')
            else:
                mat = matrix_from_mask(mask)

                v = np.ones((mat.shape[0], 1))
                ml = pyamg.smoothed_aggregation_solver(mat, v, max_coarse=10)

            shape = src.shape

            with rio.open(dstpath, 'w', **profile) as dst:

                for bidx in range(1, 4):

                    source = src.read(bidx)
                    reference = ref.read(bidx)

                    arr = blend(source, reference, mask, ml)
                    dst.write_band(bidx, arr)

@click.command('create-solver')
@click.argument('srcpath', type=click.Path(exists=True))
@click.argument('dstpath', type=click.Path(exists=False))
def create_solver(srcpath, dstpath):
    """
    Create a multi-level solver and save to disk.
    """

    mask = get_mask(srcpath)
    indices = np.nonzero(mask)

    with rio.drivers():
        with rio.open(srcpath) as src:
            mat = matrix_from_mask(mask)

    create_multilevel_solver(dstpath, mat)

geoblend.add_command(poisson)
geoblend.add_command(create_solver)

if __name__ == '__main__':
    geoblend()
<|endoftext|>


<commit_before>print "Warning: Using prototype cython.inline code..."

import tempfile
import sys, os, re, inspect

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.dist import Distribution
from Cython.Distutils.extension import Extension
from Cython.Distutils import build_ext

from Cython.Compiler.Main import Context, CompilationOptions, default_options

from Cython.Compiler.ParseTreeTransforms import CythonTransform, SkipDeclarations, AnalyseDeclarationsTransform
from Cython.Compiler.TreeFragment import parse_from_strings

_code_cache = {}

class AllSymbols(CythonTransform, SkipDeclarations):
    def __init__(self):
        CythonTransform.__init__(self, None)
        self.names = set()
    def visit_NameNode(self, node):
        self.names.add(node.name)

def unbound_symbols(code, context=None):
    if context is None:
        context = Context([], default_options)
    from Cython.Compiler.ParseTreeTransforms import AnalyseDeclarationsTransform
    if isinstance(code, str):
        code = code.decode('ascii')
    tree = parse_from_strings('(tree fragment)', code)
    for phase in context.create_pipeline(pxd=False):
        if phase is None:
            continue
        tree = phase(tree)
        if isinstance(phase, AnalyseDeclarationsTransform):
            break
    symbol_collector = AllSymbols()
    symbol_collector(tree)
    unbound = []
    import __builtin__
    for name in symbol_collector.names:
        if not tree.scope.lookup(name) and not hasattr(__builtin__, name):
            unbound.append(name)
    return unbound

def get_type(arg, context=None):
    py_type = type(arg)
    if py_type in [list, tuple, dict, str]:
        return py_type.__name__
    elif py_type is float:
        return 'double'
    elif py_type is bool:
        return 'bint'
    elif py_type is int:
        return 'long'
    elif 'numpy' in sys.modules and isinstance(arg, sys.modules['numpy'].ndarray):
        return 'numpy.ndarray[numpy.%s_t, ndim=%s]' % (arg.dtype.name, arg.ndim)
    else:
        for base_type in py_type.mro():
            if base_type.__module__ == '__builtin__':
                return 'object'
            module = context.find_module(base_type.__module__, need_pxd=False)
            if module:
                entry = module.lookup(base_type.__name__)
                if entry.is_type:
                    return '%s.%s' % (base_type.__module__, base_type.__name__)
        return 'object'

# TODO: use locals/globals for unbound variables
def cython_inline(code,
                  types='aggressive',
                  lib_dir=os.path.expanduser('~/.cython/inline'),
                  include_dirs=['.'],
                  locals=None,
                  globals=None,
                  **kwds):
    ctx = Context(include_dirs, default_options)
    if locals is None:
        locals = inspect.currentframe().f_back.f_back.f_locals
    if globals is None:
        globals = inspect.currentframe().f_back.f_back.f_globals
    try:
        for symbol in unbound_symbols(code):
            if symbol in kwds:
                continue
            elif symbol in locals:
                kwds[symbol] = locals[symbol]
            elif symbol in globals:
                kwds[symbol] = globals[symbol]
            else:
                print "Couldn't find ", symbol
    except AssertionError:
        # Parsing from strings not fully supported (e.g. cimports).
        print "Could not parse code as a string (to extract unbound symbols)."
    arg_names = kwds.keys()
    arg_names.sort()
    arg_sigs = tuple([(get_type(kwds[arg], ctx), arg) for arg in arg_names])
    key = code, arg_sigs
    module = _code_cache.get(key)
    if not module:
        cimports = []
        qualified = re.compile(r'([.\w]+)[.]')
        for type, _ in arg_sigs:
            m = qualified.match(type)
            if m:
                cimports.append('\ncimport %s' % m.groups()[0])
        module_body, func_body = extract_func_code(code)
        params = ', '.join('%s %s' % a for a in arg_sigs)
        module_code = """
%(cimports)s
%(module_body)s
def __invoke(%(params)s):
%(func_body)s
        """ % {'cimports': '\n'.join(cimports), 'module_body': module_body, 'params': params, 'func_body': func_body }
#        print module_code
        _, pyx_file = tempfile.mkstemp('.pyx')
        open(pyx_file, 'w').write(module_code)
        module = "_" + hashlib.md5(code + str(arg_sigs)).hexdigest()
        extension = Extension(
            name = module,
            sources = [pyx_file],
            pyrex_include_dirs = include_dirs)
        build_extension = build_ext(Distribution())
        build_extension.finalize_options()
        build_extension.extensions = [extension]
        build_extension.build_temp = os.path.dirname(pyx_file)
        if lib_dir not in sys.path:
            sys.path.append(lib_dir)
        build_extension.build_lib  = lib_dir
        build_extension.run()
        _code_cache[key] = module
    arg_list = [kwds[arg] for arg in arg_names]
    return __import__(module).__invoke(*arg_list)

non_space = re.compile('[^ ]')
def strip_common_indent(lines):
    min_indent = None
    for line in lines:
        if not line:
            continue # empty
        indent = non_space.search(line).start()
        if indent == len(line):
            continue # blank
        elif line[indent] == '#':
            continue # comment
        elif min_indent is None or min_indent > indent:
            min_indent = indent
    for line in lines:
        if not line:
            continue
        indent = non_space.search(line).start()
        if indent == len(line):
            continue
        elif line[indent] == '#':
            yield line
        else:
            yield line[min_indent:]

module_statement = re.compile(r'^((cdef +(extern|class))|cimport|(from .+ cimport)|(from .+ import +[*]))')
def extract_func_code(code):
    module = []
    function = []
    # TODO: string literals, backslash
    current = function
    code = code.replace('\t', ' ')
    lines = strip_common_indent(code.split('\n'))
    for line in lines:
        if not line.startswith(' '):
            if module_statement.match(line):
                current = module
            else:
                current = function
        current.append(line)
    return '\n'.join(module), '    ' + '\n    '.join(function)
<commit_msg>Fix indent stripping.<commit_after>print "Warning: Using prototype cython.inline code..."

import tempfile
import sys, os, re, inspect

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.dist import Distribution
from Cython.Distutils.extension import Extension
from Cython.Distutils import build_ext

from Cython.Compiler.Main import Context, CompilationOptions, default_options

from Cython.Compiler.ParseTreeTransforms import CythonTransform, SkipDeclarations, AnalyseDeclarationsTransform
from Cython.Compiler.TreeFragment import parse_from_strings

_code_cache = {}

class AllSymbols(CythonTransform, SkipDeclarations):
    def __init__(self):
        CythonTransform.__init__(self, None)
        self.names = set()
    def visit_NameNode(self, node):
        self.names.add(node.name)

def unbound_symbols(code, context=None):
    if context is None:
        context = Context([], default_options)
    from Cython.Compiler.ParseTreeTransforms import AnalyseDeclarationsTransform
    if isinstance(code, str):
        code = code.decode('ascii')
    tree = parse_from_strings('(tree fragment)', code)
    for phase in context.create_pipeline(pxd=False):
        if phase is None:
            continue
        tree = phase(tree)
        if isinstance(phase, AnalyseDeclarationsTransform):
            break
    symbol_collector = AllSymbols()
    symbol_collector(tree)
    unbound = []
    import __builtin__
    for name in symbol_collector.names:
        if not tree.scope.lookup(name) and not hasattr(__builtin__, name):
            unbound.append(name)
    return unbound

def get_type(arg, context=None):
    py_type = type(arg)
    if py_type in [list, tuple, dict, str]:
        return py_type.__name__
    elif py_type is float:
        return 'double'
    elif py_type is bool:
        return 'bint'
    elif py_type is int:
        return 'long'
    elif 'numpy' in sys.modules and isinstance(arg, sys.modules['numpy'].ndarray):
        return 'numpy.ndarray[numpy.%s_t, ndim=%s]' % (arg.dtype.name, arg.ndim)
    else:
        for base_type in py_type.mro():
            if base_type.__module__ == '__builtin__':
                return 'object'
            module = context.find_module(base_type.__module__, need_pxd=False)
            if module:
                entry = module.lookup(base_type.__name__)
                if entry.is_type:
                    return '%s.%s' % (base_type.__module__, base_type.__name__)
        return 'object'

# TODO: use locals/globals for unbound variables
def cython_inline(code,
                  types='aggressive',
                  lib_dir=os.path.expanduser('~/.cython/inline'),
                  include_dirs=['.'],
                  locals=None,
                  globals=None,
                  **kwds):
    code = strip_common_indent(code)
    ctx = Context(include_dirs, default_options)
    if locals is None:
        locals = inspect.currentframe().f_back.f_back.f_locals
    if globals is None:
        globals = inspect.currentframe().f_back.f_back.f_globals
    try:
        for symbol in unbound_symbols(code):
            if symbol in kwds:
                continue
            elif symbol in locals:
                kwds[symbol] = locals[symbol]
            elif symbol in globals:
                kwds[symbol] = globals[symbol]
            else:
                print "Couldn't find ", symbol
    except AssertionError:
        # Parsing from strings not fully supported (e.g. cimports).
        print "Could not parse code as a string (to extract unbound symbols)."
    arg_names = kwds.keys()
    arg_names.sort()
    arg_sigs = tuple([(get_type(kwds[arg], ctx), arg) for arg in arg_names])
    key = code, arg_sigs
    module = _code_cache.get(key)
    if not module:
        cimports = []
        qualified = re.compile(r'([.\w]+)[.]')
        for type, _ in arg_sigs:
            m = qualified.match(type)
            if m:
                cimports.append('\ncimport %s' % m.groups()[0])
        module_body, func_body = extract_func_code(code)
        params = ', '.join('%s %s' % a for a in arg_sigs)
        module_code = """
%(cimports)s
%(module_body)s
def __invoke(%(params)s):
%(func_body)s
        """ % {'cimports': '\n'.join(cimports), 'module_body': module_body, 'params': params, 'func_body': func_body }
#        print module_code
        _, pyx_file = tempfile.mkstemp('.pyx')
        open(pyx_file, 'w').write(module_code)
        module = "_" + hashlib.md5(code + str(arg_sigs)).hexdigest()
        extension = Extension(
            name = module,
            sources = [pyx_file],
            pyrex_include_dirs = include_dirs)
        build_extension = build_ext(Distribution())
        build_extension.finalize_options()
        build_extension.extensions = [extension]
        build_extension.build_temp = os.path.dirname(pyx_file)
        if lib_dir not in sys.path:
            sys.path.append(lib_dir)
        build_extension.build_lib  = lib_dir
        build_extension.run()
        _code_cache[key] = module
    arg_list = [kwds[arg] for arg in arg_names]
    return __import__(module).__invoke(*arg_list)

non_space = re.compile('[^ ]')
def strip_common_indent(code):
    min_indent = None
    lines = code.split('\n')
    for line in lines:
        match = non_space.search(line)
        if not match:
            continue # blank
        indent = match.start()
        if line[indent] == '#':
            continue # comment
        elif min_indent is None or min_indent > indent:
            min_indent = indent
    for ix, line in enumerate(lines):
        match = non_space.search(line)
        if not match or line[indent] == '#':
            continue
        else:
            lines[ix] = line[min_indent:]
    return '\n'.join(lines)

module_statement = re.compile(r'^((cdef +(extern|class))|cimport|(from .+ cimport)|(from .+ import +[*]))')
def extract_func_code(code):
    module = []
    function = []
    # TODO: string literals, backslash
    current = function
    code = code.replace('\t', ' ')
    lines = code.split('\n')
    for line in lines:
        if not line.startswith(' '):
            if module_statement.match(line):
                current = module
            else:
                current = function
        current.append(line)
    return '\n'.join(module), '    ' + '\n    '.join(function)
<|endoftext|>