<commit_before>
from collections import defaultdict

from django.contrib.auth.decorators import login_required
from django.contrib.messages import info
from django.core.urlresolvers import get_callable, reverse
from django.http import Http404, HttpResponse
from django.shortcuts import get_object_or_404, redirect
from django.template import RequestContext
from django.template.defaultfilters import slugify
from django.template.loader import get_template
from django.utils import simplejson
from django.utils.translation import ugettext as _
from django.views.decorators.cache import never_cache

from mezzanine.conf import settings
from mezzanine.utils.importing import import_dotted_path
from mezzanine.utils.views import render, set_cookie, paginate

from cartridge.shop import checkout
from cartridge.shop.forms import AddProductForm, DiscountForm, CartItemFormSet
from cartridge.shop.models import Product, ProductVariation, Order, OrderItem
from cartridge.shop.models import DiscountCode
from cartridge.shop.utils import recalculate_discount, sign

# Set up checkout handlers.
handler = lambda s: import_dotted_path(s) if s else lambda *args: None
billship_handler = handler(settings.SHOP_HANDLER_BILLING_SHIPPING)
tax_handler = handler(settings.SHOP_HANDLER_TAX)
payment_handler = handler(settings.SHOP_HANDLER_PAYMENT)
order_handler = handler(settings.SHOP_HANDLER_ORDER)

def product(request, slug, template="shop/product.html"):
    """
    Display a product - convert the product variations to JSON as well as
    handling adding the product to either the cart or the wishlist.
    """
    published_products = Product.objects.published(for_user=request.user)
    product = get_object_or_404(published_products, slug=slug)
    fields = [f.name for f in ProductVariation.option_fields()]
    variations = product.variations.all()
    variations_json = simplejson.dumps([dict([(f, getattr(v, f))
                                        for f in fields + ["sku", "image_id"]])
                                        for v in variations])
    to_cart = (request.method == "POST" and
               request.POST.get("add_wishlist") is None)
    initial_data = {}
    if variations:
        initial_data = dict([(f, getattr(variations[0], f)) for f in fields])
    initial_data["quantity"] = 1
    add_product_form = AddProductForm(request.POST or None, product=product,
                                      initial=initial_data, to_cart=to_cart)
    if request.method == "POST":
        if add_product_form.is_valid():
            if to_cart:
                quantity = add_product_form.cleaned_data["quantity"]
                request.cart.add_item(add_product_form.variation, quantity)
                recalculate_discount(request)
                info(request, _("Item added to cart"))
                return redirect("shop_cart")
            else:
                skus = request.wishlist
                sku = add_product_form.variation.sku
                if sku not in skus:
                    skus.append(sku)
                info(request, _("Item added to wishlist"))
                response = redirect("shop_wishlist")
                set_cookie(response, "wishlist", ",".join(skus))
                return response
    context = {
        "product": product,
        "editable_obj": product,
        "images": product.images.all(),
        "variations": variations,
        "variations_json": variations_json,
        "has_available_variations": any([v.has_price() for v in variations]),
        "related_products": product.related_products.published(
                                                      for_user=request.user),
        "add_product_form": add_product_form
    }
    return render(request, template, context)

@never_cache
def wishlist(request, template="shop/wishlist.html"):
    """
    Display the wishlist and handle removing items from the wishlist and
    adding them to the cart.
    """

    skus = request.wishlist
    error = None
    if request.method == "POST":
        to_cart = request.POST.get("add_cart")
        add_product_form = AddProductForm(request.POST or None,
                                          to_cart=to_cart)
        if to_cart:
            if add_product_form.is_valid():
                request.cart.add_item(add_product_form.variation, 1)
                recalculate_discount(request)
                message = _("Item added to cart")
                url = "shop_cart"
            else:
                error = add_product_form.errors.values()[0]
        else:
            message = _("Item removed from wishlist")
            url = "shop_wishlist"
        sku = request.POST.get("sku")
        if sku in skus:
            skus.remove(sku)
        if not error:
            info(request, message)
            response = redirect(url)
            set_cookie(response, "wishlist", ",".join(skus))
            return response

    # Remove skus from the cookie that no longer exist.
    published_products = Product.objects.published(for_user=request.user)
    f = {"product__in": published_products, "sku__in": skus}
    wishlist = ProductVariation.objects.filter(**f).select_related(depth=1)
    wishlist = sorted(wishlist, key=lambda v: skus.index(v.sku))
    context = {"wishlist_items": wishlist, "error": error}
    response = render(request, template, context)
    if len(wishlist) < len(skus):
        skus = [variation.sku for variation in wishlist]
        set_cookie(response, "wishlist", ",".join(skus))
    return response

@never_cache
def cart(request, template="shop/cart.html"):
    """
    Display cart and handle removing items from the cart.
    """
    cart_formset = CartItemFormSet(instance=request.cart)
    discount_form = DiscountForm(request, request.POST or None)
    if request.method == "POST":
        valid = True
        if request.POST.get("update_cart"):
            valid = request.cart.has_items()
            if not valid:
                # Session timed out.
                info(request, _("Your cart has expired"))
            else:
                cart_formset = CartItemFormSet(request.POST,
                                               instance=request.cart)
                valid = cart_formset.is_valid()
                if valid:
                    cart_formset.save()
                    recalculate_discount(request)
                    info(request, _("Cart updated"))
        else:
            valid = discount_form.is_valid()
            if valid:
                discount_form.set_discount()
        if valid:
            return redirect("shop_cart")
    context = {"cart_formset": cart_formset}
    settings.use_editable()
    if (settings.SHOP_DISCOUNT_FIELD_IN_CART and
        DiscountCode.objects.active().count() > 0):
        context["discount_form"] = discount_form
    return render(request, template, context)

@never_cache
def checkout_steps(request):
    """
    Display the order form and handle processing of each step.
    """

    # Do the authentication check here rather than using standard
    # login_required decorator. This means we can check for a custom
    # LOGIN_URL and fall back to our own login view.
    authenticated = request.user.is_authenticated()
    if settings.SHOP_CHECKOUT_ACCOUNT_REQUIRED and not authenticated:
        url = "%s?next=%s" % (settings.LOGIN_URL, reverse("shop_checkout"))
        return redirect(url)

    # Determine the Form class to use during the checkout process
    form_class = get_callable(settings.SHOP_CHECKOUT_FORM_CLASS)

    step = int(request.POST.get("step", checkout.CHECKOUT_STEP_FIRST))
    initial = checkout.initial_order_data(request, form_class)
    form = form_class(request, step, initial=initial)
    data = request.POST
    checkout_errors = []

    if request.POST.get("back") is not None:
        # Back button in the form was pressed - load the order form
        # for the previous step and maintain the field values entered.
        step -= 1
        form = form_class(request, step, initial=initial)
    elif request.method == "POST" and request.cart.has_items():
        form = form_class(request, step, initial=initial, data=data)
        if form.is_valid():
            # Copy the current form fields to the session so that
            # they're maintained if the customer leaves the checkout
            # process, but remove sensitive fields from the session
            # such as the credit card fields so that they're never
            # stored anywhere.
            request.session["order"] = dict(form.cleaned_data)
            sensitive_card_fields = ("card_number", "card_expiry_month",
                                     "card_expiry_year", "card_ccv")
            for field in sensitive_card_fields:
                if field in request.session["order"]:
                    del request.session["order"][field]

            # FIRST CHECKOUT STEP - handle shipping and discount code.
            if step == checkout.CHECKOUT_STEP_FIRST:
                try:
                    billship_handler(request, form)
                    tax_handler(request, form)
                except checkout.CheckoutError, e:
                    checkout_errors.append(e)
                form.set_discount()

            # FINAL CHECKOUT STEP - handle payment and process order.
            if step == checkout.CHECKOUT_STEP_LAST and not checkout_errors:
                # Create and save the inital order object so that
                # the payment handler has access to all of the order
                # fields. If there is a payment error then delete the
                # order, otherwise remove the cart items from stock
                # and send the order reciept email.
                order = form.save(commit=False)
                order.setup(request)
                # Try payment.
                try:
                    transaction_id = payment_handler(request, form, order)
                except checkout.CheckoutError, e:
                    # Error in payment handler.
                    order.delete()
                    checkout_errors.append(e)
                    if settings.SHOP_CHECKOUT_STEPS_CONFIRMATION:
                        step -= 1
                else:
                    # Finalize order - ``order.complete()`` performs
                    # final cleanup of session and cart.
                    # ``order_handler()`` can be defined by the
                    # developer to implement custom order processing.
                    # Then send the order email to the customer.
                    order.transaction_id = transaction_id
                    order.complete(request)
                    order_handler(request, form, order)
                    checkout.send_order_email(request, order)
                    # Set the cookie for remembering address details
                    # if the "remember" checkbox was checked.
                    response = redirect("shop_complete")
                    if form.cleaned_data.get("remember") is not None:
                        remembered = "%s:%s" % (sign(order.key), order.key)
                        set_cookie(response, "remember", remembered,
                                   secure=request.is_secure())
                    else:
                        response.delete_cookie("remember")
                    return response

            # If any checkout errors, assign them to a new form and
            # re-run is_valid. If valid, then set form to the next step.
            form = form_class(request, step, initial=initial, data=data,
                             errors=checkout_errors)
            if form.is_valid():
                step += 1
                form = form_class(request, step, initial=initial)

    step_vars = checkout.CHECKOUT_STEPS[step - 1]
    template = "shop/%s.html" % step_vars["template"]
    CHECKOUT_STEP_FIRST = step == checkout.CHECKOUT_STEP_FIRST
    context = {"form": form, "CHECKOUT_STEP_FIRST": CHECKOUT_STEP_FIRST,
               "step_title": step_vars["title"], "step_url": step_vars["url"],
               "steps": checkout.CHECKOUT_STEPS, "step": step}
    return render(request, template, context)

