<commit_before>#!/usr/bin/env python
#
# Copyright 2009 Facebook
# Copyright 2010 Alice Bevan-McGregor
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

"""A utility class to write to and read from a non-blocking socket."""

from __future__ import unicode_literals

import errno
import socket
import ssl

from marrow.io import ioloop
from marrow.util.compat import exception

log = __import__('logging').getLogger(__name__)
__all__ = ['IOStream', 'SSLIOStream']

class IOStream(object):
    """A utility class to write to and read from a non-blocking socket.

    We support three methods: write(), read_until(), and read_bytes().
    All of the methods take callbacks (since writing and reading are
    non-blocking and asynchronous). read_until() reads the socket until
    a given delimiter, and read_bytes() reads until a specified number
    of bytes have been read from the socket.

    A very simple (and broken) HTTP client using this class:

        import ioloop
        import iostream
        import socket

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        s.connect(("friendfeed.com", 80))
        stream = IOStream(s)

        def on_headers(data):
            headers = {}
            for line in data.split("\r\n"):
               parts = line.split(":")
               if len(parts) == 2:
                   headers[parts[0].strip()] = parts[1].strip()
            stream.read_bytes(int(headers["Content-Length"]), on_body)

        def on_body(data):
            print data
            stream.close()
            ioloop.IOLoop.instance().stop()

        stream.write("GET / HTTP/1.0\r\n\r\n")
        stream.read_until("\r\n\r\n", on_headers)
        ioloop.IOLoop.instance().start()

    """
    def __init__(self, socket, io_loop=None, max_buffer_size=104857600, read_chunk_size=4096):
        self.socket = socket
        self.address = socket.getpeername()
        self.socket.setblocking(False)
        self.io_loop = io_loop or ioloop.IOLoop.instance()
        self.max_buffer_size = max_buffer_size
        self.read_chunk_size = read_chunk_size
        self._read_buffer = b""
        self._write_buffer = b""
        self._read_delimiter = None
        self._read_bytes = None
        self._read_callback = None
        self._write_callback = None
        self._close_callback = None
        self._state = self.io_loop.ERROR
        self.io_loop.add_handler(self.socket.fileno(), self._handle_events, self._state)

    def read_until(self, delimiter, callback):
        """Call callback when we read the given delimiter."""
        assert not self._read_callback, "Already reading."

        loc = self._read_buffer.find(delimiter)

        if loc != -1:
            self._run_callback(callback, self._consume(loc + len(delimiter)))
            return

        # TODO: Allow multiple registered callbacks.
        self._check_closed()
        self._read_delimiter = delimiter
        self._read_callback = callback
        self._add_io_state(self.io_loop.READ)

    def read_bytes(self, num_bytes, callback):
        """Call callback when we read the given number of bytes."""
        assert not self._read_callback, "Already reading"

        if len(self._read_buffer) >= num_bytes:
            self._run_callback(callback, self._consume(num_bytes))
            return

        # TODO: Allow multiple registered callbacks.
        self._check_closed()
        self._read_bytes = num_bytes
        self._read_callback = callback
        self._add_io_state(self.io_loop.READ)

    def write(self, data, callback=None):
        """Write the given data to this stream.

        If callback is given, we call it when all of the buffered write
        data has been successfully written to the stream. If there was
        previously buffered write data and an old write callback, that
        callback is simply overwritten with this new callback.
        """
        self._check_closed()
        self._write_buffer += data
        self._add_io_state(self.io_loop.WRITE)

        # TODO: Allow multiple registered callbacks.
        self._write_callback = callback

    def set_close_callback(self, callback):
        """Call the given callback when the stream is closed."""
        # TODO: Allow multiple registered callbacks.
        self._close_callback = callback

    # TODO: set_exception_callback

    def close(self):
        """Close this stream."""
        if self.socket is not None:
            self.io_loop.remove_handler(self.socket.fileno())
            self.socket.shutdown(socket.SHUT_RDWR)
            self.socket.close()
            self.socket = None

            # TODO: Allow multiple registered callbacks.
            if self._close_callback:
                self._run_callback(self._close_callback)

    def reading(self):
        """Returns true if we are currently reading from the stream."""
        return self._read_callback is not None

    def writing(self):
        """Returns true if we are currently writing to the stream."""
        return len(self._write_buffer) # w/o > 0 = 27% speed increase if False, 5% if True

    def closed(self):
        return self.socket is None

    def _handle_events(self, fd, events):
        if not self.socket:
            log.warning("Got events for closed stream %d", fd)
            return

        if events & self.io_loop.READ:
            self._handle_read()

            if not self.socket:
                return

        if events & self.io_loop.WRITE:
            self._handle_write()

            if not self.socket:
                return

        if events & self.io_loop.ERROR:
            self.close()
            return

        state = self.io_loop.ERROR
        if self._read_delimiter or self._read_bytes:
            state |= self.io_loop.READ

        if self._write_buffer:
            state |= self.io_loop.WRITE

        if state != self._state:
            self._state = state
            self.io_loop.update_handler(self.socket.fileno(), self._state)

    def _run_callback(self, callback, *args, **kwargs):
        try:
            callback(*args, **kwargs)

        except:
            # Close the socket on an uncaught exception from a user callback
            # (It would eventually get closed when the socket object is
            # gc'd, but we don't want to rely on gc happening before we
            # run out of file descriptors)
            self.close()
            # Re-raise the exception so that IOLoop.handle_callback_exception
            # can see it and log the error
            raise

    def _handle_read(self):
        try:
            chunk = self.socket.recv(self.read_chunk_size)

        except socket.error:
            e = exception().exception
            if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                return

            else:
                log.warning("Read error on %d: %s", self.socket.fileno(), e)
                self.close()
                return

        if not chunk:
            self.close()
            return

        self._read_buffer += chunk
        rblen = len(self._read_buffer)

        if rblen >= self.max_buffer_size:
            log.error("Connection %d reached maximum read buffer size.", self.socket.fileno())
            self.close()
            return

        if self._read_bytes:
            if rblen >= self._read_bytes:
                num_bytes = self._read_bytes
                callback = self._read_callback
                self._read_callback = None
                self._read_bytes = None
                self._run_callback(callback, self._consume(num_bytes))

        elif self._read_delimiter:
            loc = self._read_buffer.find(self._read_delimiter)

            if loc != -1:
                callback = self._read_callback
                delimiter_len = len(self._read_delimiter)
                self._read_callback = None
                self._read_delimiter = None
                self._run_callback(callback, self._consume(loc + delimiter_len))

    def _handle_write(self):
        while self._write_buffer:
            try:
                num_bytes = self.socket.send(self._write_buffer)
                self._write_buffer = self._write_buffer[num_bytes:]
            except socket.error:
                e = exception().exception
                if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                    break
                else:
                    log.warning("Write error on %d: %s", self.socket.fileno(), e)
                    self.close()
                    return
        if not self._write_buffer and self._write_callback:
            callback = self._write_callback
            self._write_callback = None
            self._run_callback(callback)

    def _consume(self, loc):
        result = self._read_buffer[:loc]
        self._read_buffer = self._read_buffer[loc:]
        return result

    def _check_closed(self):
        if not self.socket:
            raise IOError("Stream is closed")

    def _add_io_state(self, state):
        if not self._state & state:
            self._state = self._state | state
            self.io_loop.update_handler(self.socket.fileno(), self._state)

class SSLIOStream(IOStream):
    """Sets up an SSL connection in a non-blocking manner"""
    def __init__(self, *args, **kwargs):
        super(SSLIOStream, self).__init__(*args, **kwargs)
        self._ssl_accepting = True
        self._do_ssl_handshake()

    def _do_ssl_handshake(self):
        # Based on code from test_ssl.py in the python stdlib
        try:
            self.socket.do_handshake()
        except ssl.SSLError:
            err = exception().exception
            if err.args[0] == ssl.SSL_ERROR_WANT_READ:
                self._add_io_state(self.io_loop.READ)
                return
            elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
                self._add_io_state(self.io_loop.WRITE)
                return
            elif err.args[0] in (ssl.SSL_ERROR_EOF,
                                 ssl.SSL_ERROR_ZERO_RETURN):
                return self.close()
            raise
        except socket.error:
            err = exception().exception
            if err.args[0] == errno.ECONNABORTED:
                return self.close()
        else:
            self._ssl_accepting = False

    def _handle_read(self):
        if self._ssl_accepting:
            self._do_ssl_handshake()
            return
        super(SSLIOStream, self)._handle_read()

    def _handle_write(self):
        if self._ssl_accepting:
            self._do_ssl_handshake()
            return
        super(SSLIOStream, self)._handle_write()
<commit_msg>Further code cleanup.<commit_after>#!/usr/bin/env python
#
# Copyright 2009 Facebook
# Copyright 2010 Alice Bevan-McGregor
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

"""A utility class to write to and read from a non-blocking socket."""

from __future__ import unicode_literals

import errno
import socket
import ssl

from marrow.io import ioloop
from marrow.util.compat import exception

log = __import__('logging').getLogger(__name__)
__all__ = ['IOStream', 'SSLIOStream']

