# -*- coding: utf-8 -*-
# as a copy of http://twistedmatrix.com/trac/ticket/990 . The licence of this
# file is for twisted under the MIT:
#       http://www.opensource.org/licenses/mit-license.php
# 
# Copyright (c) 2001-2008
# Allen Short
# Andrew Bennetts
# Apple Computer, Inc.
# Benjamin Bruheim
# Bob Ippolito
# Canonical Limited
# Christopher Armstrong
# David Reid
# Donovan Preston
# Eric Mangold
# Itamar Shtull-Trauring
# James Knight
# Jason A. Mobarak
# Jonathan Lange
# Jonathan D. Simms
# Jp Calderone
# Jürgen Hermann
# Kevin Turner
# Mary Gardiner
# Matthew Lefkowitz
# Massachusetts Institute of Technology
# Moshe Zadka
# Paul Swartz
# Pavel Pergamenshchik
# Ralph Meijer
# Sean Riley
# Travis B. Hartwell
# Thomas Herve
# Eyal Lotem
# Antoine Pitrou
# Andy Gayton
# 
# 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.
# 

from twisted.internet import defer, task, reactor
from twisted.python import failure
from twisted.python.util import mergeFunctionMetadata

"""
Implements the Cancellable Deferred as discussed at
http://twistedmatrix.com/trac/ticket/990
"""

class CancelledError(Exception):
    pass

class CancellableDeferred(defer.Deferred):
    """
    see L{twisted.internet.defer.Deferred}.

    When creating this Deferred, you may provide a canceller function,
    which will be called by d.cancel() to let you do any cleanup necessary
    if the user decides not to wait for the deferred to complete.
    """

    suppressAlreadyCalled = 0

    def __init__(self, canceller=None):
        self.canceller = canceller
        defer.Deferred.__init__(self)
    
    def cancel(self):
        """Cancel this deferred.

        If the deferred is waiting on another deferred, forward the
        cancellation to the other deferred.

        If the deferred has not yet been errback'd/callback'd, call
        the canceller function provided to the constructor. If that
        function does not do a callback/errback, or if no canceller
        function was provided, errback with CancelledError.

        Otherwise, raise AlreadyCalledError.
        """
        canceller=self.canceller
        if not self.called:
            if canceller:
                canceller(self)
            else:
                # Eat the callback that will eventually be fired
                # since there was no real canceller.
                self.suppressAlreadyCalled = 1

            if not self.called:
                # The canceller didn't do an errback of its own
                try:
                    raise CancelledError
                except:
                    self.errback(failure.Failure())
        elif isinstance(self.result, CancellableDeferred):
            # Waiting for another deferred -- cancel it instead
            # FIXME: possible problem with old API?
            self.result.cancel()
        else:
            # Called and not waiting for another deferred
            raise defer.AlreadyCalledError

    def _startRunCallbacks(self, result):
        # Canceller is no longer relevant
        self.canceller=None
        
        if self.called:
            if self.suppressAlreadyCalled:
                self.suppressAlreadyCalled = False
                return
            raise defer.AlreadyCalledError

        defer.Deferred._startRunCallbacks(self, result)

class CancellableDeferredIterator(object):
    def __init__(self, iterator, *args, **kw):
        self.iterator = iterator
        self.args = args
        self.kw = kw
        self.deferred = None
        self.cancelled = False

    def cancel(self):
        self.cancelled = True
        if self.deferred is not None:
            self.deferred.cancel()
            self.deferred = None

    def _called(self, result_or_failure):
        self.deferred = None

        return result_or_failure

    def __iter__(self):
        if self.cancelled:
            yield iter([])
            return

        for deferred in self.iterator(*self.args, **self.kw):
            self.deferred = deferred or defer.succeed(None)
            self.deferred.addBoth(self._called)
            yield self.deferred

            if self.cancelled:
                break

def cancellable_deferred_iterator(iterator):
    """
    Decorator for iterators that internally use cancellable deferreds. This
    one allows you to cancel the iterator operation and it propagates the
    cancel call to the deferred that is currently hold.
    """
    def wrapper(*args, **kw):
        return CancellableDeferredIterator(iterator, *args, **kw)

    wrapper = mergeFunctionMetadata(iterator, wrapper)
    return wrapper

def _cancellable_coiterate_impl(coiterate, iterator_factory, *args, **kw):
    """
    Wrapper around @coiterate to iterate and return a usual
    L{CancellableDeferred} that can cancel the iteration.
    """
    cancel_iter = CancellableDeferredIterator(iterator_factory, *args, **kw)
    def cancel(deferred):
        cancel_iter.cancel()

    cancel_dfr = CancellableDeferred(cancel)
    dfr = coiterate(iter(cancel_iter))
    dfr.chainDeferred(cancel_dfr)
    return cancel_dfr

def cancellable_coiterate(iterator_factory, *args, **kw):
    """
    Wrapper around task.coiterate to iterate and return a usual
    L{CancellableDeferred} that can cancel the iteration.
    """
    return _cancellable_coiterate_impl(task.coiterate, iterator_factory,
                                       *args, **kw)
 
def delay_coiterate(delay, iterator):
    """
    Does the same thing as task.coiterate(), but only run an iteration every
    @delay seconds
    """

    if delay == 0:
        # We use task.coiterate, else our delay_scheduler would make us starve
        # the reactor
        return task.coiterate(iterator)

    def delay_scheduler(x):
        # Determines when the next cooperator step will be run
        return reactor.callLater(delay, x)

    def always_terminate_factory():
        # Returns a function that will determine whether to stop iterating
        # for the current step of the cooperator. We always return true, so
        # that the cooperator does only one iteration at each step
        return lambda: True

    cooperator = task.Cooperator(scheduler=delay_scheduler,
                                 terminationPredicateFactory=always_terminate_factory)
    return cooperator.coiterate(iterator)

def cancellable_delay_coiterate(delay, iterator_factory, *args, **kw):
    """
    Wrapper around delay_coiterate to iterate and return a usual
    L{CancellableDeferred} that can cancel the iteration.
    """
    def coiterate(iterator):
        return delay_coiterate(delay, iterator)

    return _cancellable_coiterate_impl(coiterate, iterator_factory, *args, **kw)