@never_cache
def complete(request, template="shop/complete.html"):
    """
    Redirected to once an order is complete - pass the order object
    for tracking items via Google Anayltics, and displaying in
    the template if required.
    """
    try:
        order = Order.objects.from_request(request)
    except Order.DoesNotExist:
        raise Http404
    items = order.items.all()
    # Assign product names to each of the items since they're not
    # stored.
    skus = [item.sku for item in items]
    variations = ProductVariation.objects.filter(sku__in=skus)
    names = {}
    for variation in variations.select_related(depth=1):
        names[variation.sku] = variation.product.title
    for i, item in enumerate(items):
        setattr(items[i], "name", names[item.sku])
    context = {"order": order, "items": items,
               "steps": checkout.CHECKOUT_STEPS}
    return render(request, template, context)

def invoice(request, order_id, template="shop/order_invoice.html"):
    """
    Display a plain text invoice for the given order. The order must
    belong to the user which is checked via session or ID if
    authenticated, or if the current user is staff.
    """
    lookup = {"id": order_id}
    if not request.user.is_authenticated():
        lookup["key"] = request.session.session_key
    elif not request.user.is_staff:
        lookup["user_id"] = request.user.id
    order = get_object_or_404(Order, **lookup)
    context = {"order": order}
    context.update(order.details_as_dict())
    context = RequestContext(request, context)
    if request.GET.get("format") == "pdf":
        response = HttpResponse(mimetype="application/pdf")
        name = slugify("%s-invoice-%s" % (settings.SITE_TITLE, order.id))
        response["Content-Disposition"] = "attachment; filename=%s.pdf" % name
        html = get_template(template).render(context)
        import ho.pisa
        ho.pisa.CreatePDF(html, response)
        return response
    return render(request, template, context)

@login_required
def order_history(request, template="shop/order_history.html"):
    """
    Display a list of the currently logged-in user's past orders.
    """
    all_orders = Order.objects.filter(user_id=request.user.id)
    orders = paginate(all_orders.order_by('-time'),
                      request.GET.get("page", 1),
                      settings.SHOP_PER_PAGE_CATEGORY,
                      settings.MAX_PAGING_LINKS)
    # Add the total quantity to each order - this can probably be
    # replaced with fetch_related and Sum when we drop Django 1.3
    order_quantities = defaultdict(int)
    for item in OrderItem.objects.filter(order__user_id=request.user.id):
        order_quantities[item.order_id] += item.quantity
    for order in orders.object_list:
        setattr(order, "quantity_total", order_quantities[order.id])
    context = {"orders": orders}
    return render(request, template, context)
<commit_msg>Comment typos & indentation fix<commit_after>from collections import defaultdict

from django.contrib.auth.decorators import login_required
from django.contrib.messages import info
from django.core.urlresolvers import get_callable, reverse
from django.http import Http404, HttpResponse
from django.shortcuts import get_object_or_404, redirect
from django.template import RequestContext
from django.template.defaultfilters import slugify
from django.template.loader import get_template
from django.utils import simplejson
from django.utils.translation import ugettext as _
from django.views.decorators.cache import never_cache

from mezzanine.conf import settings
from mezzanine.utils.importing import import_dotted_path
from mezzanine.utils.views import render, set_cookie, paginate

from cartridge.shop import checkout
from cartridge.shop.forms import AddProductForm, DiscountForm, CartItemFormSet
from cartridge.shop.models import Product, ProductVariation, Order, OrderItem
from cartridge.shop.models import DiscountCode
from cartridge.shop.utils import recalculate_discount, sign

# Set up checkout handlers.
handler = lambda s: import_dotted_path(s) if s else lambda *args: None
billship_handler = handler(settings.SHOP_HANDLER_BILLING_SHIPPING)
tax_handler = handler(settings.SHOP_HANDLER_TAX)
payment_handler = handler(settings.SHOP_HANDLER_PAYMENT)
order_handler = handler(settings.SHOP_HANDLER_ORDER)

def product(request, slug, template="shop/product.html"):
    """
    Display a product - convert the product variations to JSON as well as
    handling adding the product to either the cart or the wishlist.
    """
    published_products = Product.objects.published(for_user=request.user)
    product = get_object_or_404(published_products, slug=slug)
    fields = [f.name for f in ProductVariation.option_fields()]
    variations = product.variations.all()
    variations_json = simplejson.dumps([dict([(f, getattr(v, f))
                                        for f in fields + ["sku", "image_id"]])
                                        for v in variations])
    to_cart = (request.method == "POST" and
               request.POST.get("add_wishlist") is None)
    initial_data = {}
    if variations:
        initial_data = dict([(f, getattr(variations[0], f)) for f in fields])
    initial_data["quantity"] = 1
    add_product_form = AddProductForm(request.POST or None, product=product,
                                      initial=initial_data, to_cart=to_cart)
    if request.method == "POST":
        if add_product_form.is_valid():
            if to_cart:
                quantity = add_product_form.cleaned_data["quantity"]
                request.cart.add_item(add_product_form.variation, quantity)
                recalculate_discount(request)
                info(request, _("Item added to cart"))
                return redirect("shop_cart")
            else:
                skus = request.wishlist
                sku = add_product_form.variation.sku
                if sku not in skus:
                    skus.append(sku)
                info(request, _("Item added to wishlist"))
                response = redirect("shop_wishlist")
                set_cookie(response, "wishlist", ",".join(skus))
                return response
    context = {
        "product": product,
        "editable_obj": product,
        "images": product.images.all(),
        "variations": variations,
        "variations_json": variations_json,
        "has_available_variations": any([v.has_price() for v in variations]),
        "related_products": product.related_products.published(
                                                      for_user=request.user),
        "add_product_form": add_product_form
    }
    return render(request, template, context)

@never_cache
def wishlist(request, template="shop/wishlist.html"):
    """
    Display the wishlist and handle removing items from the wishlist and
    adding them to the cart.
    """

    skus = request.wishlist
    error = None
    if request.method == "POST":
        to_cart = request.POST.get("add_cart")
        add_product_form = AddProductForm(request.POST or None,
                                          to_cart=to_cart)
        if to_cart:
            if add_product_form.is_valid():
                request.cart.add_item(add_product_form.variation, 1)
                recalculate_discount(request)
                message = _("Item added to cart")
                url = "shop_cart"
            else:
                error = add_product_form.errors.values()[0]
        else:
            message = _("Item removed from wishlist")
            url = "shop_wishlist"
        sku = request.POST.get("sku")
        if sku in skus:
            skus.remove(sku)
        if not error:
            info(request, message)
            response = redirect(url)
            set_cookie(response, "wishlist", ",".join(skus))
            return response

    # Remove skus from the cookie that no longer exist.
    published_products = Product.objects.published(for_user=request.user)
    f = {"product__in": published_products, "sku__in": skus}
    wishlist = ProductVariation.objects.filter(**f).select_related(depth=1)
    wishlist = sorted(wishlist, key=lambda v: skus.index(v.sku))
    context = {"wishlist_items": wishlist, "error": error}
    response = render(request, template, context)
    if len(wishlist) < len(skus):
        skus = [variation.sku for variation in wishlist]
        set_cookie(response, "wishlist", ",".join(skus))
    return response