class IOStream(object):
    """A utility class to write to and read from a non-blocking socket.

    We support three methods: write(), read_until(), and read_bytes().
    All of the methods take callbacks (since writing and reading are
    non-blocking and asynchronous). read_until() reads the socket until
    a given delimiter, and read_bytes() reads until a specified number
    of bytes have been read from the socket.

    A very simple (and broken) HTTP client using this class:

        import ioloop
        import iostream
        import socket

        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        s.connect(("friendfeed.com", 80))
        stream = IOStream(s)

        def on_headers(data):
            headers = {}
            for line in data.split("\r\n"):
               parts = line.split(":")
               if len(parts) == 2:
                   headers[parts[0].strip()] = parts[1].strip()
            stream.read_bytes(int(headers["Content-Length"]), on_body)

        def on_body(data):
            print data
            stream.close()
            ioloop.IOLoop.instance().stop()

        stream.write("GET / HTTP/1.0\r\n\r\n")
        stream.read_until("\r\n\r\n", on_headers)
        ioloop.IOLoop.instance().start()

    """
    def __init__(self, socket, io_loop=None, max_buffer_size=104857600, read_chunk_size=4096):
        self.socket = socket
        self.address = socket.getpeername()
        self.socket.setblocking(False)
        self.io_loop = io_loop or ioloop.IOLoop.instance()
        self.max_buffer_size = max_buffer_size
        self.read_chunk_size = read_chunk_size
        self._read_buffer = b""
        self._write_buffer = b""
        self._read_delimiter = None
        self._read_bytes = None
        self._read_callback = None
        self._write_callback = None
        self._close_callback = None
        self._state = self.io_loop.ERROR
        self.io_loop.add_handler(self.socket.fileno(), self._handle_events, self._state)

    def read_until(self, delimiter, callback):
        """Call callback when we read the given delimiter."""
        assert not self._read_callback, "Already reading."

        loc = self._read_buffer.find(delimiter)

        if loc != -1:
            self._run_callback(callback, self._consume(loc + len(delimiter)))
            return

        # TODO: Allow multiple registered callbacks.
        self._check_closed()
        self._read_delimiter = delimiter
        self._read_callback = callback
        self._add_io_state(self.io_loop.READ)

    def read_bytes(self, num_bytes, callback):
        """Call callback when we read the given number of bytes."""
        assert not self._read_callback, "Already reading"

        if len(self._read_buffer) >= num_bytes:
            self._run_callback(callback, self._consume(num_bytes))
            return

        # TODO: Allow multiple registered callbacks.
        self._check_closed()
        self._read_bytes = num_bytes
        self._read_callback = callback
        self._add_io_state(self.io_loop.READ)

    def write(self, data, callback=None):
        """Write the given data to this stream.

        If callback is given, we call it when all of the buffered write
        data has been successfully written to the stream. If there was
        previously buffered write data and an old write callback, that
        callback is simply overwritten with this new callback.
        """
        self._check_closed()
        self._write_buffer += data
        self._add_io_state(self.io_loop.WRITE)

        # TODO: Allow multiple registered callbacks.
        self._write_callback = callback

    def set_close_callback(self, callback):
        """Call the given callback when the stream is closed."""
        # TODO: Allow multiple registered callbacks.
        self._close_callback = callback

    # TODO: set_exception_callback

    def close(self):
        """Close this stream."""
        if self.socket is not None:
            self.io_loop.remove_handler(self.socket.fileno())
            self.socket.shutdown(socket.SHUT_RDWR)
            self.socket.close()
            self.socket = None

            # TODO: Allow multiple registered callbacks.
            if self._close_callback:
                self._run_callback(self._close_callback)

    def reading(self):
        """Returns true if we are currently reading from the stream."""
        return self._read_callback is not None

    def writing(self):
        """Returns true if we are currently writing to the stream."""
        return len(self._write_buffer) # w/o > 0 = 27% speed increase if False, 5% if True

    def closed(self):
        return self.socket is None

    def _handle_events(self, fd, events):
        if not self.socket:
            log.warning("Got events for closed stream %d", fd)
            return

        if events & self.io_loop.READ:
            self._handle_read()

            if not self.socket:
                return

        if events & self.io_loop.WRITE:
            self._handle_write()

            if not self.socket:
                return

        if events & self.io_loop.ERROR:
            self.close()
            return

        state = self.io_loop.ERROR
        if self._read_delimiter or self._read_bytes:
            state |= self.io_loop.READ

        if self._write_buffer:
            state |= self.io_loop.WRITE

        if state != self._state:
            self._state = state
            self.io_loop.update_handler(self.socket.fileno(), self._state)

    def _run_callback(self, callback, *args, **kwargs):
        try:
            callback(*args, **kwargs)

        except:
            # Close the socket on an uncaught exception from a user callback
            # (It would eventually get closed when the socket object is
            # gc'd, but we don't want to rely on gc happening before we
            # run out of file descriptors)
            self.close()
            # Re-raise the exception so that IOLoop.handle_callback_exception
            # can see it and log the error
            raise

    def _handle_read(self):
        try:
            chunk = self.socket.recv(self.read_chunk_size)

        except socket.error:
            e = exception().exception
            if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                return

            else:
                log.warning("Read error on %d: %s", self.socket.fileno(), e)
                self.close()
                return

        if not chunk:
            self.close()
            return

        self._read_buffer += chunk
        rblen = len(self._read_buffer)

        if rblen >= self.max_buffer_size:
            log.error("Connection %d reached maximum read buffer size.", self.socket.fileno())
            self.close()
            return

        if self._read_bytes:
            if rblen >= self._read_bytes:
                num_bytes = self._read_bytes
                callback = self._read_callback
                self._read_callback = None
                self._read_bytes = None
                self._run_callback(callback, self._consume(num_bytes))

        elif self._read_delimiter:
            loc = self._read_buffer.find(self._read_delimiter)

            if loc != -1:
                callback = self._read_callback
                delimiter_len = len(self._read_delimiter)
                self._read_callback = None
                self._read_delimiter = None
                self._run_callback(callback, self._consume(loc + delimiter_len))

    def _handle_write(self):
        while self._write_buffer:
            try:
                num_bytes = self.socket.send(self._write_buffer)
                self._write_buffer = self._write_buffer[num_bytes:]

            except socket.error:
                e = exception().exception

                if e[0] in (errno.EWOULDBLOCK, errno.EAGAIN):
                    break

                else:
                    log.warning("Write error on %d: %s", self.socket.fileno(), e)
                    self.close()
                    return

        # TODO: Allow multiple callbacks.
        if not self._write_buffer and self._write_callback:
            callback = self._write_callback
            self._write_callback = None
            self._run_callback(callback)

    def _consume(self, loc):
        result = self._read_buffer[:loc]
        self._read_buffer = self._read_buffer[loc:]
        return result

    def _check_closed(self):
        if not self.socket:
            raise IOError("Stream is closed.")

    def _add_io_state(self, state):
        if not self._state & state:
            self._state = self._state | state
            self.io_loop.update_handler(self.socket.fileno(), self._state)

class SSLIOStream(IOStream):
    """Sets up an SSL connection in a non-blocking manner"""
    def __init__(self, *args, **kwargs):
        super(SSLIOStream, self).__init__(*args, **kwargs)
        self._ssl_accepting = True
        self._do_ssl_handshake()

    def _do_ssl_handshake(self):
        # Based on code from test_ssl.py in the python stdlib
        try:
            self.socket.do_handshake()
        except ssl.SSLError:
            err = exception().exception
            if err.args[0] == ssl.SSL_ERROR_WANT_READ:
                self._add_io_state(self.io_loop.READ)
                return
            elif err.args[0] == ssl.SSL_ERROR_WANT_WRITE:
                self._add_io_state(self.io_loop.WRITE)
                return
            elif err.args[0] in (ssl.SSL_ERROR_EOF,
                                 ssl.SSL_ERROR_ZERO_RETURN):
                return self.close()
            raise
        except socket.error:
            err = exception().exception
            if err.args[0] == errno.ECONNABORTED:
                return self.close()
        else:
            self._ssl_accepting = False

    def _handle_read(self):
        if self._ssl_accepting:
            self._do_ssl_handshake()
            return
        super(SSLIOStream, self)._handle_read()

    def _handle_write(self):
        if self._ssl_accepting:
            self._do_ssl_handshake()
            return
        super(SSLIOStream, self)._handle_write()
<|endoftext|>


<commit_before># Lint as: python3
# Copyright 2019, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Library for loading CIFAR-100 training and testing data."""

import collections
import functools

import tensorflow as tf
import tensorflow_federated as tff

CIFAR_SHAPE = (32, 32, 3)
TOTAL_FEATURE_SIZE = 32 * 32 * 3
NUM_EXAMPLES_PER_CLIENT = 500
TEST_BATCH_SIZE = 100

