import os
import shutil
import tempfile

from chainer.serializers import npz
from chainer.training import extension


def snapshot_object(target, filename, savefun=npz.save_npz):
    """Returns a trainer extension to take snapshots of a given object.

    This extension serializes the given object and saves it to the output
    directory.

    This extension is called once per epoch by default. To take a
    snapshot at a different interval, a trigger object specifying the
    required interval can be passed along with this extension
    to the `extend()` method of the trainer.

    The default priority is -100, which is lower than that of most
    built-in extensions.

    Args:
        target: Object to serialize.
        filename (str): Name of the file into which the object is serialized.
            It can be a format string, where the trainer object is passed to
            the :meth:`str.format` method. For example,
            ``'snapshot_{.updater.iteration}'`` is converted to
            ``'snapshot_10000'`` at the 10,000th iteration.
        savefun: Function to save the object. It takes two arguments: the
            output file path and the object to serialize.

    Returns:
        An extension function.

    """
    @extension.make_extension(trigger=(1, 'epoch'), priority=-100)
    def snapshot_object(trainer):
        _snapshot_object(trainer, target, filename.format(trainer), savefun)

    return snapshot_object


def snapshot(savefun=npz.save_npz,
             filename='snapshot_iter_{.updater.iteration}'):
    """Returns a trainer extension to take snapshots of the trainer.

    This extension serializes the trainer object and saves it to the output
    directory. It is used to support resuming the training loop from the saved
    state.

    This extension is called once per epoch by default. To take a
    snapshot at a different interval, a trigger object specifying the
    required interval can be passed along with this extension
    to the `extend()` method of the trainer.

    The default priority is -100, which is lower than that of most
    built-in extensions.

    .. note::
       This extension first writes the serialized object to a temporary file
       and then rename it to the target file name. Thus, if the program stops
       right before the renaming, the temporary file might be left in the
       output directory.

    Args:
        savefun: Function to save the trainer. It takes two arguments: the
            output file path and the trainer object.
        filename (str): Name of the file into which the trainer is serialized.
            It can be a format string, where the trainer object is passed to
            the :meth:`str.format` method.

    """
    @extension.make_extension(trigger=(1, 'epoch'), priority=-100)
    def snapshot(trainer):
        _snapshot_object(trainer, trainer, filename.format(trainer), savefun)

    return snapshot


def _snapshot_object(trainer, target, filename, savefun):
    fn = filename.format(trainer)
    prefix = 'tmp' + fn
    fd, tmppath = tempfile.mkstemp(prefix=prefix, dir=trainer.out)
    try:
        savefun(tmppath, target)
    except Exception:
        os.close(fd)
        os.remove(tmppath)
        raise
    os.close(fd)
    shutil.move(tmppath, os.path.join(trainer.out, fn))

@@ -109,16 +109,246 @@
 tension%0A
+from chainer.training.extensions.snapshots import condition as condition_module%0Afrom chainer.training.extensions.snapshots import handler as handler_module%0Afrom chainer.training.extensions.snapshots import writer as writer_module%0A
 %0A%0Adef sn
@@ -3540,8 +3540,1823 @@
 t, fn))%0A
+%0A%0Aclass Snapshot(extension.Extension):%0A    %22%22%22Takes a snapshot.%0A%0A    Args:%0A        target: Object to serialize. If not specified, it will%0A            be trainer object.%0A        condition: Condition object. It must be a callable object that%0A            returns boolean in its call. If it returns True the snapshot will%0A            be done. If not it will be skipped.%0A        writer: Writer object. It need to be a callable object.%0A        handler: Serializer handler object.%0A        filename (str): Name of the file into which the object is serialized.%0A            It can be a format string, where the trainer object is passed to%0A            the :meth:%60str.format%60 method. For example,%0A            %60%60'snapshot_%7B.updater.iteration%7D'%60%60 is converted to%0A            %60%60'snapshot_10000'%60%60 at the 10,000th iteration.%0A%0A    %22%22%22%0A%0A    def __init__(self,%0A                 target=None,%0A                 condition=condition_module.Always(),%0A                 writer=writer_module.SimpleWriter(),%0A                 handler=handler_module.SerializerHandler(),%0A                 filename='snapshot_iter_%7B.updater.iteration%7D'):%0A        self._target = target%0A        self._filename = filename%0A        self.condition = condition%0A        self.writer = writer%0A        self.handler = handler%0A%0A    def __call__(self, trainer):%0A        if self.condition(trainer, self):%0A            target = trainer if self._target is None else self._target%0A            self.handler.serialize(target)%0A            filename = self._filename%0A            if callable(filename):%0A                filename = filename(trainer)%0A            filename = filename.format(trainer)%0A            outdir = trainer.out%0A            self.writer(filename, outdir, self.handler)%0A%0A    def finalize(self):%0A        if hasattr(self.writer, 'finalize'):%0A            self.writer.finalize()%0A

Update _snapshot file.