@never_cache
def cart(request, template="shop/cart.html"):
    """
    Display cart and handle removing items from the cart.
    """
    cart_formset = CartItemFormSet(instance=request.cart)
    discount_form = DiscountForm(request, request.POST or None)
    if request.method == "POST":
        valid = True
        if request.POST.get("update_cart"):
            valid = request.cart.has_items()
            if not valid:
                # Session timed out.
                info(request, _("Your cart has expired"))
            else:
                cart_formset = CartItemFormSet(request.POST,
                                               instance=request.cart)
                valid = cart_formset.is_valid()
                if valid:
                    cart_formset.save()
                    recalculate_discount(request)
                    info(request, _("Cart updated"))
        else:
            valid = discount_form.is_valid()
            if valid:
                discount_form.set_discount()
        if valid:
            return redirect("shop_cart")
    context = {"cart_formset": cart_formset}
    settings.use_editable()
    if (settings.SHOP_DISCOUNT_FIELD_IN_CART and
        DiscountCode.objects.active().count() > 0):
        context["discount_form"] = discount_form
    return render(request, template, context)

@never_cache
def checkout_steps(request):
    """
    Display the order form and handle processing of each step.
    """

    # Do the authentication check here rather than using standard
    # login_required decorator. This means we can check for a custom
    # LOGIN_URL and fall back to our own login view.
    authenticated = request.user.is_authenticated()
    if settings.SHOP_CHECKOUT_ACCOUNT_REQUIRED and not authenticated:
        url = "%s?next=%s" % (settings.LOGIN_URL, reverse("shop_checkout"))
        return redirect(url)

    # Determine the Form class to use during the checkout process
    form_class = get_callable(settings.SHOP_CHECKOUT_FORM_CLASS)

    step = int(request.POST.get("step", checkout.CHECKOUT_STEP_FIRST))
    initial = checkout.initial_order_data(request, form_class)
    form = form_class(request, step, initial=initial)
    data = request.POST
    checkout_errors = []

    if request.POST.get("back") is not None:
        # Back button in the form was pressed - load the order form
        # for the previous step and maintain the field values entered.
        step -= 1
        form = form_class(request, step, initial=initial)
    elif request.method == "POST" and request.cart.has_items():
        form = form_class(request, step, initial=initial, data=data)
        if form.is_valid():
            # Copy the current form fields to the session so that
            # they're maintained if the customer leaves the checkout
            # process, but remove sensitive fields from the session
            # such as the credit card fields so that they're never
            # stored anywhere.
            request.session["order"] = dict(form.cleaned_data)
            sensitive_card_fields = ("card_number", "card_expiry_month",
                                     "card_expiry_year", "card_ccv")
            for field in sensitive_card_fields:
                if field in request.session["order"]:
                    del request.session["order"][field]

            # FIRST CHECKOUT STEP - handle shipping and discount code.
            if step == checkout.CHECKOUT_STEP_FIRST:
                try:
                    billship_handler(request, form)
                    tax_handler(request, form)
                except checkout.CheckoutError, e:
                    checkout_errors.append(e)
                form.set_discount()

            # FINAL CHECKOUT STEP - handle payment and process order.
            if step == checkout.CHECKOUT_STEP_LAST and not checkout_errors:
                # Create and save the initial order object so that
                # the payment handler has access to all of the order
                # fields. If there is a payment error then delete the
                # order, otherwise remove the cart items from stock
                # and send the order receipt email.
                order = form.save(commit=False)
                order.setup(request)
                # Try payment.
                try:
                    transaction_id = payment_handler(request, form, order)
                except checkout.CheckoutError, e:
                    # Error in payment handler.
                    order.delete()
                    checkout_errors.append(e)
                    if settings.SHOP_CHECKOUT_STEPS_CONFIRMATION:
                        step -= 1
                else:
                    # Finalize order - ``order.complete()`` performs
                    # final cleanup of session and cart.
                    # ``order_handler()`` can be defined by the
                    # developer to implement custom order processing.
                    # Then send the order email to the customer.
                    order.transaction_id = transaction_id
                    order.complete(request)
                    order_handler(request, form, order)
                    checkout.send_order_email(request, order)
                    # Set the cookie for remembering address details
                    # if the "remember" checkbox was checked.
                    response = redirect("shop_complete")
                    if form.cleaned_data.get("remember") is not None:
                        remembered = "%s:%s" % (sign(order.key), order.key)
                        set_cookie(response, "remember", remembered,
                                   secure=request.is_secure())
                    else:
                        response.delete_cookie("remember")
                    return response

            # If any checkout errors, assign them to a new form and
            # re-run is_valid. If valid, then set form to the next step.
            form = form_class(request, step, initial=initial, data=data,
                              errors=checkout_errors)
            if form.is_valid():
                step += 1
                form = form_class(request, step, initial=initial)

    step_vars = checkout.CHECKOUT_STEPS[step - 1]
    template = "shop/%s.html" % step_vars["template"]
    CHECKOUT_STEP_FIRST = step == checkout.CHECKOUT_STEP_FIRST
    context = {"form": form, "CHECKOUT_STEP_FIRST": CHECKOUT_STEP_FIRST,
               "step_title": step_vars["title"], "step_url": step_vars["url"],
               "steps": checkout.CHECKOUT_STEPS, "step": step}
    return render(request, template, context)

@never_cache
def complete(request, template="shop/complete.html"):
    """
    Redirected to once an order is complete - pass the order object
    for tracking items via Google Anayltics, and displaying in
    the template if required.
    """
    try:
        order = Order.objects.from_request(request)
    except Order.DoesNotExist:
        raise Http404
    items = order.items.all()
    # Assign product names to each of the items since they're not
    # stored.
    skus = [item.sku for item in items]
    variations = ProductVariation.objects.filter(sku__in=skus)
    names = {}
    for variation in variations.select_related(depth=1):
        names[variation.sku] = variation.product.title
    for i, item in enumerate(items):
        setattr(items[i], "name", names[item.sku])
    context = {"order": order, "items": items,
               "steps": checkout.CHECKOUT_STEPS}
    return render(request, template, context)

def invoice(request, order_id, template="shop/order_invoice.html"):
    """
    Display a plain text invoice for the given order. The order must
    belong to the user which is checked via session or ID if
    authenticated, or if the current user is staff.
    """
    lookup = {"id": order_id}
    if not request.user.is_authenticated():
        lookup["key"] = request.session.session_key
    elif not request.user.is_staff:
        lookup["user_id"] = request.user.id
    order = get_object_or_404(Order, **lookup)
    context = {"order": order}
    context.update(order.details_as_dict())
    context = RequestContext(request, context)
    if request.GET.get("format") == "pdf":
        response = HttpResponse(mimetype="application/pdf")
        name = slugify("%s-invoice-%s" % (settings.SITE_TITLE, order.id))
        response["Content-Disposition"] = "attachment; filename=%s.pdf" % name
        html = get_template(template).render(context)
        import ho.pisa
        ho.pisa.CreatePDF(html, response)
        return response
    return render(request, template, context)

@login_required
def order_history(request, template="shop/order_history.html"):
    """
    Display a list of the currently logged-in user's past orders.
    """
    all_orders = Order.objects.filter(user_id=request.user.id)
    orders = paginate(all_orders.order_by('-time'),
                      request.GET.get("page", 1),
                      settings.SHOP_PER_PAGE_CATEGORY,
                      settings.MAX_PAGING_LINKS)
    # Add the total quantity to each order - this can probably be
    # replaced with fetch_related and Sum when we drop Django 1.3
    order_quantities = defaultdict(int)
    for item in OrderItem.objects.filter(order__user_id=request.user.id):
        order_quantities[item.order_id] += item.quantity
    for order in orders.object_list:
        setattr(order, "quantity_total", order_quantities[order.id])
    context = {"orders": orders}
    return render(request, template, context)