def preprocess_cifar_example(example, crop_shape, distort=False):
  """Preprocesses a CIFAR-100 example by cropping, flipping, and normalizing."""
  image = tf.cast(example['image'], tf.float32)
  if distort:
    image = tf.image.random_crop(image, size=crop_shape)
    image = tf.image.random_flip_left_right(image)
  else:
    image = tf.image.resize_with_crop_or_pad(
        image, target_height=crop_shape[1], target_width=crop_shape[2])
  image = tf.image.per_image_standardization(image)
  return (image, example['label'])

def get_federated_cifar100(client_epochs_per_round,
                           train_batch_size,
                           crop_shape=CIFAR_SHAPE):
  """Loads and preprocesses federated CIFAR100 training and testing sets.

  Args:
    client_epochs_per_round: An integer specifying the number of local epochs
      performed per training round by each client. Used to replicate each client
      dataset an appropriate number of times.
    train_batch_size: The batch size for the training dataset.
    crop_shape: An iterable of integers specifying the desired crop
      shape for pre-processing. Must be convertable to a tuple of integers
      (CROP_HEIGHT, CROP_WIDTH, NUM_CHANNELS) which cannot have elements that
      exceed (32, 32, 3), element-wise. The element in the last index should be
      set to 3 to maintain the RGB image structure of the elements.

  Returns:
    A tuple of `tff.simulation.ClientData` and `tf.data.Datset` objects.
  """
  if not isinstance(crop_shape, collections.Iterable):
    raise TypeError(
        'Argument crop_shape must be an iterable.')
  crop_shape = tuple(crop_shape)
  if len(crop_shape) != 3:
    raise ValueError('The crop_shape must have length 3, corresponding to a '
                     'tensor of shape [height, width, channels].')
  cifar_train, cifar_test = tff.simulation.datasets.cifar100.load_data()
  train_crop_shape = (train_batch_size,) + crop_shape
  test_crop_shape = (TEST_BATCH_SIZE,) + crop_shape
  train_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=train_crop_shape, distort=True)
  test_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=test_crop_shape, distort=False)

  def preprocess_train_dataset(dataset):
    """Preprocess CIFAR100 training dataset."""
    return dataset.shuffle(buffer_size=NUM_EXAMPLES_PER_CLIENT).repeat(
        client_epochs_per_round).batch(
            train_batch_size, drop_remainder=True).map(train_image_map)

  def preprocess_test_dataset(dataset):
    """Preprocess CIFAR100 testing dataset."""
    return dataset.batch(
        TEST_BATCH_SIZE, drop_remainder=False).map(test_image_map)

  cifar_train = cifar_train.preprocess(preprocess_train_dataset)
  cifar_test = preprocess_test_dataset(
      cifar_test.create_tf_dataset_from_all_clients()).cache()
  return cifar_train, cifar_test

def get_centralized_cifar100(train_batch_size, crop_shape=CIFAR_SHAPE):
  """Loads and preprocesses centralized CIFAR100 training and testing sets.

  Args:
    train_batch_size: The batch size for the training dataset.
    crop_shape: An iterable of integers specifying the desired crop
      shape for pre-processing. Must be convertable to a tuple of integers
      (CROP_HEIGHT, CROP_WIDTH, NUM_CHANNELS) which cannot have elements that
      exceed (32, 32, 3), element-wise. The element in the last index should be
      set to 3 to maintain the RGB image structure of the elements.

  Returns:
    A length two tuple of `tf.data.Dataset` objects.
  """
  try:
    crop_shape = tuple(crop_shape)
  except:
    raise ValueError(
        'Argument crop_shape must be able to coerced into a length 3 tuple.')
  if len(crop_shape) != 3:
    raise ValueError('The crop_shape must have length 3, corresponding to a '
                     'tensor of shape [height, width, channels].')
  cifar_train, cifar_test = tff.simulation.datasets.cifar100.load_data()
  train_crop_shape = (train_batch_size,) + crop_shape
  test_crop_shape = (TEST_BATCH_SIZE,) + crop_shape
  train_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=train_crop_shape, distort=True)
  test_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=test_crop_shape, distort=False)

  cifar_train = cifar_train.create_tf_dataset_from_all_clients().shuffle(
      buffer_size=10000).batch(train_batch_size, drop_remainder=True).map(
          train_image_map, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  cifar_test = cifar_test.create_tf_dataset_from_all_clients().batch(
      TEST_BATCH_SIZE, drop_remainder=False).map(test_image_map).cache()

  return cifar_train, cifar_test
<commit_msg>Extends Cifar100 dataset to allow for serialization of returned federated datasets.<commit_after># Lint as: python3
# Copyright 2019, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Library for loading CIFAR-100 training and testing data."""

import collections
import functools

import tensorflow as tf
import tensorflow_federated as tff

CIFAR_SHAPE = (32, 32, 3)
TOTAL_FEATURE_SIZE = 32 * 32 * 3
NUM_EXAMPLES_PER_CLIENT = 500
TEST_BATCH_SIZE = 100

def preprocess_cifar_example(example, crop_shape, distort=False):
  """Preprocesses a CIFAR-100 example by cropping, flipping, and normalizing."""
  image = tf.cast(example['image'], tf.float32)
  if distort:
    image = tf.image.random_crop(image, size=crop_shape)
    image = tf.image.random_flip_left_right(image)
  else:
    image = tf.image.resize_with_crop_or_pad(
        image, target_height=crop_shape[1], target_width=crop_shape[2])
  image = tf.image.per_image_standardization(image)
  return (image, example['label'])

def get_federated_cifar100(client_epochs_per_round,
                           train_batch_size,
                           crop_shape=CIFAR_SHAPE,
                           serializable=False):
  """Loads and preprocesses federated CIFAR100 training and testing sets.

  Args:
    client_epochs_per_round: An integer specifying the number of local epochs
      performed per training round by each client. Used to replicate each client
      dataset an appropriate number of times.
    train_batch_size: The batch size for the training dataset.
    crop_shape: An iterable of integers specifying the desired crop
      shape for pre-processing. Must be convertable to a tuple of integers
      (CROP_HEIGHT, CROP_WIDTH, NUM_CHANNELS) which cannot have elements that
      exceed (32, 32, 3), element-wise. The element in the last index should be
      set to 3 to maintain the RGB image structure of the elements.
    serializable: Boolean indicating whether the returned datasets are intended
      to be serialized and shipped across RPC channels. If `True`, stateful
      transformations will be disallowed.

  Returns:
    A tuple of `tff.simulation.ClientData` and `tf.data.Datset` objects.
  """
  if not isinstance(crop_shape, collections.Iterable):
    raise TypeError(
        'Argument crop_shape must be an iterable.')
  crop_shape = tuple(crop_shape)
  if len(crop_shape) != 3:
    raise ValueError('The crop_shape must have length 3, corresponding to a '
                     'tensor of shape [height, width, channels].')
  if not isinstance(serializable, bool):
    raise TypeError(
        'serializable must be a Boolean; you passed {} of type {}.'.format(
            serializable, type(serializable)))
  cifar_train, cifar_test = tff.simulation.datasets.cifar100.load_data()
  train_crop_shape = (train_batch_size,) + crop_shape
  test_crop_shape = (TEST_BATCH_SIZE,) + crop_shape
  train_image_map = functools.partial(
      preprocess_cifar_example,
      crop_shape=train_crop_shape,
      distort=not serializable)
  test_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=test_crop_shape, distort=False)

  def preprocess_train_dataset(dataset):
    """Preprocess CIFAR100 training dataset."""
    if not serializable:
      dataset = dataset.shuffle(buffer_size=NUM_EXAMPLES_PER_CLIENT)
    return dataset.repeat(client_epochs_per_round).batch(
        train_batch_size, drop_remainder=True).map(train_image_map)

  def preprocess_test_dataset(dataset):
    """Preprocess CIFAR100 testing dataset."""
    return dataset.batch(
        TEST_BATCH_SIZE, drop_remainder=False).map(test_image_map)

  cifar_train = cifar_train.preprocess(preprocess_train_dataset)
  cifar_test = preprocess_test_dataset(
      cifar_test.create_tf_dataset_from_all_clients()).cache()
  return cifar_train, cifar_test