<|endoftext|>


<commit_before>from collections import defaultdict, deque
import errno
from functools import partial, wraps
import logging
import socket
from threading import RLock, Event, Lock, Thread
import traceback
from Queue import Queue

import pyev

from cassandra import ConsistencyLevel
from cassandra.marshal import (int8_unpack, int32_unpack)
from cassandra.decoder import (OptionsMessage, ReadyMessage, AuthenticateMessage,
                               StartupMessage, ErrorMessage, CredentialsMessage,
                               QueryMessage, ResultMessage, decode_response)

log = logging.getLogger(__name__)

locally_supported_compressions = {}
try:
    import snappy
except ImportError:
    pass
else:
    # work around apparently buggy snappy decompress
    def decompress(byts):
        if byts == '\x00':
            return ''
        return snappy.decompress(byts)
    locally_supported_compressions['snappy'] = (snappy.compress, decompress)

MAX_STREAM_PER_CONNECTION = 128

PROTOCOL_VERSION = 0x01
PROTOCOL_VERSION_MASK = 0x7f

HEADER_DIRECTION_FROM_CLIENT = 0x00
HEADER_DIRECTION_TO_CLIENT = 0x80
HEADER_DIRECTION_MASK = 0x80

NONBLOCKING = (errno.EAGAIN, errno.EWOULDBLOCK)

class ConnectionException(Exception):

    def __init__(self, message, host=None):
        Exception.__init__(self, message)
        self.host = host

class ConnectionBusy(Exception):
    pass

class ProgrammingError(Exception):
    pass

class ProtocolError(Exception):
    pass

_loop = pyev.default_loop(pyev.EVBACKEND_SELECT)

_loop_notifier = _loop.async(lambda *a, **kw: None)
_loop_notifier.start()

# prevent _loop_notifier from keeping the loop from returning
_loop.unref()

_loop_started = None
_loop_lock = Lock()

def _run_loop():
    while True:
        end_condition = _loop.start()
        # there are still active watchers, no deadlock
        with _loop_lock:
            if end_condition:
                log.debug("Restarting event loop")
                continue
            else:
                # all Connections have been closed, no active watchers
                log.debug("All Connections currently closed, event loop ended")
                global _loop_started
                _loop_started = False
                break

def _start_loop():
    global _loop_started
    should_start = False
    with _loop_lock:
        if not _loop_started:
            log.debug("Starting pyev event loop")
            _loop_started = True
            should_start = True

    if should_start:
        t = Thread(target=_run_loop, name="async_event_loop")
        t.daemon = True
        t.start()

    return should_start

def defunct_on_error(f):

    @wraps(f)
    def wrapper(self, *args, **kwargs):
        try:
            return f(self, *args, **kwargs)
        except Exception, exc:
            self.defunct(exc)

    return wrapper

class Connection(object):

    in_buffer_size = 4096
    out_buffer_size = 4096

    cql_version = None

    keyspace = None
    compression = True
    compressor = None
    decompressor = None

    last_error = None
    in_flight = 0
    is_defunct = False
    is_closed = False
    lock = None

    _buf = ""
    _total_reqd_bytes = 0
    _read_watcher = None
    _write_watcher = None
    _socket = None

    @classmethod
    def factory(cls, *args, **kwargs):
        conn = cls(*args, **kwargs)
        conn.connected_event.wait()
        if conn.last_error:
            raise conn.last_error
        else:
            return conn

    def __init__(self, host='127.0.0.1', port=9042, credentials=None, sockopts=None, compression=True, cql_version=None):
        self.host = host
        self.port = port
        self.credentials = credentials
        self.compression = compression
        self.cql_version = cql_version

        self.connected_event = Event()

        self._id_queue = Queue(MAX_STREAM_PER_CONNECTION)
        for i in range(MAX_STREAM_PER_CONNECTION):
            self._id_queue.put_nowait(i)

        self._callbacks = {}
        self._push_watchers = defaultdict(set)
        self.lock = RLock()
        self.id_lock = Lock()

        self.deque = deque()

        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.connect((host, port))
        self._socket.setblocking(0)

        if sockopts:
            for args in sockopts:
                self._socket.setsockopt(*args)

        self._read_watcher = pyev.Io(self._socket._sock, pyev.EV_READ, _loop, self.handle_read)
        self._write_watcher = pyev.Io(self._socket._sock, pyev.EV_WRITE, _loop, self.handle_write)
        with _loop_lock:
            self._read_watcher.start()
            self._write_watcher.start()

        log.debug("Sending initial options message for new Connection to %s" % (host,))
        self.send_msg(OptionsMessage(), self._handle_options_response)

        # start the global event loop if needed
        if not _start_loop():
            # if the loop was already started, notify it
            with _loop_lock:
                _loop_notifier.send()

    def close(self):
        with self.lock:
            if self.is_closed:
                return
            self.is_closed = True

        log.debug("Closing connection to %s" % (self.host,))
        if self._read_watcher:
            self._read_watcher.stop()
        if self._write_watcher:
            self._write_watcher.stop()
        self._socket.close()
        with _loop_lock:
            _loop_notifier.send()

        # don't leave in-progress operations hanging
        if not self.is_defunct:
            self._error_all_callbacks(
                ConnectionException("Connection to %s was closed" % self.host))

    def __del__(self):
        self.close()

    def defunct(self, exc):
        log.debug("Defuncting connection to %s: %s %s" %
                  (self.host, exc, traceback.format_exc(exc)))
        self.last_error = exc
        self.is_defunct = True
        self._error_all_callbacks(exc)
        self.connected_event.set()
        return exc

    def _error_all_callbacks(self, exc):
        for cb in self._callbacks.values():
            cb(exc)

    def handle_write(self, watcher, revents):
        try:
            next_msg = self.deque.popleft()
        except IndexError:
            self._write_watcher.stop()
            return

        try:
            sent = self._socket.send(next_msg)
        except socket.error, err:
            if (err.args[0] in NONBLOCKING):
                self.deque.appendleft(next_msg)
            else:
                self.defunct(err)
            return
        else:
            if sent < len(next_msg):
                self.deque.appendleft(next_msg[sent:])

            if not self.deque:
                self._write_watcher.stop()

    def handle_read(self, watcher, revents):
        try:
            buf = self._socket.recv(self.in_buffer_size)
        except socket.error, err:
            if err.args[0] not in NONBLOCKING:
                self.defunct(err)
            return

        if buf:
            self._buf += buf
            while True:
                if len(self._buf) < 8:
                    # we don't have a complete header yet
                    break
                elif self._total_reqd_bytes and len(self._buf) < self._total_reqd_bytes:
                    # we already saw a header, but we don't have a complete message yet
                    break
                else:
                    body_len = int32_unpack(self._buf[4:8])
                    if len(self._buf) - 8 >= body_len:
                        msg = self._buf[:8 + body_len]
                        self._buf = self._buf[8 + body_len:]
                        self._total_reqd_bytes = 0
                        self.process_msg(msg, body_len)
                    else:
                        self._total_reqd_bytes = body_len + 8
        else:
            logging.debug("connection closed by server")
            self.close()

    @defunct_on_error
    def process_msg(self, msg, body_len):
        version, flags, stream_id, opcode = map(int8_unpack, msg[:4])
        if stream_id < 0:
            callback = None
        else:
            callback = self._callbacks.pop(stream_id)
            self._id_queue.put_nowait(stream_id)

        # check that the protocol version is supported
        given_version = version & PROTOCOL_VERSION_MASK
        if given_version != PROTOCOL_VERSION:
            raise ProtocolError("Unsupported CQL protocol version: %d" % given_version)

        # check that the header direction is correct
        if version & HEADER_DIRECTION_MASK != HEADER_DIRECTION_TO_CLIENT:
            raise ProtocolError(
                "Header direction in response is incorrect; opcode %04x, stream id %r"
                % (opcode, stream_id))

        if body_len > 0:
            body = msg[8:]
        elif body_len == 0:
            body = ""
        else:
            raise ProtocolError("Got negative body length: %r" % body_len)

        response = decode_response(stream_id, flags, opcode, body, self.decompressor)

        try:
            if stream_id < 0:
                self.handle_pushed(response)
            elif callback is not None:
                callback(response)
        except:
            log.exception("Callback handler errored, ignoring:")

    def handle_pushed(self, response):
        for cb in self._push_watchers[response.type]:
            try:
                cb(response)
            except:
                log.exception("Pushed event handler errored, ignoring:")

    def push(self, data):
        sabs = self.out_buffer_size
        if len(data) > sabs:
            chunks = []
            for i in xrange(0, len(data), sabs):
                chunks.append(data[i:i + sabs])
        else:
            chunks = [data]

        with self.lock:
            self.deque.extend(chunks)

            if not self._write_watcher.active:
                with _loop_lock:
                    self._write_watcher.start()
                    _loop_notifier.send()

    def send_msg(self, msg, cb):
        if self.is_defunct:
            raise ConnectionException("Connection to %s is defunct" % self.host)
        elif self.is_closed:
            raise ConnectionException("Connection to %s is closed" % self.host)

        try:
            request_id = self._id_queue.get_nowait()
        except Queue.EMPTY:
            raise ConnectionBusy(
                "Connection to %s is at the max number of requests" % self.host)

        self._callbacks[request_id] = cb
        self.push(msg.to_string(request_id, compression=self.compressor))
        return request_id

    def wait_for_response(self, msg):
        return self.wait_for_responses(msg)[0]

    def wait_for_responses(self, *msgs):
        waiter = ResponseWaiter(len(msgs))
        for i, msg in enumerate(msgs):
            self.send_msg(msg, partial(waiter.got_response, index=i))

        return waiter.deliver()

    def register_watcher(self, event_type, callback):
        self._push_watchers[event_type].add(callback)

    def register_watchers(self, type_callback_dict):
        for event_type, callback in type_callback_dict.items():
            self.register_watcher(event_type, callback)

    @defunct_on_error
    def _handle_options_response(self, options_response):
        if self.is_defunct:
            return
        log.debug("Received options response on new Connection from %s" % self.host)
        self.supported_cql_versions = options_response.cql_versions
        self.remote_supported_compressions = options_response.options['COMPRESSION']

        if self.cql_version:
            if self.cql_version not in self.supported_cql_versions:
                raise ProtocolError(
                    "cql_version %r is not supported by remote (w/ native "
                    "protocol). Supported versions: %r"
                    % (self.cql_version, self.supported_cql_versions))
        else:
            self.cql_version = self.supported_cql_versions[0]

        opts = {}
        self._compressor = None
        if self.compression:
            overlap = (set(locally_supported_compressions.keys()) &
                       set(self.remote_supported_compressions))
            if len(overlap) == 0:
                log.debug("No available compression types supported on both ends."
                          " locally supported: %r. remotely supported: %r"
                          % (locally_supported_compressions.keys(),
                             self.remote_supported_compressions))
            else:
                compression_type = iter(overlap).next() # choose any
                opts['COMPRESSION'] = compression_type
                # set the decompressor here, but set the compressor only after
                # a successful Ready message
                self._compressor, self.decompressor = \
                    locally_supported_compressions[compression_type]

        sm = StartupMessage(cqlversion=self.cql_version, options=opts)
        self.send_msg(sm, cb=self._handle_startup_response)

    @defunct_on_error
    def _handle_startup_response(self, startup_response):
        if self.is_defunct:
            return
        if isinstance(startup_response, ReadyMessage):
            log.debug("Got ReadyMessage on new Connection from %s" % self.host)
            if self._compressor:
                self.compressor = self._compressor
            self.connected_event.set()
        elif isinstance(startup_response, AuthenticateMessage):
            log.debug("Got AuthenticateMessage on new Connection from %s" % self.host)

            if self.credentials is None:
                raise ProgrammingError('Remote end requires authentication.')

            self.authenticator = startup_response.authenticator
            cm = CredentialsMessage(creds=self.credentials)
            self.send_msg(cm, cb=self._handle_startup_response)
        elif isinstance(startup_response, ErrorMessage):
            log.debug("Received ErrorMessage on new Connection from %s: %s"
                      % (self.host, startup_response.summary_msg()))
            raise ConnectionException(
                "Failed to initialize new connection to %s: %s"
                % (self.host, startup_response.summary_msg()))
        else:
            msg = "Unexpected response during Connection setup: %r" % (startup_response,)
            log.error(msg)
            raise ProtocolError(msg)

    def set_keyspace(self, keyspace):
        if not keyspace or keyspace == self.keyspace:
            return

        with self.lock:
            query = 'USE "%s"' % (keyspace,)
            try:
                result = self.wait_for_response(
                    QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE))
                if isinstance(result, ResultMessage):
                    self.keyspace = keyspace
                else:
                    raise self.defunct(ConnectionException(
                        "Problem while setting keyspace: %r" % (result,), self.host))
            except Exception, exc:
                raise self.defunct(ConnectionException(
                    "Problem while setting keyspace: %r" % (exc,), self.host))

class ResponseWaiter(object):

    def __init__(self, num_responses):
        self.pending = num_responses
        self.error = None
        self.responses = [None] * num_responses
        self.event = Event()

    def got_response(self, response, index):
        if isinstance(response, Exception):
            self.error = response
            self.event.set()
        else:
            self.responses[index] = response
            self.pending -= 1
            if not self.pending:
                self.event.set()

    def deliver(self):
        self.event.wait()
        if self.error:
            raise self.error
        else:
            return self.responses
<commit_msg>Correctly handle errors when decoding responses<commit_after>from collections import defaultdict, deque
import errno
from functools import partial, wraps
import logging
import socket
from threading import RLock, Event, Lock, Thread
import traceback
from Queue import Queue

import pyev

from cassandra import ConsistencyLevel
from cassandra.marshal import (int8_unpack, int32_unpack)
from cassandra.decoder import (OptionsMessage, ReadyMessage, AuthenticateMessage,
                               StartupMessage, ErrorMessage, CredentialsMessage,
                               QueryMessage, ResultMessage, decode_response)

log = logging.getLogger(__name__)

locally_supported_compressions = {}
try:
    import snappy
except ImportError:
    pass
else:
    # work around apparently buggy snappy decompress
    def decompress(byts):
        if byts == '\x00':
            return ''
        return snappy.decompress(byts)
    locally_supported_compressions['snappy'] = (snappy.compress, decompress)

MAX_STREAM_PER_CONNECTION = 128

PROTOCOL_VERSION = 0x01
PROTOCOL_VERSION_MASK = 0x7f

HEADER_DIRECTION_FROM_CLIENT = 0x00
HEADER_DIRECTION_TO_CLIENT = 0x80
HEADER_DIRECTION_MASK = 0x80

NONBLOCKING = (errno.EAGAIN, errno.EWOULDBLOCK)

class ConnectionException(Exception):

    def __init__(self, message, host=None):
        Exception.__init__(self, message)
        self.host = host

class ConnectionBusy(Exception):
    pass

class ProgrammingError(Exception):
    pass

class ProtocolError(Exception):
    pass

_loop = pyev.default_loop(pyev.EVBACKEND_SELECT)

_loop_notifier = _loop.async(lambda *a, **kw: None)
_loop_notifier.start()

# prevent _loop_notifier from keeping the loop from returning
_loop.unref()

_loop_started = None
_loop_lock = Lock()

def _run_loop():
    while True:
        end_condition = _loop.start()
        # there are still active watchers, no deadlock
        with _loop_lock:
            if end_condition:
                log.debug("Restarting event loop")
                continue
            else:
                # all Connections have been closed, no active watchers
                log.debug("All Connections currently closed, event loop ended")
                global _loop_started
                _loop_started = False
                break

def _start_loop():
    global _loop_started
    should_start = False
    with _loop_lock:
        if not _loop_started:
            log.debug("Starting pyev event loop")
            _loop_started = True
            should_start = True

    if should_start:
        t = Thread(target=_run_loop, name="async_event_loop")
        t.daemon = True
        t.start()

    return should_start