def get_centralized_cifar100(train_batch_size, crop_shape=CIFAR_SHAPE):
  """Loads and preprocesses centralized CIFAR100 training and testing sets.

  Args:
    train_batch_size: The batch size for the training dataset.
    crop_shape: An iterable of integers specifying the desired crop
      shape for pre-processing. Must be convertable to a tuple of integers
      (CROP_HEIGHT, CROP_WIDTH, NUM_CHANNELS) which cannot have elements that
      exceed (32, 32, 3), element-wise. The element in the last index should be
      set to 3 to maintain the RGB image structure of the elements.

  Returns:
    A length two tuple of `tf.data.Dataset` objects.
  """
  try:
    crop_shape = tuple(crop_shape)
  except:
    raise ValueError(
        'Argument crop_shape must be able to coerced into a length 3 tuple.')
  if len(crop_shape) != 3:
    raise ValueError('The crop_shape must have length 3, corresponding to a '
                     'tensor of shape [height, width, channels].')
  cifar_train, cifar_test = tff.simulation.datasets.cifar100.load_data()
  train_crop_shape = (train_batch_size,) + crop_shape
  test_crop_shape = (TEST_BATCH_SIZE,) + crop_shape
  train_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=train_crop_shape, distort=True)
  test_image_map = functools.partial(
      preprocess_cifar_example, crop_shape=test_crop_shape, distort=False)

  cifar_train = cifar_train.create_tf_dataset_from_all_clients().shuffle(
      buffer_size=10000).batch(train_batch_size, drop_remainder=True).map(
          train_image_map, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  cifar_test = cifar_test.create_tf_dataset_from_all_clients().batch(
      TEST_BATCH_SIZE, drop_remainder=False).map(test_image_map).cache()

  return cifar_train, cifar_test
<|endoftext|>


<commit_before># coding=utf-8
import concurrent.futures
import itertools
import logging
import os

from colorama import Fore, Style

import settings
from weibo.api import WeiboApi

class Crawler(object):
    def __init__(self, target_url):
        """
        初始化
        :param target_url: 目标微博主页url
        """
        # 目标数据
        self.target = WeiboApi.fetch_user_info(target_url)
        self.uid, self.name = self.target['oid'], self.target['onick']

        # 本地预处理
        self.root = self.__init_folder()
        self.logger = logging.getLogger(self.__class__.__name__)

    def start(self):
        """
        依次下载每一个相册
        :return: None
        """
        self.logger.info(Fore.BLUE + Style.BRIGHT + '开始下载 "%s" 的微博相册' % self.name)

        # 获取每一页的相册列表
        page_size, album_count = 20, 0
        for page in itertools.count(1):
            total, album_list = WeiboApi.fetch_album_list(self.uid, page, page_size)
            if not album_list:
                break

            for album in album_list:
                album_count += 1
                msg = '开始下载第 %d / %d 个微博相册《%s》' % (album_count, total, album['caption'])
                self.logger.info(Fore.BLUE + msg)
                self.__download_album(album)

    def __download_album(self, album):
        """
        下载单个相册
        :param album: 相册数据
        :return: None
        """
        # 相册所有图片的id
        all_photo_ids = WeiboApi.fetch_photo_ids(self.uid, album['album_id'], album['type'])
        self.logger.info(Fore.BLUE + '检测到 %d 张图片' % len(all_photo_ids))

        # 相册所有大图的数据
        all_large_pics = self.__fetch_large_pics(album, all_photo_ids)
        total = len(all_large_pics)

        # 下载所有大图
        with concurrent.futures.ThreadPoolExecutor() as executor:
            album_path = self.__make_album_path(album)

            future_to_large = {
                executor.submit(self.__download_pic, large, album_path): large
                for large in all_large_pics
            }

            for i, future in enumerate(concurrent.futures.as_completed(future_to_large)):
                large = future_to_large[future]
                count_msg = '%d/%d ' % (i + 1, total)
                try:
                    result, path = future.result()
                except Exception as exc:
                    err = '在下载图片 %r 时抛出了异常: %s'
                    self.logger.error(''.join([Fore.RED, count_msg, err % (large['photo_id'], exc)]))
                else:
                    style = result and Style.NORMAL or Style.DIM
                    self.logger.info(''.join([Fore.GREEN, style, count_msg, path]))
            else:
                self.logger.info(Fore.BLUE + '《%s》 已完成' % album['caption'])

    def __fetch_large_pics(self, album, ids):
        """
        获取某相册所有的大图数据
        :param album: 相册
        :param ids: 所有图片的id
        :return: list
        """
        chunk_size, all_large_pics = 50, []

        with concurrent.futures.ThreadPoolExecutor() as executor:
            future_to_chunk = {
                executor.submit(
                    WeiboApi.fetch_large_list,
                    self.uid, ids[i: i + chunk_size], album['type']
                ): i
                for i in range(0, len(ids), chunk_size)
            }

            for future in concurrent.futures.as_completed(future_to_chunk):
                chunk = future_to_chunk[future]
                try:
                    large_list = future.result()
                except Exception as exc:
                    err = '在查询第 %d 块的 %d 个图片的大图时抛出了异常: %s'
                    self.logger.error(Fore.RED + err % (chunk, chunk_size, exc))
                else:
                    all_large_pics.extend(large_list)

        self.logger.info(Fore.BLUE + '检测到 %d 张大图' % len(all_large_pics))
        return all_large_pics

    def __download_pic(self, pic, path):
        """
        下载单个图片
        :param pic: 图片数据
        :param path: 存储路径目录
        :return: bool 下载结果, str 下载路径
        """
        path = os.path.join(path, self.__make_photo_name(pic))
        if not os.path.exists(path):
            url = WeiboApi.make_large_url(pic)
            with open(path, 'wb') as fp:
                fp.write(WeiboApi.get(url).content)
            return True, path
        return False, path

    def __make_album_path(self, album):
        """
        生成并创建相册下载路径
        :param album: 相册数据
        :return: str 下载路径
        """
        album_path = os.path.join(self.root, album['caption'])
        if not os.path.exists(album_path):
            os.mkdir(album_path)
        return album_path

    def __make_photo_name(self, large):
        """
        生成图片文件名
        :param large: 图片数据
        :return: str 文件名
        """
        f, p = large.get('feed_id'), large['pic_name']
        return '_'.join(f and [f, p] or [p])

    def __init_folder(self):
        """
        准备文件夹
        需要检测是否存在相同uid、不同name的情况
        :return: str 该用户的存储文件夹名
        """
        # 根目录
        root = settings.STORE_PATH
        if not os.path.exists(root):
            os.mkdir(root)

        # 用户目录名
        home = '-'.join([self.uid, self.name])
        home_path = os.path.join(root, home)

        # 处理用户更改微博名的情况
        for dir in os.listdir(root):
            if dir.startswith(self.uid):
                if dir != home:
                    src = os.path.join(root, dir)
                    dst = os.path.join(root, home)
                    os.rename(src, dst)
                break
        else:  # 没有已知uid的文件夹
            os.mkdir(home_path)

        return home_path
<commit_msg>优化了下载异常的处理<commit_after># coding=utf-8
import concurrent.futures
import itertools
import logging
import os

from colorama import Fore, Style

import settings
from weibo.api import WeiboApi

class Crawler(object):
    def __init__(self, target_url):
        """
        初始化
        :param target_url: 目标微博主页url
        """
        # 目标数据
        self.target = WeiboApi.fetch_user_info(target_url)
        self.uid, self.name = self.target['oid'], self.target['onick']

        # 本地预处理
        self.root = self.__init_folder()
        self.logger = logging.getLogger(self.__class__.__name__)

    def start(self):
        """
        依次下载每一个相册
        :return: None
        """
        self.logger.info(Fore.BLUE + Style.BRIGHT + '开始下载 "%s" 的微博相册' % self.name)

        # 获取每一页的相册列表
        page_size, album_count = 20, 0
        for page in itertools.count(1):
            total, album_list = WeiboApi.fetch_album_list(self.uid, page, page_size)
            if not album_list:
                break

            for album in album_list:
                album_count += 1
                msg = '开始下载第 %d / %d 个微博相册《%s》' % (album_count, total, album['caption'])
                self.logger.info(Fore.BLUE + msg)
                self.__download_album(album)

    def __download_album(self, album):
        """
        下载单个相册
        :param album: 相册数据
        :return: None
        """
        # 相册所有图片的id
        all_photo_ids = WeiboApi.fetch_photo_ids(self.uid, album['album_id'], album['type'])
        self.logger.info(Fore.BLUE + '检测到 %d 张图片' % len(all_photo_ids))

        # 相册所有大图的数据
        all_large_pics = self.__fetch_large_pics(album, all_photo_ids)
        total = len(all_large_pics)

        # 下载所有大图
        with concurrent.futures.ThreadPoolExecutor() as executor:
            album_path = self.__make_album_path(album)

            future_to_large = {
                executor.submit(self.__download_pic, large, album_path): large
                for large in all_large_pics
            }

            for i, future in enumerate(concurrent.futures.as_completed(future_to_large)):
                large = future_to_large[future]
                count_msg = '%d/%d ' % (i + 1, total)
                try:
                    result, path = future.result()
                except Exception as exc:
                    err = '%s 抛出了异常: %s' % (WeiboApi.make_large_url(large), exc)
                    self.logger.error(''.join([Fore.RED, count_msg, err]))
                else:
                    style = result and Style.NORMAL or Style.DIM
                    self.logger.info(''.join([Fore.GREEN, style, count_msg, path]))
            else:
                self.logger.info(Fore.BLUE + '《%s》 已完成' % album['caption'])

    def __fetch_large_pics(self, album, ids):
        """
        获取某相册所有的大图数据
        :param album: 相册
        :param ids: 所有图片的id
        :return: list
        """
        chunk_size, all_large_pics = 50, []

        with concurrent.futures.ThreadPoolExecutor() as executor:
            future_to_chunk = {
                executor.submit(
                    WeiboApi.fetch_large_list,
                    self.uid, ids[i: i + chunk_size], album['type']
                ): i
                for i in range(0, len(ids), chunk_size)
            }

            for future in concurrent.futures.as_completed(future_to_chunk):
                chunk = future_to_chunk[future]
                try:
                    large_list = future.result()
                except Exception as exc:
                    err = '在查询第 %d 块的 %d 个图片的大图时抛出了异常: %s'
                    self.logger.error(Fore.RED + err % (chunk, chunk_size, exc))
                else:
                    all_large_pics.extend(large_list)

        self.logger.info(Fore.BLUE + '检测到 %d 张大图' % len(all_large_pics))
        return all_large_pics

    def __download_pic(self, pic, path):
        """
        下载单个图片
        :param pic: 图片数据
        :param path: 存储路径目录
        :return: bool 下载结果, str 下载路径
        """
        path = os.path.join(path, self.__make_photo_name(pic))
        if not os.path.exists(path):
            url = WeiboApi.make_large_url(pic)
            response = WeiboApi.get(url, timeout=2)
            with open(path, 'wb') as fp:
                fp.write(response.content)
            return True, path
        return False, path

    def __make_album_path(self, album):
        """
        生成并创建相册下载路径
        :param album: 相册数据
        :return: str 下载路径
        """
        album_path = os.path.join(self.root, album['caption'])
        if not os.path.exists(album_path):
            os.mkdir(album_path)
        return album_path

    def __make_photo_name(self, large):
        """
        生成图片文件名
        :param large: 图片数据
        :return: str 文件名
        """
        f, p = large.get('feed_id'), large['pic_name']
        return '_'.join(f and [f, p] or [p])

    def __init_folder(self):
        """
        准备文件夹
        需要检测是否存在相同uid、不同name的情况
        :return: str 该用户的存储文件夹名
        """
        # 根目录
        root = settings.STORE_PATH
        if not os.path.exists(root):
            os.mkdir(root)

        # 用户目录名
        home = '-'.join([self.uid, self.name])
        home_path = os.path.join(root, home)

        # 处理用户更改微博名的情况
        for dir in os.listdir(root):
            if dir.startswith(self.uid):
                if dir != home:
                    src = os.path.join(root, dir)
                    dst = os.path.join(root, home)
                    os.rename(src, dst)
                break
        else:  # 没有已知uid的文件夹
            os.mkdir(home_path)

        return home_path
<|endoftext|>


<commit_before>import atexit, sys, os, time

from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer

from .client import BaseClient, MQTTException

from ambulances.models import Ambulance
from ambulances.serializers import AmbulanceSerializer

from hospital.models import Equipment, HospitalEquipment, Hospital
from hospital.serializers import HospitalSerializer, \
    HospitalEquipmentSerializer, EquipmentSerializer

from login.serializers import ExtendedProfileSerializer

# MessagePublishClient class

class MessagePublishClient(BaseClient):

    def publish_profile(self, profile, **kwargs):
        pass

    def publish_ambulance(self, ambulance, **kwargs):
        pass

    def remove_ambulance(self, ambulance, **kwargs):
        pass

    def publish_hospital(self, hospital, **kwargs):
        pass

    def remove_hospital(self, hospital, **kwargs):
        pass

    def publish_hospital_metadata(self, hospital, **kwargs):
        pass

    def publish_hospital_equipment(self, hospital, **kwargs):
        pass

    def remove_hospital_equipment(self, hospital, **kwargs):
        pass

# Uses Alex Martelli's Borg for making PublishClient act like a singleton

class PublishClient(MessagePublishClient):

    def on_disconnect(self, client, userdata, rc):
        # Exception is generated only if never connected
        if not self.connected and rc:
            raise MQTTException('Disconnected',
                                rc)

    def publish_topic(self, topic, serializer, qos=0, retain=False):
        # Publish to topic
        self.publish(topic,
                     JSONRenderer().render(serializer.data),
                     qos=qos,
                     retain=retain)

    def remove_topic(self, topic, serializer, qos=0):
        # Publish null to retained topic
        self.publish(topic,
                     null,
                     qos=qos,
                     retain=True)

    def publish_profile(self, profile, qos=2, retain=True):
        self.publish_topic('user/{}/profile'.format(profile.user.username),
                          ExtendedProfileSerializer(profile),
                          qos=qos,
                          retain=retain)

    def publish_ambulance(self, ambulance, qos=2, retain=True):
        self.publish_topic('ambulance/{}/data'.format(ambulance.id),
                          AmbulanceSerializer(ambulance),
                          qos=qos,
                          retain=retain)

    def remove_ambulance(self, ambulance):
        self.remove_topic('ambulance/{}/data'.format(ambulance.id))

    def publish_hospital(self, hospital, qos=2, retain=True):
        self.publish_topic('hospital/{}/data'.format(hospital.id),
                          HospitalSerializer(hospital),
                          qos=qos,
                          retain=retain)

    def remove_hospital(self, hospital):
        self.remove_topic('hospital/{}/data'.format(hospital.id))
        self.remove_topic('hospital/{}/metadata'.format(hospital.id))

    def publish_hospital_metadata(self, hospital, qos=2, retain=True):
        hospital_equipment = hospital.hospitalequipment_set.values('equipment')
        equipment = Equipment.objects.filter(id__in=hospital_equipment)
        self.publish_topic('hospital/{}/metadata'.format(hospital.id),
                          EquipmentSerializer(equipment, many=True),
                          qos=qos,
                          retain=retain)

    def publish_hospital_equipment(self, equipment, qos=2, retain=True):
        self.publish_topic('hospital/{}/equipment/{}/data'.format(equipment.hospital.id,
                                                                 equipment.equipment.name),
                          HospitalEquipmentSerializer(equipment),
                          qos=qos,
                          retain=retain)

    def remove_hospital_equipment(self, equipment):
        self.remove_topic('hospital/{}/equipment/{}/data'.format(equipment.hospital.id,
                                                                 equipment.equipment.name))

class SingletonPublishClient(PublishClient):

    _shared_state = {}

    def __init__(self, **kwargs):

        print('SingletonPublishClient')

        # Makes sure it is a singleton
        self.__dict__ = self._shared_state

        # Do not initialize if already initialized
        if not self.__dict__ == {}:
            print('*** SKIP INITIALIZATION ***')
            return

        # initialize BaseClient
        super().__init__(**kwargs)

# to be used with the lazy constructor

def get_client():

    # Start client
    from django.core.management.base import OutputWrapper
    from django.core.management.color import color_style, no_style
    from django.conf import settings

    stdout = OutputWrapper(sys.stdout)
    style = color_style()

    # Instantiate broker
    broker = {
        'HOST': 'localhost',
        'PORT': 1883,
        'KEEPALIVE': 60,
        'CLEAN_SESSION': True
    }

    broker.update(settings.MQTT)
    broker['CLIENT_ID'] = 'mqtt_publish_' + str(os.getpid())

    try:

        # try to connect
        print('Connecting to MQTT brocker...')
        client = SingletonPublishClient(broker, stdout, style,
                                        verbosity=1, debug=True)

        # wait for connection
        while not client.connected:
            client.loop()

        # start loop
        client.loop_start()

        # register atexit handler to make sure it disconnects at exit
        atexit.register(client.disconnect)

    except Exception as e:

        print('Exception {}'.format(e))
        print('Could not connect to MQTT brocker. Using dummy client...')

        client = MessagePublishClient()

    return client

<commit_msg>Create separate hospital app<commit_after>import atexit, sys, os, time

from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer

from .client import BaseClient, MQTTException

from ambulances.models import Ambulance
from ambulances.serializers import AmbulanceSerializer

from hospital.models import Equipment, HospitalEquipment, Hospital
from hospital.serializers import HospitalSerializer, \
    HospitalEquipmentSerializer, EquipmentSerializer

from login.serializers import ExtendedProfileSerializer

# MessagePublishClient class

class MessagePublishClient(BaseClient):

    def publish_profile(self, profile, **kwargs):
        pass

    def publish_ambulance(self, ambulance, **kwargs):
        pass

    def remove_ambulance(self, ambulance, **kwargs):
        pass

    def publish_hospital(self, hospital, **kwargs):
        pass

    def remove_hospital(self, hospital, **kwargs):
        pass

    def publish_hospital_metadata(self, hospital, **kwargs):
        pass

    def publish_hospital_equipment(self, hospital, **kwargs):
        pass

    def remove_hospital_equipment(self, hospital, **kwargs):
        pass

# Uses Alex Martelli's Borg for making PublishClient act like a singleton

class PublishClient(MessagePublishClient):

    def on_disconnect(self, client, userdata, rc):
        # Exception is generated only if never connected
        if not self.connected and rc:
            raise MQTTException('Disconnected',
                                rc)

    def publish_topic(self, topic, serializer, qos=0, retain=False):
        # Publish to topic
        self.publish(topic,
                     JSONRenderer().render(serializer.data),
                     qos=qos,
                     retain=retain)

    def remove_topic(self, topic, serializer, qos=0):
        # Publish null to retained topic
        self.publish(topic,
                     null,
                     qos=qos,
                     retain=True)

    def publish_profile(self, profile, qos=2, retain=True):
        self.publish_topic('user/{}/profile'.format(profile.user.username),
                          ExtendedProfileSerializer(profile),
                          qos=qos,
                          retain=retain)

    def publish_ambulance(self, ambulance, qos=2, retain=True):
        self.publish_topic('ambulance/{}/data'.format(ambulance.id),
                          AmbulanceSerializer(ambulance),
                          qos=qos,
                          retain=retain)

    def remove_ambulance(self, ambulance):
        self.remove_topic('ambulance/{}/data'.format(ambulance.id))

    def publish_hospital(self, hospital, qos=2, retain=True):
        self.publish_topic('hospital/{}/data'.format(hospital.id),
                          HospitalSerializer(hospital),
                          qos=qos,
                          retain=retain)

    def remove_hospital(self, hospital):
        self.remove_topic('hospital/{}/data'.format(hospital.id))
        self.remove_topic('hospital/{}/metadata'.format(hospital.id))

    def publish_hospital_metadata(self, hospital, qos=2, retain=True):
        hospital_equipment = hospital.hospitalequipment_set.values('equipment')
        equipment = Equipment.objects.filter(id__in=hospital_equipment)
        self.publish_topic('hospital/{}/metadata'.format(hospital.id),
                          EquipmentSerializer(equipment, many=True),
                          qos=qos,
                          retain=retain)

    def publish_hospital_equipment(self, equipment, qos=2, retain=True):
        self.publish_topic('hospital/{}/equipment/{}/data'.format(equipment.hospital.id,
                                                                 equipment.equipment.name),
                          HospitalEquipmentSerializer(equipment),
                          qos=qos,
                          retain=retain)

    def remove_hospital_equipment(self, equipment):
        self.remove_topic('hospital/{}/equipment/{}/data'.format(equipment.hospital.id,
                                                                 equipment.equipment.name))

class SingletonPublishClient(PublishClient):

    _shared_state = {}

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

        print('SingletonPublishClient')

        # Makes sure it is a singleton
        self.__dict__ = self._shared_state

        # Do not initialize if already initialized
        if not self.__dict__ == {}:
            print('*** SKIP INITIALIZATION ***')
            return

        # initialize BaseClient
        super().__init__(*args, **kwargs)

# to be used with the lazy constructor

def get_client():

    # Start client
    from django.core.management.base import OutputWrapper
    from django.core.management.color import color_style, no_style
    from django.conf import settings

    stdout = OutputWrapper(sys.stdout)
    style = color_style()

    # Instantiate broker
    broker = {
        'HOST': 'localhost',
        'PORT': 1883,
        'KEEPALIVE': 60,
        'CLEAN_SESSION': True
    }

    broker.update(settings.MQTT)
    broker['CLIENT_ID'] = 'mqtt_publish_' + str(os.getpid())

    try:

        # try to connect
        print('Connecting to MQTT brocker...')
        client = SingletonPublishClient(broker, stdout, style,
                                        verbosity=1, debug=True)

        # wait for connection
        while not client.connected:
            client.loop()

        # start loop
        client.loop_start()

        # register atexit handler to make sure it disconnects at exit
        atexit.register(client.disconnect)

    except Exception as e:

        print('Exception {}'.format(e))
        print('Could not connect to MQTT brocker. Using dummy client...')

        client = MessagePublishClient()

    return client

<|endoftext|>


<commit_before>import django
from django import template
from django.db.models import Q
from django.template.loader import render_to_string
from threadedcomments.compat import BASE_APP, django_comments as comments
from threadedcomments.util import annotate_tree_properties, fill_tree as real_fill_tree

if BASE_APP == 'django.contrib.comments':
    from django.contrib.comments.templatetags.comments import BaseCommentNode
elif BASE_APP == 'django_comments':
    from django_comments.templatetags.comments import BaseCommentNode
else:
    raise NotImplementedError()

register = template.Library()

class BaseThreadedCommentNode(BaseCommentNode):
    def __init__(self, parent=None, flat=False, root_only=False, newest=False, **kwargs):
        self.parent = parent
        self.flat = flat
        self.root_only = root_only
        self.newest = newest
        super(BaseThreadedCommentNode, self).__init__(**kwargs)

    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if len(tokens) > 2:
            if tokens[1] != 'for':
                raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        extra_kw = {}
        if tokens[-1] in ('flat', 'root_only', 'newest'):
            extra_kw[str(tokens.pop())] = True

        if len(tokens) == 5:
            # {% get_whatever for obj as varname %}
            if tokens[3] != 'as':
                raise template.TemplateSyntaxError("Fourth argument in %r must be 'as'" % tokens[0])

            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                as_varname=tokens[4],
                **extra_kw
            )
        elif len(tokens) == 6:
            # {% get_whatever for app.model pk as varname %}
            if tokens[4] != 'as':
                raise template.TemplateSyntaxError("Fourth argument in %r must be 'as'" % tokens[0])

            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                as_varname=tokens[5],
                **extra_kw
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes either 5 or 6 arguments" % (tokens[0],))

    def get_queryset(self, context):
        qs = super(BaseThreadedCommentNode, self).get_queryset(context)
        if self.flat:
            qs = qs.order_by('-submit_date')
        elif self.root_only:
            qs = qs.exclude(parent__isnull=False).order_by('-submit_date')
        elif self.newest:
            qs = qs.order_by('-newest_activity', 'tree_path')
        return qs

    # For older Django (1.5) versions:
    def get_query_set(self, context):
        qs = super(BaseThreadedCommentNode, self).get_query_set(context)
        if self.flat:
            qs = qs.order_by('-submit_date')
        elif self.root_only:
            qs = qs.exclude(parent__isnull=False).order_by('-submit_date')
        elif self.newest:
            qs = qs.order_by('-newest_activity', 'tree_path')

        return qs