def defunct_on_error(f):

    @wraps(f)
    def wrapper(self, *args, **kwargs):
        try:
            return f(self, *args, **kwargs)
        except Exception, exc:
            self.defunct(exc)

    return wrapper

class Connection(object):

    in_buffer_size = 4096
    out_buffer_size = 4096

    cql_version = None

    keyspace = None
    compression = True
    compressor = None
    decompressor = None

    last_error = None
    in_flight = 0
    is_defunct = False
    is_closed = False
    lock = None

    _buf = ""
    _total_reqd_bytes = 0
    _read_watcher = None
    _write_watcher = None
    _socket = None

    @classmethod
    def factory(cls, *args, **kwargs):
        conn = cls(*args, **kwargs)
        conn.connected_event.wait()
        if conn.last_error:
            raise conn.last_error
        else:
            return conn

    def __init__(self, host='127.0.0.1', port=9042, credentials=None, sockopts=None, compression=True, cql_version=None):
        self.host = host
        self.port = port
        self.credentials = credentials
        self.compression = compression
        self.cql_version = cql_version

        self.connected_event = Event()

        self._id_queue = Queue(MAX_STREAM_PER_CONNECTION)
        for i in range(MAX_STREAM_PER_CONNECTION):
            self._id_queue.put_nowait(i)

        self._callbacks = {}
        self._push_watchers = defaultdict(set)
        self.lock = RLock()
        self.id_lock = Lock()

        self.deque = deque()

        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.connect((host, port))
        self._socket.setblocking(0)

        if sockopts:
            for args in sockopts:
                self._socket.setsockopt(*args)

        self._read_watcher = pyev.Io(self._socket._sock, pyev.EV_READ, _loop, self.handle_read)
        self._write_watcher = pyev.Io(self._socket._sock, pyev.EV_WRITE, _loop, self.handle_write)
        with _loop_lock:
            self._read_watcher.start()
            self._write_watcher.start()

        log.debug("Sending initial options message for new Connection to %s" % (host,))
        self.send_msg(OptionsMessage(), self._handle_options_response)

        # start the global event loop if needed
        if not _start_loop():
            # if the loop was already started, notify it
            with _loop_lock:
                _loop_notifier.send()

    def close(self):
        with self.lock:
            if self.is_closed:
                return
            self.is_closed = True

        log.debug("Closing connection to %s" % (self.host,))
        if self._read_watcher:
            self._read_watcher.stop()
        if self._write_watcher:
            self._write_watcher.stop()
        self._socket.close()
        with _loop_lock:
            _loop_notifier.send()

        # don't leave in-progress operations hanging
        if not self.is_defunct:
            self._error_all_callbacks(
                ConnectionException("Connection to %s was closed" % self.host))

    def __del__(self):
        self.close()

    def defunct(self, exc):
        log.debug("Defuncting connection to %s: %s\n%s" %
                  (self.host, exc, traceback.format_exc(exc)))
        self.last_error = exc
        self.is_defunct = True
        self._error_all_callbacks(exc)
        self.connected_event.set()
        return exc

    def _error_all_callbacks(self, exc):
        for cb in self._callbacks.values():
            cb(exc)

    def handle_write(self, watcher, revents):
        try:
            next_msg = self.deque.popleft()
        except IndexError:
            self._write_watcher.stop()
            return

        try:
            sent = self._socket.send(next_msg)
        except socket.error, err:
            if (err.args[0] in NONBLOCKING):
                self.deque.appendleft(next_msg)
            else:
                self.defunct(err)
            return
        else:
            if sent < len(next_msg):
                self.deque.appendleft(next_msg[sent:])

            if not self.deque:
                self._write_watcher.stop()

    def handle_read(self, watcher, revents):
        try:
            buf = self._socket.recv(self.in_buffer_size)
        except socket.error, err:
            if err.args[0] not in NONBLOCKING:
                self.defunct(err)
            return

        if buf:
            self._buf += buf
            while True:
                if len(self._buf) < 8:
                    # we don't have a complete header yet
                    break
                elif self._total_reqd_bytes and len(self._buf) < self._total_reqd_bytes:
                    # we already saw a header, but we don't have a complete message yet
                    break
                else:
                    body_len = int32_unpack(self._buf[4:8])
                    if len(self._buf) - 8 >= body_len:
                        msg = self._buf[:8 + body_len]
                        self._buf = self._buf[8 + body_len:]
                        self._total_reqd_bytes = 0
                        self.process_msg(msg, body_len)
                    else:
                        self._total_reqd_bytes = body_len + 8
        else:
            logging.debug("connection closed by server")
            self.close()

    @defunct_on_error
    def process_msg(self, msg, body_len):
        version, flags, stream_id, opcode = map(int8_unpack, msg[:4])
        if stream_id < 0:
            callback = None
        else:
            callback = self._callbacks.pop(stream_id)
            self._id_queue.put_nowait(stream_id)

        try:
            # check that the protocol version is supported
            given_version = version & PROTOCOL_VERSION_MASK
            if given_version != PROTOCOL_VERSION:
                raise ProtocolError("Unsupported CQL protocol version: %d" % given_version)

            # check that the header direction is correct
            if version & HEADER_DIRECTION_MASK != HEADER_DIRECTION_TO_CLIENT:
                raise ProtocolError(
                    "Header direction in response is incorrect; opcode %04x, stream id %r"
                    % (opcode, stream_id))

            if body_len > 0:
                body = msg[8:]
            elif body_len == 0:
                body = ""
            else:
                raise ProtocolError("Got negative body length: %r" % body_len)

            response = decode_response(stream_id, flags, opcode, body, self.decompressor)
        except Exception, exc:
            log.exception("Error decoding response from Cassandra. "
                          "Opcode: 0x%X, message contents: %r" % (opcode, body))
            callback(exc)
            self.defunct(exc)
            return

        try:
            if stream_id < 0:
                self.handle_pushed(response)
            elif callback is not None:
                callback(response)
        except:
            log.exception("Callback handler errored, ignoring:")

    def handle_pushed(self, response):
        for cb in self._push_watchers[response.type]:
            try:
                cb(response)
            except:
                log.exception("Pushed event handler errored, ignoring:")

    def push(self, data):
        sabs = self.out_buffer_size
        if len(data) > sabs:
            chunks = []
            for i in xrange(0, len(data), sabs):
                chunks.append(data[i:i + sabs])
        else:
            chunks = [data]

        with self.lock:
            self.deque.extend(chunks)

            if not self._write_watcher.active:
                with _loop_lock:
                    self._write_watcher.start()
                    _loop_notifier.send()

    def send_msg(self, msg, cb):
        if self.is_defunct:
            raise ConnectionException("Connection to %s is defunct" % self.host)
        elif self.is_closed:
            raise ConnectionException("Connection to %s is closed" % self.host)

        try:
            request_id = self._id_queue.get_nowait()
        except Queue.EMPTY:
            raise ConnectionBusy(
                "Connection to %s is at the max number of requests" % self.host)

        self._callbacks[request_id] = cb
        self.push(msg.to_string(request_id, compression=self.compressor))
        return request_id

    def wait_for_response(self, msg):
        return self.wait_for_responses(msg)[0]

    def wait_for_responses(self, *msgs):
        waiter = ResponseWaiter(len(msgs))
        for i, msg in enumerate(msgs):
            self.send_msg(msg, partial(waiter.got_response, index=i))

        return waiter.deliver()

    def register_watcher(self, event_type, callback):
        self._push_watchers[event_type].add(callback)

    def register_watchers(self, type_callback_dict):
        for event_type, callback in type_callback_dict.items():
            self.register_watcher(event_type, callback)

    @defunct_on_error
    def _handle_options_response(self, options_response):
        if self.is_defunct:
            return
        log.debug("Received options response on new Connection from %s" % self.host)
        self.supported_cql_versions = options_response.cql_versions
        self.remote_supported_compressions = options_response.options['COMPRESSION']

        if self.cql_version:
            if self.cql_version not in self.supported_cql_versions:
                raise ProtocolError(
                    "cql_version %r is not supported by remote (w/ native "
                    "protocol). Supported versions: %r"
                    % (self.cql_version, self.supported_cql_versions))
        else:
            self.cql_version = self.supported_cql_versions[0]

        opts = {}
        self._compressor = None
        if self.compression:
            overlap = (set(locally_supported_compressions.keys()) &
                       set(self.remote_supported_compressions))
            if len(overlap) == 0:
                log.debug("No available compression types supported on both ends."
                          " locally supported: %r. remotely supported: %r"
                          % (locally_supported_compressions.keys(),
                             self.remote_supported_compressions))
            else:
                compression_type = iter(overlap).next() # choose any
                opts['COMPRESSION'] = compression_type
                # set the decompressor here, but set the compressor only after
                # a successful Ready message
                self._compressor, self.decompressor = \
                    locally_supported_compressions[compression_type]

        sm = StartupMessage(cqlversion=self.cql_version, options=opts)
        self.send_msg(sm, cb=self._handle_startup_response)

    @defunct_on_error
    def _handle_startup_response(self, startup_response):
        if self.is_defunct:
            return
        if isinstance(startup_response, ReadyMessage):
            log.debug("Got ReadyMessage on new Connection from %s" % self.host)
            if self._compressor:
                self.compressor = self._compressor
            self.connected_event.set()
        elif isinstance(startup_response, AuthenticateMessage):
            log.debug("Got AuthenticateMessage on new Connection from %s" % self.host)

            if self.credentials is None:
                raise ProgrammingError('Remote end requires authentication.')

            self.authenticator = startup_response.authenticator
            cm = CredentialsMessage(creds=self.credentials)
            self.send_msg(cm, cb=self._handle_startup_response)
        elif isinstance(startup_response, ErrorMessage):
            log.debug("Received ErrorMessage on new Connection from %s: %s"
                      % (self.host, startup_response.summary_msg()))
            raise ConnectionException(
                "Failed to initialize new connection to %s: %s"
                % (self.host, startup_response.summary_msg()))
        else:
            msg = "Unexpected response during Connection setup: %r" % (startup_response,)
            log.error(msg)
            raise ProtocolError(msg)

    def set_keyspace(self, keyspace):
        if not keyspace or keyspace == self.keyspace:
            return

        with self.lock:
            query = 'USE "%s"' % (keyspace,)
            try:
                result = self.wait_for_response(
                    QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE))
                if isinstance(result, ResultMessage):
                    self.keyspace = keyspace
                else:
                    raise self.defunct(ConnectionException(
                        "Problem while setting keyspace: %r" % (result,), self.host))
            except Exception, exc:
                raise self.defunct(ConnectionException(
                    "Problem while setting keyspace: %r" % (exc,), self.host))