class CommentListNode(BaseThreadedCommentNode):
    """
    Insert a list of comments into the context.
    """
    def get_context_value_from_queryset(self, context, qs):
        return list(qs)

class CommentCountNode(CommentListNode):
    """
    Insert a count of comments into the context.
    """
    def get_context_value_from_queryset(self, context, qs):
        return qs.count()

class CommentFormNode(BaseThreadedCommentNode):
    """
    Insert a form for the comment model into the context.
    """
    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % (tokens[0],))

        if len(tokens) < 7:
            # Default get_comment_form code
            return super(CommentFormNode, cls).handle_token(parser, token)
        elif len(tokens) == 7:
            # {% get_comment_form for [object] as [varname] with [parent_id] %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have a 'with' as the last but one argument." % (tokens[0],))
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                as_varname=tokens[4],
                parent=parser.compile_filter(tokens[6]),
            )
        elif len(tokens) == 8:
            # {% get_comment_form for [app].[model] [object_id] as [varname] with [parent_id] %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have a 'with' as the last but one argument." % (tokens[0],))
            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                as_varname=tokens[5],
                parent=parser.compile_filter(tokens[7]),
            )

    def get_form(self, context):
        parent_id = None
        if self.parent:
            parent_id = self.parent.resolve(context, ignore_failures=True)

        obj = self.get_object(context)
        if obj:
            return comments.get_form()(obj, parent=parent_id)
        else:
            return None

    def get_object(self, context):
        if self.object_expr:
            try:
                return self.object_expr.resolve(context)
            except template.VariableDoesNotExist:
                return None
        else:
            object_pk = self.object_pk_expr.resolve(context, ignore_failures=True)
            return self.ctype.get_object_for_this_type(pk=object_pk)

    def render(self, context):
        context[self.as_varname] = self.get_form(context)
        return ''