class ResponseWaiter(object):

    def __init__(self, num_responses):
        self.pending = num_responses
        self.error = None
        self.responses = [None] * num_responses
        self.event = Event()

    def got_response(self, response, index):
        if isinstance(response, Exception):
            self.error = response
            self.event.set()
        else:
            self.responses[index] = response
            self.pending -= 1
            if not self.pending:
                self.event.set()

    def deliver(self):
        self.event.wait()
        if self.error:
            raise self.error
        else:
            return self.responses
<|endoftext|>


<commit_before># The MIT License (MIT)
# Copyright (c) 2016, 2017 by the ESA CCI Toolbox development team and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Description
===========

Arithmetic operations

Functions
=========
"""

import sys

from cate.core.op import op
from xarray import ufuncs as xu
import xarray as xr

@op(tags=['arithmetic'])
def ds_arithmetics(ds: xr.Dataset,
                op: str) -> xr.Dataset:
    """
    Do arithmetic operations on the given dataset by providing a list of
    arithmetic operations and the corresponding constant. The operations will
    be applied to the dataset in the order in which they appear in the list.
    For example:
    'log,+5,-2,/3,*2'

    Currently supported arithmetic operations:
    log,log10,log2,log1p,exp,+,-,/,*

    where:
        log - natural logarithm
        log10 - base 10 logarithm
        log2 - base 2 logarithm
        log1p - log(1+x)
        exp - the exponential

    The operations will be applied element-wise to all arrays of the dataset.

    :param ds: The dataset to which to apply arithmetic operations
    :param op: A comma separated list of arithmetic operations to apply
    :return: The dataset with given arithmetic operations applied
    """
    retset = ds
    for item in op.split(','):
        item = item.strip()
        if item[0] == '+':
            retset = retset + float(item[1:])
        elif item[0] == '-':
            retset = retset - float(item[1:])
        elif item[0] == '*':
            retset = retset * float(item[1:])
        elif item[0] == '/':
            retset = retset / float(item[1:])
        elif item[:] == 'log':
            retset = xu.log(retset)
        elif item[:] == 'log10':
            retset = xu.log10(retset)
        elif item[:] == 'log2':
            retset = xu.log2(retset)
        elif item[:] == 'log1p':
            retset = xu.log1p(retset)
        elif item[:] == 'exp':
            retset = xu.exp(retset)
        else:
            raise ValueError('Arithmetic operation {} not'
                             ' implemented.'.format(item[0]))

    return retset

@op(tags=['arithmetic'])
def diff(ds: xr.Dataset, ds2:xr.Dataset):
    """
    Calculate the difference of two datasets (ds - ds2). This is done by
    matching variable names in the two datasets against each other and taking
    the difference of matching variables.

    If lat/lon/time extents differ between the datasets, the default behavior
    is to take the intersection of the datasets and run subtraction on that.
    However, broadcasting is possible. E.g. ds(lat/lon/time) - ds(lat/lon) is
    valid. In this case the subtrahend will be stretched to the size of
    ds(lat/lon/time) so that it can be subtracted. This also works if the
    subtrahend is a single time slice of arbitrary temporal position. In this
    case, the time dimension will be squeezed out leaving a lat/lon dataset.

    :param ds: The minuend dataset
    :param ds2: The subtrahend dataset
    :return: The difference dataset
    """
    try:
        # Times do not intersect
        if 0 == len(ds.time - ds2.time) and\
           len(ds.time) == len(ds2.time): # Times are the same length
            # If the datasets don't intersect in time dimension, a naive difference
            # would return empty data variables. Hence, the time coordinate has to
            # be dropped beforehand
            ds = ds.drop('time')
            ds2 = ds2.drop('time')
            return ds - ds2
    except AttributeError:
        # It is likely that the one operand is a lat/lon array that can be
        # broadcast against the other operand
        pass

    try:
        if 1 ==  len(ds2.time):
            # The subtrahend is a single time-slice -> squeeze 'time' dimension to
            # be able to broadcast is along minuend
            ds2 = ds2.squeeze('time', drop=True)
    except AttributeError:
        # Doesn't have a time dimension already
        pass
    except TypeError as e:
        if 'unsized object' in str(e):
            # The 'time' variable is a scalar
            pass
        else:
            raise TypeError(str(e))

    return ds - ds2
<commit_msg>Fix PEP8 type in arithmetics<commit_after># The MIT License (MIT)
# Copyright (c) 2016, 2017 by the ESA CCI Toolbox development team and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Description
===========

Arithmetic operations

Functions
=========
"""

import sys

from cate.core.op import op
from xarray import ufuncs as xu
import xarray as xr

@op(tags=['arithmetic'])
def ds_arithmetics(ds: xr.Dataset,
                   op: str) -> xr.Dataset:
    """
    Do arithmetic operations on the given dataset by providing a list of
    arithmetic operations and the corresponding constant. The operations will
    be applied to the dataset in the order in which they appear in the list.
    For example:
    'log,+5,-2,/3,*2'

    Currently supported arithmetic operations:
    log,log10,log2,log1p,exp,+,-,/,*

    where:
        log - natural logarithm
        log10 - base 10 logarithm
        log2 - base 2 logarithm
        log1p - log(1+x)
        exp - the exponential

    The operations will be applied element-wise to all arrays of the dataset.

    :param ds: The dataset to which to apply arithmetic operations
    :param op: A comma separated list of arithmetic operations to apply
    :return: The dataset with given arithmetic operations applied
    """
    retset = ds
    for item in op.split(','):
        item = item.strip()
        if item[0] == '+':
            retset = retset + float(item[1:])
        elif item[0] == '-':
            retset = retset - float(item[1:])
        elif item[0] == '*':
            retset = retset * float(item[1:])
        elif item[0] == '/':
            retset = retset / float(item[1:])
        elif item[:] == 'log':
            retset = xu.log(retset)
        elif item[:] == 'log10':
            retset = xu.log10(retset)
        elif item[:] == 'log2':
            retset = xu.log2(retset)
        elif item[:] == 'log1p':
            retset = xu.log1p(retset)
        elif item[:] == 'exp':
            retset = xu.exp(retset)
        else:
            raise ValueError('Arithmetic operation {} not'
                             ' implemented.'.format(item[0]))

    return retset

@op(tags=['arithmetic'])
def diff(ds: xr.Dataset, ds2:xr.Dataset):
    """
    Calculate the difference of two datasets (ds - ds2). This is done by
    matching variable names in the two datasets against each other and taking
    the difference of matching variables.

    If lat/lon/time extents differ between the datasets, the default behavior
    is to take the intersection of the datasets and run subtraction on that.
    However, broadcasting is possible. E.g. ds(lat/lon/time) - ds(lat/lon) is
    valid. In this case the subtrahend will be stretched to the size of
    ds(lat/lon/time) so that it can be subtracted. This also works if the
    subtrahend is a single time slice of arbitrary temporal position. In this
    case, the time dimension will be squeezed out leaving a lat/lon dataset.

    :param ds: The minuend dataset
    :param ds2: The subtrahend dataset
    :return: The difference dataset
    """
    try:
        # Times do not intersect
        if 0 == len(ds.time - ds2.time) and\
           len(ds.time) == len(ds2.time): # Times are the same length
            # If the datasets don't intersect in time dimension, a naive difference
            # would return empty data variables. Hence, the time coordinate has to
            # be dropped beforehand
            ds = ds.drop('time')
            ds2 = ds2.drop('time')
            return ds - ds2
    except AttributeError:
        # It is likely that the one operand is a lat/lon array that can be
        # broadcast against the other operand
        pass

    try:
        if 1 ==  len(ds2.time):
            # The subtrahend is a single time-slice -> squeeze 'time' dimension to
            # be able to broadcast is along minuend
            ds2 = ds2.squeeze('time', drop=True)
    except AttributeError:
        # Doesn't have a time dimension already
        pass
    except TypeError as e:
        if 'unsized object' in str(e):
            # The 'time' variable is a scalar
            pass
        else:
            raise TypeError(str(e))

    return ds - ds2
<|endoftext|>


<commit_before># The MIT License (MIT)
# Copyright (c) 2016, 2017 by the ESA CCI Toolbox development team and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np

from ..image import OpImage, ImagePyramid, create_ndarray_downsampling_image
from ..tilingscheme import pow2_2d_subdivision

__author__ = "Norman Fomferra (Brockmann Consult GmbH)"

class H5PyDatasetImage(OpImage):
    """
    An `OpImage` implementation for h5py datasets.

    :param h5_dataset: The h5py dataset
    :param tile_size: Tile size, if None, it is derived from chunk sizes
    """

    def __init__(self, h5_dataset, tile_size=None):
        self._h5_dataset = h5_dataset
        size = h5_dataset.shape[-1], h5_dataset.shape[-2]
        width, height = size
        if tile_size is None:
            if h5_dataset.chunks:
                chunk_width, chunk_height = h5_dataset.chunks[-1], h5_dataset.chunks[-2]
            else:
                chunk_width, chunk_height = min(width, 512), min(height, 512)
            try:
                _, tile_size, _, _ = pow2_2d_subdivision(width, height, tw_opt=chunk_width, th_opt=chunk_height)
            except ValueError:
                tile_size = width, height
        mode = str(h5_dataset.dtype)
        num_tiles = width // tile_size[0], height // tile_size[1]
        super().__init__(size, tile_size, num_tiles, format='ndarray', mode=mode)

    @property
    def h5_dataset(self):
        return self._h5_dataset

    def compute_tile(self, tile_x, tile_y, rectangle):
        x, y, w, h = rectangle
        tile = self._h5_dataset[:, y:y + h, x:x + w]
        _, dh, dw = tile.shape
        fill_value = self._h5_dataset.fillvalue

        if dh < h or dw < w:
            # if original size is less than tile size, force tile size and fill with suitable background value
            if fill_value is not None:
                background_value = fill_value
            else:
                if np.issubdtype(tile.dtype, float) or np.issubdtype(tile.dtype, complex):
                    background_value = np.nan
                else:
                    background_value = 0
            new_data = np.full((1, h, w), background_value, dtype=tile.dtype)
            new_data[:, 0:dh, 0:dw] = tile
            tile = new_data

        if not np.ma.is_masked(tile):
            # if tile is not masked
            if fill_value is not None:
                # and we have a fill value, return a masked tile
                tile = np.ma.masked_equal(tile, fill_value)
            elif np.issubdtype(tile.dtype, float) or np.issubdtype(tile.dtype, complex):
                # and it is of float type, return a masked tile with a mask from invalids, i.e. NaN, -Inf, +Inf
                tile = np.ma.masked_invalid(tile)

        return tile

    def create_pyramid(self, **kwargs):
        return ImagePyramid.create_from_image(self, create_ndarray_downsampling_image, **kwargs)
<commit_msg>float --> np.floating<commit_after># The MIT License (MIT)
# Copyright (c) 2016, 2017 by the ESA CCI Toolbox development team and contributors
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np

from ..image import OpImage, ImagePyramid, create_ndarray_downsampling_image
from ..tilingscheme import pow2_2d_subdivision

__author__ = "Norman Fomferra (Brockmann Consult GmbH)"

class H5PyDatasetImage(OpImage):
    """
    An `OpImage` implementation for h5py datasets.

    :param h5_dataset: The h5py dataset
    :param tile_size: Tile size, if None, it is derived from chunk sizes
    """

    def __init__(self, h5_dataset, tile_size=None):
        self._h5_dataset = h5_dataset
        size = h5_dataset.shape[-1], h5_dataset.shape[-2]
        width, height = size
        if tile_size is None:
            if h5_dataset.chunks:
                chunk_width, chunk_height = h5_dataset.chunks[-1], h5_dataset.chunks[-2]
            else:
                chunk_width, chunk_height = min(width, 512), min(height, 512)
            try:
                _, tile_size, _, _ = pow2_2d_subdivision(width, height, tw_opt=chunk_width, th_opt=chunk_height)
            except ValueError:
                tile_size = width, height
        mode = str(h5_dataset.dtype)
        num_tiles = width // tile_size[0], height // tile_size[1]
        super().__init__(size, tile_size, num_tiles, format='ndarray', mode=mode)

    @property
    def h5_dataset(self):
        return self._h5_dataset

    def compute_tile(self, tile_x, tile_y, rectangle):
        x, y, w, h = rectangle
        tile = self._h5_dataset[:, y:y + h, x:x + w]
        _, dh, dw = tile.shape
        fill_value = self._h5_dataset.fillvalue

        if dh < h or dw < w:
            # if original size is less than tile size, force tile size and fill with suitable background value
            if fill_value is not None:
                background_value = fill_value
            else:
                if np.issubdtype(tile.dtype, np.floating) or np.issubdtype(tile.dtype, np.complexfloating):
                    background_value = np.nan
                else:
                    background_value = 0
            new_data = np.full((1, h, w), background_value, dtype=tile.dtype)
            new_data[:, 0:dh, 0:dw] = tile
            tile = new_data

        if not np.ma.is_masked(tile):
            # if tile is not masked
            if fill_value is not None:
                # and we have a fill value, return a masked tile
                tile = np.ma.masked_equal(tile, fill_value)
            elif np.issubdtype(tile.dtype, np.floating) or np.issubdtype(tile.dtype, np.complexfloating):
                # and it is of float type, return a masked tile with a mask from invalids, i.e. NaN, -Inf, +Inf
                tile = np.ma.masked_invalid(tile)

        return tile

    def create_pyramid(self, **kwargs):
        return ImagePyramid.create_from_image(self, create_ndarray_downsampling_image, **kwargs)
<|endoftext|>