class RenderCommentFormNode(CommentFormNode):
    @classmethod
    def handle_token(cls, parser, token):
        """
        Class method to parse render_comment_form and return a Node.
        """
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        if len(tokens) == 3:
            # {% render_comment_form for obj %}
            return cls(object_expr=parser.compile_filter(tokens[2]))
        elif len(tokens) == 4:
            # {% render_comment_form for app.model object_pk %}
            return cls(
                ctype=BaseCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3])
            )
        elif len(tokens) == 5:
            # {% render_comment_form for obj with parent_id %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have 'with' as the last but one argument" % (tokens[0],))
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                parent=parser.compile_filter(tokens[4])
            )
        elif len(tokens) == 6:
            # {% render_comment_form for app.model object_pk with parent_id %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have 'with' as the last but one argument" % (tokens[0],))
            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                parent=parser.compile_filter(tokens[5])
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes 2 to 5 arguments" % (tokens[0],))

    def render(self, context):
        ctype, object_pk = self.get_target_ctype_pk(context)
        if object_pk:
            template_search_list = (
                "comments/%s/%s/form.html" % (ctype.app_label, ctype.model),
                "comments/%s/form.html" % ctype.app_label,
                "comments/form.html",
            )
            context.push()
            form_str = render_to_string(
                template_search_list,
                {"form" : self.get_form(context)},
                context
            )
            context.pop()
            return form_str
        else:
            return ''

class RenderCommentListNode(CommentListNode):
    """
    Render the comments list.
    """
    # By having this class added, this module is a drop-in replacement for ``{% load comments %}``.

    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        extra_kw = {}
        if tokens[-1] in ('flat', 'root_only', 'newest'):
            extra_kw[str(tokens.pop())] = True

        if len(tokens) == 3:
            # {% render_comment_list for obj %}
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                **extra_kw
            )
        elif len(tokens) == 4:
            # {% render_comment_list for app.models pk %}
            return cls(
                ctype = BaseCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr = parser.compile_filter(tokens[3]),
                **extra_kw
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes either 2 or 3 arguments" % (tokens[0],))

    def render(self, context):
        ctype, object_pk = self.get_target_ctype_pk(context)
        if object_pk:
            template_search_list = [
                "comments/%s/%s/list.html" % (ctype.app_label, ctype.model),
                "comments/%s/list.html" % ctype.app_label,
                "comments/list.html"
            ]
            # For newer Django (1.6) versions
            if django.VERSION >= (1,6):
                qs = self.get_queryset(context)
            else:
                qs = self.get_query_set(context)
            context.push()
            liststr = render_to_string(template_search_list, {
                "comment_list" : self.get_context_value_from_queryset(context, qs)
            }, context)
            context.pop()
            return liststr
        else:
            return ''

@register.tag
def get_comment_count(parser, token):
    """
    Gets the comment count for the given params and populates the template
    context with a variable containing that value, whose name is defined by the
    'as' clause.

    Syntax::

        {% get_comment_count for [object] as [varname]  %}
        {% get_comment_count for [app].[model] [object_id] as [varname]  %}

    Example usage::

        {% get_comment_count for event as comment_count %}
        {% get_comment_count for calendar.event event.id as comment_count %}
        {% get_comment_count for calendar.event 17 as comment_count %}

    """
    return CommentCountNode.handle_token(parser, token)

@register.tag
def get_comment_list(parser, token):
    """
    Gets the list of comments for the given params and populates the template
    context with a variable containing that value, whose name is defined by the
    'as' clause.

    Syntax::

        {% get_comment_list for [object] as [varname] %}
        {% get_comment_list for [object] as [varname] [flat|root_only] %}
        {% get_comment_list for [app].[model] [object_id] as [varname] %}
        {% get_comment_list for [app].[model] [object_id] as [varname] [flat|root_only] %}

    Example usage::

        {% get_comment_list for event as comment_list %}
        {% for comment in comment_list %}
            ...
        {% endfor %}
        {% get_comment_list for event as comment_list flat %}

    """
    return CommentListNode.handle_token(parser, token)

@register.tag
def render_comment_list(parser, token):
    """
    Render the comment list (as returned by ``{% get_comment_list %}``)
    through the ``comments/list.html`` template

    Syntax::

        {% render_comment_list for [object] %}
        {% render_comment_list for [app].[model] [object_id] %}

        {% render_comment_list for [object] [flat|root_only] %}
        {% render_comment_list for [app].[model] [object_id] [flat|root_only] %}

    Example usage::

        {% render_comment_list for event %}

    """
    return RenderCommentListNode.handle_token(parser, token)

@register.tag
def get_comment_form(parser, token):
    """
    Get a (new) form object to post a new comment.

    Syntax::

        {% get_comment_form for [object] as [varname] %}
        {% get_comment_form for [object] as [varname] with [parent_id] %}
        {% get_comment_form for [app].[model] [object_id] as [varname] %}
        {% get_comment_form for [app].[model] [object_id] as [varname] with [parent_id] %}
    """
    return CommentFormNode.handle_token(parser, token)

@register.tag
def render_comment_form(parser, token):
    """
    Render the comment form (as returned by ``{% render_comment_form %}``)
    through the ``comments/form.html`` template.

    Syntax::

        {% render_comment_form for [object] %}
        {% render_comment_form for [object] with [parent_id] %}
        {% render_comment_form for [app].[model] [object_id] %}
        {% render_comment_form for [app].[model] [object_id] with [parent_id] %}
    """
    return RenderCommentFormNode.handle_token(parser, token)

@register.filter
def annotate_tree(comments):
    """
    Add ``open``, ``close`` properties to the comments, to render the tree.

    Syntax::

        {% for comment in comment_list|annotate_tree %}
            {% ifchanged comment.parent_id %}{% else %}
{% endifchanged %}
            {% if not comment.open and not comment.close %}
{% endif %}
            {% if comment.open %}
{% endif %}

                ...
            {% for close in comment.close %}

{% endfor %}
        {% endfor %}

    When the :func:`fill_tree` filter, place the ``annotate_tree`` code after it::

        {% for comment in comment_list|fill_tree|annotate_tree %}
            ...
        {% endfor %}
    """
    return annotate_tree_properties(comments)

@register.filter
def fill_tree(comments):
    """
    When paginating the comments, insert the parent nodes of the first comment.

    Syntax::

        {% for comment in comment_list|annotate_tree %}
            ...
        {% endfor %}
    """
    return real_fill_tree(comments)
<commit_msg>'limit' and 'admin' parameter for commentlist<commit_after>import django
from django import template
from django.conf import settings
from django.db.models import Q
from django.template.loader import render_to_string
from threadedcomments.compat import BASE_APP, django_comments as comments
from threadedcomments.util import annotate_tree_properties, fill_tree as real_fill_tree
from django.utils.encoding import smart_text

if BASE_APP == 'django.contrib.comments':
    from django.contrib.comments.templatetags.comments import BaseCommentNode
elif BASE_APP == 'django_comments':
    from django_comments.templatetags.comments import BaseCommentNode
else:
    raise NotImplementedError()

register = template.Library()

class AdminOverrideCommentNode(BaseCommentNode):
    def get_queryset(self, context):
        ctype, object_pk = self.get_target_ctype_pk(context)
        if not object_pk:
            return self.comment_model.objects.none()

        qs = self.comment_model.objects.filter(
            content_type=ctype,
            object_pk=smart_text(object_pk),
            site__pk=settings.SITE_ID,
        )

        # The is_public and is_removed fields are implementation details of the
        # built-in comment model's spam filtering system, so they might not
        # be present on a custom comment model subclass. If they exist, we
        # should filter on them.
        field_names = [f.name for f in self.comment_model._meta.fields]
        if not self.admin:
            if 'is_public' in field_names:
                qs = qs.filter(is_public=True)
            if getattr(settings, 'COMMENTS_HIDE_REMOVED', True) and 'is_removed' in field_names:
                qs = qs.filter(is_removed=False)

        return qs

class BaseThreadedCommentNode(AdminOverrideCommentNode):
    def __init__(self, parent=None, flat=False, root_only=False, newest=False, limit=False, admin=False, **kwargs):
        self.parent = parent
        self.flat = flat
        self.root_only = root_only
        self.newest = newest
        self.limit = limit
        self.admin = admin
        super(BaseThreadedCommentNode, self).__init__(**kwargs)

    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if len(tokens) > 2:
            if tokens[1] != 'for':
                raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        extra_kw = {}
        if tokens[-1] in ('flat', 'root_only', 'newest', 'admin'):
            extra_kw[str(tokens.pop())] = True

        if len(tokens) == 5:
            # {% get_whatever for obj as varname %}
            if tokens[3] != 'as':
                raise template.TemplateSyntaxError("Fourth argument in %r must be 'as'" % tokens[0])

            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                as_varname=tokens[4],
                **extra_kw
            )
        elif len(tokens) == 6:
            # {% get_whatever for app.model pk as varname %}
            if tokens[4] != 'as':
                raise template.TemplateSyntaxError("Fourth argument in %r must be 'as'" % tokens[0])

            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                as_varname=tokens[5],
                **extra_kw
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes either 5 or 6 arguments" % (tokens[0],))

    def get_queryset(self, context):
        qs = super(BaseThreadedCommentNode, self).get_queryset(context)
        if self.limit:
            parent_qs = qs.exclude(parent__isnull=False).order_by('-newest_activity', '-submit_date').values_list('pk', flat=True)[:self.limit]
            qs = qs.filter(Q(parent_id__in=parent_qs) | Q(pk__in=parent_qs)).distinct()
        if self.flat:
            qs = qs.order_by('-submit_date')
        elif self.root_only:
            qs = qs.exclude(parent__isnull=False).order_by('-submit_date')
        elif self.newest:
            qs = qs.order_by('-newest_activity', 'tree_path')
        return qs

    # For older Django (1.5) versions:
    def get_query_set(self, context):
        qs = super(BaseThreadedCommentNode, self).get_query_set(context)
        if self.limit:
            parent_qs = qs.exclude(parent__isnull=False).order_by('-newest_activity', '-submit_date').values_list('pk', flat=True)[:self.limit]
            qs = qs.filter(Q(parent_id__in=parent_qs) | Q(pk__in=parent_qs)).distinct()
        if self.flat:
            qs = qs.order_by('-submit_date')
        elif self.root_only:
            qs = qs.exclude(parent__isnull=False).order_by('-submit_date')
        elif self.newest:
            qs = qs.order_by('-newest_activity', 'tree_path')

        return qs

class CommentListNode(BaseThreadedCommentNode):
    """
    Insert a list of comments into the context.
    """
    def get_context_value_from_queryset(self, context, qs):
        return list(qs)

class CommentCountNode(CommentListNode):
    """
    Insert a count of comments into the context.
    """
    def get_context_value_from_queryset(self, context, qs):
        return qs.count()

class CommentFormNode(BaseThreadedCommentNode):
    """
    Insert a form for the comment model into the context.
    """
    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % (tokens[0],))

        if len(tokens) < 7:
            # Default get_comment_form code
            return super(CommentFormNode, cls).handle_token(parser, token)
        elif len(tokens) == 7:
            # {% get_comment_form for [object] as [varname] with [parent_id] %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have a 'with' as the last but one argument." % (tokens[0],))
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                as_varname=tokens[4],
                parent=parser.compile_filter(tokens[6]),
            )
        elif len(tokens) == 8:
            # {% get_comment_form for [app].[model] [object_id] as [varname] with [parent_id] %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have a 'with' as the last but one argument." % (tokens[0],))
            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                as_varname=tokens[5],
                parent=parser.compile_filter(tokens[7]),
            )

    def get_form(self, context):
        parent_id = None
        if self.parent:
            parent_id = self.parent.resolve(context, ignore_failures=True)

        obj = self.get_object(context)
        if obj:
            return comments.get_form()(obj, parent=parent_id)
        else:
            return None

    def get_object(self, context):
        if self.object_expr:
            try:
                return self.object_expr.resolve(context)
            except template.VariableDoesNotExist:
                return None
        else:
            object_pk = self.object_pk_expr.resolve(context, ignore_failures=True)
            return self.ctype.get_object_for_this_type(pk=object_pk)

    def render(self, context):
        context[self.as_varname] = self.get_form(context)
        return ''

class RenderCommentFormNode(CommentFormNode):
    @classmethod
    def handle_token(cls, parser, token):
        """
        Class method to parse render_comment_form and return a Node.
        """
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        if len(tokens) == 3:
            # {% render_comment_form for obj %}
            return cls(object_expr=parser.compile_filter(tokens[2]))
        elif len(tokens) == 4:
            # {% render_comment_form for app.model object_pk %}
            return cls(
                ctype=BaseCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3])
            )
        elif len(tokens) == 5:
            # {% render_comment_form for obj with parent_id %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have 'with' as the last but one argument" % (tokens[0],))
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                parent=parser.compile_filter(tokens[4])
            )
        elif len(tokens) == 6:
            # {% render_comment_form for app.model object_pk with parent_id %}
            if tokens[-2] != u'with':
                raise template.TemplateSyntaxError("%r tag must have 'with' as the last but one argument" % (tokens[0],))
            return cls(
                ctype=BaseThreadedCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr=parser.compile_filter(tokens[3]),
                parent=parser.compile_filter(tokens[5])
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes 2 to 5 arguments" % (tokens[0],))

    def render(self, context):
        ctype, object_pk = self.get_target_ctype_pk(context)
        if object_pk:
            template_search_list = (
                "comments/%s/%s/form.html" % (ctype.app_label, ctype.model),
                "comments/%s/form.html" % ctype.app_label,
                "comments/form.html",
            )
            context.push()
            form_str = render_to_string(
                template_search_list,
                {"form" : self.get_form(context)},
                context
            )
            context.pop()
            return form_str
        else:
            return ''

class RenderCommentListNode(CommentListNode):
    """
    Render the comments list.
    """
    # By having this class added, this module is a drop-in replacement for ``{% load comments %}``.

    @classmethod
    def handle_token(cls, parser, token):
        tokens = token.contents.split()
        if tokens[1] != 'for':
            raise template.TemplateSyntaxError("Second argument in %r tag must be 'for'" % tokens[0])

        extra_kw = {}
        if tokens[-2] == 'limit':
            if tokens[-1].isdigit():
                extra_kw['limit'] = tokens.pop(-1)  # removes limit integer
                tokens.pop(-1)  # removes 'limit'
            else:
                raise template.TemplateSyntaxError(
                    "When using 'limit' with %r tag, it needs to be followed by a positive integer" % (tokens[0],))
        if tokens[-1] in ('flat', 'root_only', 'newest', 'admin'):
            extra_kw[str(tokens.pop())] = True

        if len(tokens) == 3:
            # {% render_comment_list for obj %}
            return cls(
                object_expr=parser.compile_filter(tokens[2]),
                **extra_kw
            )
        elif len(tokens) == 4:
            # {% render_comment_list for app.models pk %}
            return cls(
                ctype = BaseCommentNode.lookup_content_type(tokens[2], tokens[0]),
                object_pk_expr = parser.compile_filter(tokens[3]),
                **extra_kw
            )
        else:
            raise template.TemplateSyntaxError("%r tag takes either 2 or 3 arguments" % (tokens[0],))

    def render(self, context):
        ctype, object_pk = self.get_target_ctype_pk(context)
        if object_pk:
            template_search_list = [
                "comments/%s/%s/list.html" % (ctype.app_label, ctype.model),
                "comments/%s/list.html" % ctype.app_label,
                "comments/list.html"
            ]
            # For newer Django (1.6) versions
            if django.VERSION >= (1,6):
                qs = self.get_queryset(context)
            else:
                qs = self.get_query_set(context)
            context.push()
            liststr = render_to_string(template_search_list, {
                "comment_list" : self.get_context_value_from_queryset(context, qs)
            }, context)
            context.pop()
            return liststr
        else:
            return ''

@register.tag
def get_comment_count(parser, token):
    """
    Gets the comment count for the given params and populates the template
    context with a variable containing that value, whose name is defined by the
    'as' clause.

    Syntax::

        {% get_comment_count for [object] as [varname]  %}
        {% get_comment_count for [app].[model] [object_id] as [varname]  %}

    Example usage::

        {% get_comment_count for event as comment_count %}
        {% get_comment_count for calendar.event event.id as comment_count %}
        {% get_comment_count for calendar.event 17 as comment_count %}

    """
    return CommentCountNode.handle_token(parser, token)

@register.tag
def get_comment_list(parser, token):
    """
    Gets the list of comments for the given params and populates the template
    context with a variable containing that value, whose name is defined by the
    'as' clause.

    Syntax::

        {% get_comment_list for [object] as [varname] %}
        {% get_comment_list for [object] as [varname] [flat|root_only] %}
        {% get_comment_list for [app].[model] [object_id] as [varname] %}
        {% get_comment_list for [app].[model] [object_id] as [varname] [flat|root_only] %}

    Example usage::

        {% get_comment_list for event as comment_list %}
        {% for comment in comment_list %}
            ...
        {% endfor %}
        {% get_comment_list for event as comment_list flat %}

    """
    return CommentListNode.handle_token(parser, token)

@register.tag
def render_comment_list(parser, token):
    """
    Render the comment list (as returned by ``{% get_comment_list %}``)
    through the ``comments/list.html`` template

    Syntax::

        {% render_comment_list for [object] %}
        {% render_comment_list for [app].[model] [object_id] %}

        {% render_comment_list for [object] [flat|root_only] %}
        {% render_comment_list for [app].[model] [object_id] [flat|root_only] %}

    Example usage::

        {% render_comment_list for event %}

    """
    return RenderCommentListNode.handle_token(parser, token)

@register.tag
def get_comment_form(parser, token):
    """
    Get a (new) form object to post a new comment.

    Syntax::

        {% get_comment_form for [object] as [varname] %}
        {% get_comment_form for [object] as [varname] with [parent_id] %}
        {% get_comment_form for [app].[model] [object_id] as [varname] %}
        {% get_comment_form for [app].[model] [object_id] as [varname] with [parent_id] %}
    """
    return CommentFormNode.handle_token(parser, token)

@register.tag
def render_comment_form(parser, token):
    """
    Render the comment form (as returned by ``{% render_comment_form %}``)
    through the ``comments/form.html`` template.

    Syntax::

        {% render_comment_form for [object] %}
        {% render_comment_form for [object] with [parent_id] %}
        {% render_comment_form for [app].[model] [object_id] %}
        {% render_comment_form for [app].[model] [object_id] with [parent_id] %}
    """
    return RenderCommentFormNode.handle_token(parser, token)

@register.filter
def annotate_tree(comments):
    """
    Add ``open``, ``close`` properties to the comments, to render the tree.

    Syntax::

        {% for comment in comment_list|annotate_tree %}
            {% ifchanged comment.parent_id %}{% else %}
{% endifchanged %}
            {% if not comment.open and not comment.close %}
{% endif %}
            {% if comment.open %}
{% endif %}

                ...
            {% for close in comment.close %}

{% endfor %}
        {% endfor %}

    When the :func:`fill_tree` filter, place the ``annotate_tree`` code after it::

        {% for comment in comment_list|fill_tree|annotate_tree %}
            ...
        {% endfor %}
    """
    return annotate_tree_properties(comments)

@register.filter
def fill_tree(comments):
    """
    When paginating the comments, insert the parent nodes of the first comment.

    Syntax::

        {% for comment in comment_list|annotate_tree %}
            ...
        {% endfor %}
    """
    return real_fill_tree(comments)
<|endoftext|>