<commit_before>#!/usr/bin/env python3
"""
Update firmware using CVRA bootloading protocol.
"""
import page
import logging
import commands
import msgpack
from zlib import crc32
from sys import exit

import utils
import progressbar
import sys

CHUNK_SIZE = 2048

def parse_commandline_args(args=None):
    """
    Parses the program commandline arguments.
    Args must be an array containing all arguments.
    """
    parser = utils.ConnectionArgumentParser(description=__doc__)
    parser.add_argument('-b', '--binary', dest='binary_file',
                        help='Path to the binary file to upload',
                        required=True,
                        metavar='FILE')

    parser.add_argument('-a', '--base-address', dest='base_address',
                        help='Base address of the firmware',
                        metavar='ADDRESS',
                        required=True,
                        # automatically convert value to hex
                        type=lambda s: int(s, 16))

    parser.add_argument('-c', '--device-class',
                        dest='device_class',
                        help='Device class to flash', required=True)
    parser.add_argument('-r', '--run',
                        help='Run application after flashing',
                        action='store_true')
    parser.add_argument("ids",
                        metavar='DEVICEID',
                        nargs='+', type=int,
                        help="Device IDs to flash")

    return parser.parse_args(args)

def flash_binary(fdesc, binary, base_address, device_class, destinations,
                 page_size=2048):
    """
    Writes a full binary to the flash using the given file descriptor.

    It also takes the binary image, the base address and the device class as
    parameters.
    """

    print("Erasing pages...")
    pbar = progressbar.ProgressBar(maxval=len(binary)).start()

    # First erase all pages
    for offset in range(0, len(binary), page_size):
        erase_command = commands.encode_erase_flash_page(base_address + offset,
                                                         device_class)
        res = utils.write_command_retry(fdesc, erase_command, destinations)

        failed_boards = [str(id) for id, success in res.items()
                         if not msgpack.unpackb(success)]

        if failed_boards:
            msg = ", ".join(failed_boards)
            msg = "Boards {} failed during page erase, aborting...".format(msg)
            logging.critical(msg)
            sys.exit(2)

        pbar.update(offset)

    pbar.finish()

    print("Writing pages...")
    pbar = progressbar.ProgressBar(maxval=len(binary)).start()

    # Then write all pages in chunks
    for offset, chunk in enumerate(page.slice_into_pages(binary, CHUNK_SIZE)):
        offset *= CHUNK_SIZE
        command = commands.encode_write_flash(chunk,
                                              base_address + offset,
                                              device_class)

        res = utils.write_command_retry(fdesc, command, destinations)
        failed_boards = [str(id) for id, success in res.items()
                         if not msgpack.unpackb(success)]

        if failed_boards:
            msg = ", ".join(failed_boards)
            msg = "Boards {} failed during page write, aborting...".format(msg)
            logging.critical(msg)
            sys.exit(2)

        pbar.update(offset)
    pbar.finish()

    # Finally update application CRC and size in config
    config = dict()
    config['application_size'] = len(binary)
    config['application_crc'] = crc32(binary)
    utils.config_update_and_save(fdesc, config, destinations)

def check_binary(fdesc, binary, base_address, destinations):
    """
    Check that the binary was correctly written to all destinations.

    Returns a list of all nodes which are passing the test.
    """
    valid_nodes = []

    expected_crc = crc32(binary)

    command = commands.encode_crc_region(base_address, len(binary))
    utils.write_command(fdesc, command, destinations)

    reader = utils.read_can_datagrams(fdesc)

    boards_checked = 0

    while boards_checked < len(destinations):
        dt = next(reader)

        if dt is None:
            continue

        answer, _, src = dt

        crc = msgpack.unpackb(answer)

        if crc == expected_crc:
            valid_nodes.append(src)

        boards_checked += 1

    return valid_nodes

def run_application(fdesc, destinations):
    """
    Asks the given node to run the application.
    """
    command = commands.encode_jump_to_main()
    utils.write_command(fdesc, command, destinations)

def verification_failed(failed_nodes):
    """
    Prints a message about the verification failing and exits
    """
    error_msg = "Verification failed for nodes {}" \
                .format(", ".join(str(x) for x in failed_nodes))
    print(error_msg)
    exit(1)

def check_online_boards(fdesc, boards):
    """
    Returns a set containing the online boards.
    """
    online_boards = set()

    utils.write_command(fdesc, commands.encode_ping(), boards)
    reader = utils.read_can_datagrams(fdesc)

    while True:
        dt = reader.read_datagram()

        if dt is None:  # Timeout
            break

        _, _, src = dt
        online_boards.add(src)

    return online_boards

def main():
    """
    Entry point of the application.
    """
    args = parse_commandline_args()
    with open(args.binary_file, 'rb') as input_file:
        binary = input_file.read()

    serial_port = utils.open_connection(args)

    online_boards = check_online_boards(serial_port, args.ids)

    if online_boards != set(args.ids):
        offline_boards = [str(i) for i in set(args.ids) - online_boards]
        print("Boards {} are offline, aborting..."
              .format(", ".join(offline_boards)))
        exit(2)

    print("Flashing firmware (size: {} bytes)".format(len(binary)))
    flash_binary(serial_port, binary, args.base_address, args.device_class,
                 args.ids)

    print("Verifying firmware...")
    valid_nodes_set = set(check_binary(serial_port, binary,
                                       args.base_address, args.ids))
    nodes_set = set(args.ids)

    if valid_nodes_set == nodes_set:
        print("OK")
    else:
        verification_failed(nodes_set - valid_nodes_set)

    if args.run:
        run_application(serial_port, args.ids)

if __name__ == "__main__":
    main()
<commit_msg>Update datagram reader in flashing tool<commit_after>#!/usr/bin/env python3
"""
Update firmware using CVRA bootloading protocol.
"""
import page
import logging
import commands
import msgpack
from zlib import crc32
from sys import exit

import utils
import progressbar
import sys

CHUNK_SIZE = 2048

def parse_commandline_args(args=None):
    """
    Parses the program commandline arguments.
    Args must be an array containing all arguments.
    """
    parser = utils.ConnectionArgumentParser(description=__doc__)
    parser.add_argument('-b', '--binary', dest='binary_file',
                        help='Path to the binary file to upload',
                        required=True,
                        metavar='FILE')

    parser.add_argument('-a', '--base-address', dest='base_address',
                        help='Base address of the firmware',
                        metavar='ADDRESS',
                        required=True,
                        # automatically convert value to hex
                        type=lambda s: int(s, 16))

    parser.add_argument('-c', '--device-class',
                        dest='device_class',
                        help='Device class to flash', required=True)
    parser.add_argument('-r', '--run',
                        help='Run application after flashing',
                        action='store_true')
    parser.add_argument("ids",
                        metavar='DEVICEID',
                        nargs='+', type=int,
                        help="Device IDs to flash")

    return parser.parse_args(args)

def flash_binary(fdesc, binary, base_address, device_class, destinations,
                 page_size=2048):
    """
    Writes a full binary to the flash using the given file descriptor.

    It also takes the binary image, the base address and the device class as
    parameters.
    """

    print("Erasing pages...")
    pbar = progressbar.ProgressBar(maxval=len(binary)).start()

    # First erase all pages
    for offset in range(0, len(binary), page_size):
        erase_command = commands.encode_erase_flash_page(base_address + offset,
                                                         device_class)
        res = utils.write_command_retry(fdesc, erase_command, destinations)

        failed_boards = [str(id) for id, success in res.items()
                         if not msgpack.unpackb(success)]

        if failed_boards:
            msg = ", ".join(failed_boards)
            msg = "Boards {} failed during page erase, aborting...".format(msg)
            logging.critical(msg)
            sys.exit(2)

        pbar.update(offset)

    pbar.finish()

    print("Writing pages...")
    pbar = progressbar.ProgressBar(maxval=len(binary)).start()

    # Then write all pages in chunks
    for offset, chunk in enumerate(page.slice_into_pages(binary, CHUNK_SIZE)):
        offset *= CHUNK_SIZE
        command = commands.encode_write_flash(chunk,
                                              base_address + offset,
                                              device_class)

        res = utils.write_command_retry(fdesc, command, destinations)
        failed_boards = [str(id) for id, success in res.items()
                         if not msgpack.unpackb(success)]

        if failed_boards:
            msg = ", ".join(failed_boards)
            msg = "Boards {} failed during page write, aborting...".format(msg)
            logging.critical(msg)
            sys.exit(2)

        pbar.update(offset)
    pbar.finish()

    # Finally update application CRC and size in config
    config = dict()
    config['application_size'] = len(binary)
    config['application_crc'] = crc32(binary)
    utils.config_update_and_save(fdesc, config, destinations)

def check_binary(fdesc, binary, base_address, destinations):
    """
    Check that the binary was correctly written to all destinations.

    Returns a list of all nodes which are passing the test.
    """
    valid_nodes = []

    expected_crc = crc32(binary)

    command = commands.encode_crc_region(base_address, len(binary))
    utils.write_command(fdesc, command, destinations)

    reader = utils.read_can_datagrams(fdesc)

    boards_checked = 0

    while boards_checked < len(destinations):
        dt = next(reader)

        if dt is None:
            continue

        answer, _, src = dt

        crc = msgpack.unpackb(answer)

        if crc == expected_crc:
            valid_nodes.append(src)

        boards_checked += 1

    return valid_nodes

def run_application(fdesc, destinations):
    """
    Asks the given node to run the application.
    """
    command = commands.encode_jump_to_main()
    utils.write_command(fdesc, command, destinations)

def verification_failed(failed_nodes):
    """
    Prints a message about the verification failing and exits
    """
    error_msg = "Verification failed for nodes {}" \
                .format(", ".join(str(x) for x in failed_nodes))
    print(error_msg)
    exit(1)

def check_online_boards(fdesc, boards):
    """
    Returns a set containing the online boards.
    """
    online_boards = set()

    utils.write_command(fdesc, commands.encode_ping(), boards)
    reader = utils.read_can_datagrams(fdesc)

    for dt in reader:
        if dt is None:
            break
        _, _, src = dt
        online_boards.add(src)

    return online_boards

def main():
    """
    Entry point of the application.
    """
    args = parse_commandline_args()
    with open(args.binary_file, 'rb') as input_file:
        binary = input_file.read()

    serial_port = utils.open_connection(args)

    online_boards = check_online_boards(serial_port, args.ids)

    if online_boards != set(args.ids):
        offline_boards = [str(i) for i in set(args.ids) - online_boards]
        print("Boards {} are offline, aborting..."
              .format(", ".join(offline_boards)))
        exit(2)

    print("Flashing firmware (size: {} bytes)".format(len(binary)))
    flash_binary(serial_port, binary, args.base_address, args.device_class,
                 args.ids)

    print("Verifying firmware...")
    valid_nodes_set = set(check_binary(serial_port, binary,
                                       args.base_address, args.ids))
    nodes_set = set(args.ids)

    if valid_nodes_set == nodes_set:
        print("OK")
    else:
        verification_failed(nodes_set - valid_nodes_set)

    if args.run:
        run_application(serial_port, args.ids)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Copyright (C) 2016 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# File: vyos-show-version
# Purpose:
#    Displays image version and system information.
#    Used by the "run show version" command.

import os
import sys
import argparse
import json

import vyos.version
import vyos.limericks

from vyos.util import cmd
from vyos.util import call
from vyos.util import run
from vyos.util import read_file
from vyos.util import DEVNULL

parser = argparse.ArgumentParser()
parser.add_argument("-a", "--all", action="store_true", help="Include individual package versions")
parser.add_argument("-f", "--funny", action="store_true", help="Add something funny to the output")
parser.add_argument("-j", "--json", action="store_true", help="Produce JSON output")

version_output_tmpl = """
Version:          VyOS {version}
Release Train:    {release_train}

Built by:         {built_by}
Built on:         {built_on}
Build UUID:       {build_uuid}
Build Commit ID:  {build_git}

Architecture:     {system_arch}
Boot via:         {boot_via}
System type:      {system_type}

Hardware vendor:  {hardware_vendor}
Hardware model:   {hardware_model}
Hardware S/N:     {hardware_serial}
Hardware UUID:    {hardware_uuid}

Copyright:        VyOS maintainers and contributors
"""

if __name__ == '__main__':
    args = parser.parse_args()

    version_data = vyos.version.get_full_version_data()

    if args.json:
        print(json.dumps(version_data))
        sys.exit(0)

    print(version_output_tmpl.format(**version_data).strip())

    if args.all:
        print("Package versions:")
        call("dpkg -l")

    if args.funny:
        print(vyos.limericks.get_random())
<commit_msg>op-mode: version: remove unused import statements<commit_after>#!/usr/bin/env python3
#
# Copyright (C) 2016-2020 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Purpose:
#    Displays image version and system information.
#    Used by the "run show version" command.

import sys
import argparse
import json

import vyos.version
import vyos.limericks

from vyos.util import call

parser = argparse.ArgumentParser()
parser.add_argument("-a", "--all", action="store_true", help="Include individual package versions")
parser.add_argument("-f", "--funny", action="store_true", help="Add something funny to the output")
parser.add_argument("-j", "--json", action="store_true", help="Produce JSON output")

version_output_tmpl = """
Version:          VyOS {version}
Release Train:    {release_train}

Built by:         {built_by}
Built on:         {built_on}
Build UUID:       {build_uuid}
Build Commit ID:  {build_git}

Architecture:     {system_arch}
Boot via:         {boot_via}
System type:      {system_type}

Hardware vendor:  {hardware_vendor}
Hardware model:   {hardware_model}
Hardware S/N:     {hardware_serial}
Hardware UUID:    {hardware_uuid}

Copyright:        VyOS maintainers and contributors
"""

if __name__ == '__main__':
    args = parser.parse_args()

    version_data = vyos.version.get_full_version_data()

    if args.json:
        print(json.dumps(version_data))
        sys.exit(0)

    print(version_output_tmpl.format(**version_data).strip())

    if args.all:
        print("Package versions:")
        call("dpkg -l")

    if args.funny:
        print(vyos.limericks.get_random())
<|endoftext|>


<commit_before># (c) Nelen & Schuurmans.  MIT licensed, see LICENSE.rst.
from __future__ import unicode_literals
from dateutil.relativedelta import relativedelta
from datetime import datetime
from pycassa.cassandra.ttypes import NotFoundException
import pandas as pd
import pycassa
import pytz

INTERNAL_TIMEZONE = pytz.UTC
COLNAME_FORMAT = '%Y-%m-%dT%H:%M:%SZ'
COLNAME_SEPERATOR = '_'

class BucketFormat:
    HOURLY = '%Y-%m-%dT%H'
    DAILY = '%Y-%m-%d'
    MONTHLY = '%Y-%m'
    YEARLY = '%Y'

def bucket_format(sensor_id):
    return BucketFormat.YEARLY

def bucket_delta(bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return relativedelta(hours=+1)
    if (bucketformat == BucketFormat.DAILY):
        return relativedelta(days=+1)
    if (bucketformat == BucketFormat.MONTHLY):
        return relativedelta(months=+1)
    if (bucketformat == BucketFormat.YEARLY):
        return relativedelta(years=+1)

def bucket_start(timestamp, bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return timestamp.replace(
            minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.DAILY):
        return timestamp.replace(
            hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.MONTHLY):
        return timestamp.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.YEARLY):
        return timestamp.replace(
            month=1, day=1, hour=0, minute=0, second=0, microsecond=0
        )

class CassandraDataStore(object):
    _column_families = {}
    _batches = {}

    def __init__(self, nodes, keyspace, queue_size):
        self.pool = pycassa.ConnectionPool(
            keyspace=keyspace, server_list=nodes, prefill=False
        )
        self.queue_size = queue_size

    def _get_column_family(self, column_family):
        if column_family not in self._column_families:
            self._column_families[column_family] = \
                pycassa.ColumnFamily(self.pool, column_family)
        return self._column_families[column_family]

    def _get_batch(self, column_family):
        if column_family not in self._batches:
            cf = self._get_column_family(column_family)
            self._batches[column_family] = cf.batch(queue_size=self.queue_size)
        return self._batches[column_family]

    def read(self, column_family, sensor_id, start, end, params=[]):
        assert start.tzinfo is not None, \
            "Start datetime must be timezone aware"
        assert end.tzinfo is not None, \
            "End datetime must be timezone aware"
        assert str(start.tzinfo.utcoffset(start))[1:] == ':00:00' or \
            str(start.tzinfo.utcoffset(start))[1:] == ':30:00', \
            "Start datetime has weird utc offset; use tz.localize"
        assert str(end.tzinfo.utcoffset(end))[1:] == ':00:00' or \
            str(end.tzinfo.utcoffset(end))[1:] == ':30:00', \
            "End datetime has weird utc offset; use tz.localize"

        # The bucket format defines how much data is on one Cassandra row.
        format = bucket_format(sensor_id)

        key_format = sensor_id + ":" + format
        stamp = bucket_start(start.astimezone(INTERNAL_TIMEZONE), format)
        delta = bucket_delta(format)

        col_start = start.astimezone(INTERNAL_TIMEZONE) \
            .strftime(COLNAME_FORMAT)
        col_end = end.astimezone(INTERNAL_TIMEZONE).strftime(COLNAME_FORMAT)

        datetimes = {}
        data = {}

        # From each bucket within in the specified range, get the columns
        # within the specified range.
        rowkeys = []
        while stamp < end:
            rowkeys.append(stamp.strftime(key_format))
            stamp += delta
        print "Get " + repr(rowkeys) + " with filter " + \
            col_start + " " + col_end
        try:
            result = self._get_column_family(column_family).multiget(
                rowkeys, column_start=col_start, column_finish=col_end
            )
            for rowkey in result:
                for col_name in result[rowkey]:
                    bits = col_name.split(COLNAME_SEPERATOR)
                    if (len(bits) > 1):
                        dt = datetime.strptime(bits[0], COLNAME_FORMAT)
                        key = col_name[len(bits[0]) + 1:]
                        if (key in params) or (len(params) == 0):
                            if not dt in datetimes.keys():
                                datetimes[dt] = {}
                            datetimes[dt][key] = result[rowkey][col_name]
                            if not key in data.keys():
                                data[key] = []
        except NotFoundException:
            pass
        except Exception as ex:
            print repr(ex)
        # Now, reform the data by series.
        for dt in sorted(datetimes):
            for key in data:
                if key in datetimes[dt]:
                    data[key].append(datetimes[dt][key])
                else:
                    data[key].append(None)
        # And create the Pandas DataFrame.
        return pd.DataFrame(data=data, index=sorted(datetimes))

    def write_frame(self, column_family, sensor_id, df):
        for timestamp, row in df.iterrows():
            self.write_row(column_family, sensor_id, timestamp, row.to_dict())

    def write_row(self, column_family, sensor_id, timestamp, row):
        ts_int = timestamp.astimezone(INTERNAL_TIMEZONE)
        key = ts_int.strftime(sensor_id + ':' + bucket_format(sensor_id))
        stamp = ts_int.strftime(COLNAME_FORMAT)
        self._get_batch(column_family).insert(key, dict(
            ("%s%s%s" % (stamp, COLNAME_SEPERATOR, k), str(v))
            for k, v in row.iteritems()
        ))

    def commit(self):
        for column_family in self._batches.keys():
            self._batches[column_family].send()
            del self._batches[column_family]
<commit_msg>Added truncate method to DataStore (needed for tests).<commit_after># (c) Nelen & Schuurmans.  MIT licensed, see LICENSE.rst.
from __future__ import unicode_literals
from dateutil.relativedelta import relativedelta
from datetime import datetime
from pycassa.cassandra.ttypes import NotFoundException
import pandas as pd
import pycassa
import pytz

INTERNAL_TIMEZONE = pytz.UTC
COLNAME_FORMAT = '%Y-%m-%dT%H:%M:%SZ'
COLNAME_SEPERATOR = '_'

class BucketFormat:
    HOURLY = '%Y-%m-%dT%H'
    DAILY = '%Y-%m-%d'
    MONTHLY = '%Y-%m'
    YEARLY = '%Y'

def bucket_format(sensor_id):
    return BucketFormat.YEARLY

def bucket_delta(bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return relativedelta(hours=+1)
    if (bucketformat == BucketFormat.DAILY):
        return relativedelta(days=+1)
    if (bucketformat == BucketFormat.MONTHLY):
        return relativedelta(months=+1)
    if (bucketformat == BucketFormat.YEARLY):
        return relativedelta(years=+1)

def bucket_start(timestamp, bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return timestamp.replace(
            minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.DAILY):
        return timestamp.replace(
            hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.MONTHLY):
        return timestamp.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.YEARLY):
        return timestamp.replace(
            month=1, day=1, hour=0, minute=0, second=0, microsecond=0
        )

class CassandraDataStore(object):
    _column_families = {}
    _batches = {}

    def __init__(self, nodes, keyspace, queue_size):
        self.pool = pycassa.ConnectionPool(
            keyspace=keyspace, server_list=nodes, prefill=False
        )
        self.queue_size = queue_size

    def _get_column_family(self, column_family):
        if column_family not in self._column_families:
            self._column_families[column_family] = \
                pycassa.ColumnFamily(self.pool, column_family)
        return self._column_families[column_family]

    def _get_batch(self, column_family):
        if column_family not in self._batches:
            cf = self._get_column_family(column_family)
            self._batches[column_family] = cf.batch(queue_size=self.queue_size)
        return self._batches[column_family]

    def read(self, column_family, sensor_id, start, end, params=[]):
        assert start.tzinfo is not None, \
            "Start datetime must be timezone aware"
        assert end.tzinfo is not None, \
            "End datetime must be timezone aware"
        assert str(start.tzinfo.utcoffset(start))[1:] == ':00:00' or \
            str(start.tzinfo.utcoffset(start))[1:] == ':30:00', \
            "Start datetime has weird utc offset; use tz.localize"
        assert str(end.tzinfo.utcoffset(end))[1:] == ':00:00' or \
            str(end.tzinfo.utcoffset(end))[1:] == ':30:00', \
            "End datetime has weird utc offset; use tz.localize"

        # The bucket format defines how much data is on one Cassandra row.
        format = bucket_format(sensor_id)

        key_format = sensor_id + ":" + format
        stamp = bucket_start(start.astimezone(INTERNAL_TIMEZONE), format)
        delta = bucket_delta(format)

        col_start = start.astimezone(INTERNAL_TIMEZONE) \
            .strftime(COLNAME_FORMAT)
        col_end = end.astimezone(INTERNAL_TIMEZONE).strftime(COLNAME_FORMAT)

        datetimes = {}
        data = {}

        # From each bucket within in the specified range, get the columns
        # within the specified range.
        rowkeys = []
        while stamp < end:
            rowkeys.append(stamp.strftime(key_format))
            stamp += delta
        print "Get " + repr(rowkeys) + " with filter " + \
            col_start + " " + col_end
        try:
            result = self._get_column_family(column_family).multiget(
                rowkeys, column_start=col_start, column_finish=col_end
            )
            for rowkey in result:
                for col_name in result[rowkey]:
                    bits = col_name.split(COLNAME_SEPERATOR)
                    if (len(bits) > 1):
                        dt = datetime.strptime(bits[0], COLNAME_FORMAT)
                        key = col_name[len(bits[0]) + 1:]
                        if (key in params) or (len(params) == 0):
                            if not dt in datetimes.keys():
                                datetimes[dt] = {}
                            datetimes[dt][key] = result[rowkey][col_name]
                            if not key in data.keys():
                                data[key] = []
        except NotFoundException:
            pass
        except Exception as ex:
            print repr(ex)
        # Now, reform the data by series.
        for dt in sorted(datetimes):
            for key in data:
                if key in datetimes[dt]:
                    data[key].append(datetimes[dt][key])
                else:
                    data[key].append(None)
        # And create the Pandas DataFrame.
        return pd.DataFrame(data=data, index=sorted(datetimes))

    def write_frame(self, column_family, sensor_id, df):
        for timestamp, row in df.iterrows():
            self.write_row(column_family, sensor_id, timestamp, row.to_dict())

    def write_row(self, column_family, sensor_id, timestamp, row):
        ts_int = timestamp.astimezone(INTERNAL_TIMEZONE)
        key = ts_int.strftime(sensor_id + ':' + bucket_format(sensor_id))
        stamp = ts_int.strftime(COLNAME_FORMAT)
        self._get_batch(column_family).insert(key, dict(
            ("%s%s%s" % (stamp, COLNAME_SEPERATOR, k), str(v))
            for k, v in row.iteritems()
        ))

    def truncate(self, column_family):
        self._get_column_family(column_family).truncate()

    def commit(self):
        for column_family in self._batches.keys():
            self._batches[column_family].send()
            del self._batches[column_family]
<|endoftext|>


<commit_before># (c) Nelen & Schuurmans.  MIT licensed, see LICENSE.rst.

from __future__ import unicode_literals

from datetime import datetime
from dateutil.relativedelta import relativedelta
import logging

from pycassa.cassandra.ttypes import NotFoundException
import numpy as np
import pandas as pd
import pycassa
import pytz

INTERNAL_TIMEZONE = pytz.UTC
COLNAME_FORMAT = '%Y-%m-%dT%H:%M:%SZ'
COLNAME_FORMAT_MS = '%Y-%m-%dT%H:%M:%S.%fZ'
COLNAME_SEPERATOR = '_'
MAX_COLUMNS = 2147483647

logger = logging.getLogger(__name__)

class BucketFormat:
    HOURLY = '%Y-%m-%dT%H'
    DAILY = '%Y-%m-%d'
    MONTHLY = '%Y-%m'
    YEARLY = '%Y'

def bucket_format(sensor_id):
    return BucketFormat.YEARLY

def bucket_delta(bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return relativedelta(hours=+1)
    if (bucketformat == BucketFormat.DAILY):
        return relativedelta(days=+1)
    if (bucketformat == BucketFormat.MONTHLY):
        return relativedelta(months=+1)
    if (bucketformat == BucketFormat.YEARLY):
        return relativedelta(years=+1)

def bucket_start(timestamp, bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return timestamp.replace(
            minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.DAILY):
        return timestamp.replace(
            hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.MONTHLY):
        return timestamp.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.YEARLY):
        return timestamp.replace(
            month=1, day=1, hour=0, minute=0, second=0, microsecond=0
        )

def strptime(dt):
    # datetime.strptime is slow. This low level alternative performs better.
    if len(dt) > 26:
        return datetime(int(dt[0:4]), int(dt[5:7]), int(dt[8:10]),
            int(dt[11:13]), int(dt[14:16]), int(dt[17:19]), int(dt[20:26]))
    return datetime(int(dt[0:4]), int(dt[5:7]), int(dt[8:10]),
        int(dt[11:13]), int(dt[14:16]), int(dt[17:19]))

class CassandraDataStore(object):
    _column_families = {}
    _batches = {}

    def __init__(self, nodes, keyspace, queue_size):
        self.pool = pycassa.ConnectionPool(
            keyspace=keyspace, server_list=nodes, prefill=False, max_overflow=0
        )
        self.queue_size = queue_size
        self.read_consistency_level = pycassa.ConsistencyLevel.ONE
        self.write_consistency_level = pycassa.ConsistencyLevel.QUORUM

    def _get_column_family(self, column_family):
        if column_family not in self._column_families:
            self._column_families[column_family] = \
                pycassa.ColumnFamily(self.pool, column_family,
                    write_consistency_level=self.write_consistency_level,
                    read_consistency_level=self.read_consistency_level)
        return self._column_families[column_family]

    def _get_batch(self, column_family):
        if column_family not in self._batches:
            cf = self._get_column_family(column_family)
            self._batches[column_family] = cf.batch(queue_size=self.queue_size)
        return self._batches[column_family]

    def read(self, column_family, sensor_id, start, end, params=[],
             convert_values_to=None, ignore_rejected=None):
        if start:
            assert start.tzinfo is not None, \
                "Start datetime must be timezone aware"
            assert str(start.tzinfo.utcoffset(start))[1:] == ':00:00' or \
                str(start.tzinfo.utcoffset(start))[1:] == ':30:00', \
                "Start datetime has weird utc offset; use tz.localize"
        if end:
            assert end.tzinfo is not None, \
                "End datetime must be timezone aware"
            assert str(end.tzinfo.utcoffset(end))[1:] == ':00:00' or \
                str(end.tzinfo.utcoffset(end))[1:] == ':30:00', \
                "End datetime has weird utc offset; use tz.localize"

        if start is None or end is None:
            return pd.DataFrame()

        # The bucket format defines how much data is on one Cassandra row.
        format = bucket_format(sensor_id)

        key_format = sensor_id + ":" + format
        stamp = bucket_start(start.astimezone(INTERNAL_TIMEZONE), format)
        delta = bucket_delta(format)

        # From each bucket within in the specified range, get the columns
        # within the specified range.
        rowkeys = []
        while stamp < end:
            rowkeys.append(stamp.strftime(key_format))
            stamp += delta

        # If no Cassandra rows are in requested date range, return nothing.
        if len(rowkeys) == 0:
            return pd.DataFrame()

        col_start = start.astimezone(INTERNAL_TIMEZONE) \
            .strftime(COLNAME_FORMAT_MS)
        col_end = end.astimezone(INTERNAL_TIMEZONE).strftime(COLNAME_FORMAT_MS)

        data = {}
        keys = []

        try:
            timer_start = datetime.now()
            result = self._get_column_family(column_family).multiget(
                rowkeys,
                column_start=col_start,
                column_finish=col_end,
                column_count=MAX_COLUMNS
            )
            for rowkey in result:
                for col_name in result[rowkey]:
                    bits = col_name.split(COLNAME_SEPERATOR)
                    if (len(bits) > 1):
                        dt = strptime(bits[0])
                        key = col_name[len(bits[0]) + 1:]
                        if not params or key in params:
                            if key not in keys:
                                keys.append(key)
                            if not dt in data:
                                data[dt] = {}
                            data[dt][key] = result[rowkey][col_name]
            logger.debug("multiget in %s", datetime.now() - timer_start)
        except NotFoundException:
            pass

        # Flatten the dataset by key.
        # Missing values are converted to None.
        timer_start = datetime.now()
        datetimes = sorted(data.keys())
        notrejected = datetimes
        data_flat = {key: [] for key in keys}
        for dt in datetimes:
            row = data[dt]
            flag = row.get("flag")
            if ignore_rejected and flag == "6":
                notrejected.remove(dt)
                continue
            else:
                for key in keys:
                    value = row.get(key)
                    data_flat[key].append(value)
        # index should align with keys values.
        datetimes = notrejected
        logger.debug("flattened in %s", datetime.now() - timer_start)

        # Convert values to an appropriate Numpy array type, if requested.
        # Unknown types are kept in their current (Cassandra) form.
        timer_start = datetime.now()
        if convert_values_to is not None and 'value' in data_flat:
            dtype_map = {
                'float': np.float32,
                'integer': np.int32
            }
            dtype = dtype_map.get(convert_values_to)
            if dtype is not None:
                # There's a bug in numpy.genfromtxt causing a list of length 1
                # to produce a 0-dim array. Appending and removing a dummy
                # value resolves this.
                length = len(data_flat['value'])
                data_flat['value'].append('dummy')
                data_flat['value'] = \
                    np.genfromtxt(data_flat['value'], dtype=dtype)[:length]
        logger.debug("converted in %s", datetime.now() - timer_start)

        # And create the Pandas DataFrame.
        timer_start = datetime.now()
        result = pd.DataFrame(data=data_flat, index=datetimes)
        logger.debug("pandafied in %s", datetime.now() - timer_start)
        if len(datetimes) > 0:
            result.tz_localize(INTERNAL_TIMEZONE, copy=False)
        return result

    def write_row(self, column_family, sensor_id, timestamp, row):
        ts_int = timestamp.astimezone(INTERNAL_TIMEZONE)
        key = ts_int.strftime(sensor_id + ':' + bucket_format(sensor_id))
        stamp = ts_int.strftime(COLNAME_FORMAT_MS)
        self._get_batch(column_family).insert(key, dict(
            ("%s%s%s" % (stamp, COLNAME_SEPERATOR, k), str(v))
            for k, v in row.iteritems()
        ))

    def truncate(self, column_family):
        self._get_column_family(column_family).truncate()

    def commit(self, column_family):
        if column_family in self._batches.keys():
            self._batches[column_family].send()
            del self._batches[column_family]
<commit_msg>noobieness removed<commit_after># (c) Nelen & Schuurmans.  MIT licensed, see LICENSE.rst.

from __future__ import unicode_literals

from datetime import datetime
from dateutil.relativedelta import relativedelta
import logging

from pycassa.cassandra.ttypes import NotFoundException
import numpy as np
import pandas as pd
import pycassa
import pytz

INTERNAL_TIMEZONE = pytz.UTC
COLNAME_FORMAT = '%Y-%m-%dT%H:%M:%SZ'
COLNAME_FORMAT_MS = '%Y-%m-%dT%H:%M:%S.%fZ'
COLNAME_SEPERATOR = '_'
MAX_COLUMNS = 2147483647

logger = logging.getLogger(__name__)

class BucketFormat:
    HOURLY = '%Y-%m-%dT%H'
    DAILY = '%Y-%m-%d'
    MONTHLY = '%Y-%m'
    YEARLY = '%Y'

def bucket_format(sensor_id):
    return BucketFormat.YEARLY

def bucket_delta(bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return relativedelta(hours=+1)
    if (bucketformat == BucketFormat.DAILY):
        return relativedelta(days=+1)
    if (bucketformat == BucketFormat.MONTHLY):
        return relativedelta(months=+1)
    if (bucketformat == BucketFormat.YEARLY):
        return relativedelta(years=+1)

def bucket_start(timestamp, bucketformat):
    if (bucketformat == BucketFormat.HOURLY):
        return timestamp.replace(
            minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.DAILY):
        return timestamp.replace(
            hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.MONTHLY):
        return timestamp.replace(
            day=1, hour=0, minute=0, second=0, microsecond=0
        )
    if (bucketformat == BucketFormat.YEARLY):
        return timestamp.replace(
            month=1, day=1, hour=0, minute=0, second=0, microsecond=0
        )

def strptime(dt):
    # datetime.strptime is slow. This low level alternative performs better.
    if len(dt) > 26:
        return datetime(int(dt[0:4]), int(dt[5:7]), int(dt[8:10]),
            int(dt[11:13]), int(dt[14:16]), int(dt[17:19]), int(dt[20:26]))
    return datetime(int(dt[0:4]), int(dt[5:7]), int(dt[8:10]),
        int(dt[11:13]), int(dt[14:16]), int(dt[17:19]))

class CassandraDataStore(object):
    _column_families = {}
    _batches = {}

    def __init__(self, nodes, keyspace, queue_size):
        self.pool = pycassa.ConnectionPool(
            keyspace=keyspace, server_list=nodes, prefill=False, max_overflow=0
        )
        self.queue_size = queue_size
        self.read_consistency_level = pycassa.ConsistencyLevel.ONE
        self.write_consistency_level = pycassa.ConsistencyLevel.QUORUM

    def _get_column_family(self, column_family):
        if column_family not in self._column_families:
            self._column_families[column_family] = \
                pycassa.ColumnFamily(self.pool, column_family,
                    write_consistency_level=self.write_consistency_level,
                    read_consistency_level=self.read_consistency_level)
        return self._column_families[column_family]

    def _get_batch(self, column_family):
        if column_family not in self._batches:
            cf = self._get_column_family(column_family)
            self._batches[column_family] = cf.batch(queue_size=self.queue_size)
        return self._batches[column_family]

    def read(self, column_family, sensor_id, start, end, params=[],
             convert_values_to=None, ignore_rejected=None):
        if start:
            assert start.tzinfo is not None, \
                "Start datetime must be timezone aware"
            assert str(start.tzinfo.utcoffset(start))[1:] == ':00:00' or \
                str(start.tzinfo.utcoffset(start))[1:] == ':30:00', \
                "Start datetime has weird utc offset; use tz.localize"
        if end:
            assert end.tzinfo is not None, \
                "End datetime must be timezone aware"
            assert str(end.tzinfo.utcoffset(end))[1:] == ':00:00' or \
                str(end.tzinfo.utcoffset(end))[1:] == ':30:00', \
                "End datetime has weird utc offset; use tz.localize"

        if start is None or end is None:
            return pd.DataFrame()

        # The bucket format defines how much data is on one Cassandra row.
        format = bucket_format(sensor_id)

        key_format = sensor_id + ":" + format
        stamp = bucket_start(start.astimezone(INTERNAL_TIMEZONE), format)
        delta = bucket_delta(format)

        # From each bucket within in the specified range, get the columns
        # within the specified range.
        rowkeys = []
        while stamp < end:
            rowkeys.append(stamp.strftime(key_format))
            stamp += delta

        # If no Cassandra rows are in requested date range, return nothing.
        if len(rowkeys) == 0:
            return pd.DataFrame()

        col_start = start.astimezone(INTERNAL_TIMEZONE) \
            .strftime(COLNAME_FORMAT_MS)
        col_end = end.astimezone(INTERNAL_TIMEZONE).strftime(COLNAME_FORMAT_MS)

        data = {}
        keys = []

        try:
            timer_start = datetime.now()
            result = self._get_column_family(column_family).multiget(
                rowkeys,
                column_start=col_start,
                column_finish=col_end,
                column_count=MAX_COLUMNS
            )
            for rowkey in result:
                for col_name in result[rowkey]:
                    bits = col_name.split(COLNAME_SEPERATOR)
                    if (len(bits) > 1):
                        dt = strptime(bits[0])
                        key = col_name[len(bits[0]) + 1:]
                        if not params or key in params:
                            if key not in keys:
                                keys.append(key)
                            if not dt in data:
                                data[dt] = {}
                            data[dt][key] = result[rowkey][col_name]
            logger.debug("multiget in %s", datetime.now() - timer_start)
        except NotFoundException:
            pass

        # Flatten the dataset by key.
        # Missing values are converted to None.
        timer_start = datetime.now()
        datetimes = sorted(data.keys())
        notrejected = list(datetimes)
        data_flat = {key: [] for key in keys}
        for dt in datetimes:
            row = data[dt]
            flag = row.get("flag")
            if ignore_rejected and flag == "6":
                notrejected.remove(dt)
                continue
            else:
                for key in keys:
                    value = row.get(key)
                    data_flat[key].append(value)
        # index should align with keys values.
        datetimes = notrejected
        logger.debug("flattened in %s", datetime.now() - timer_start)

        # Convert values to an appropriate Numpy array type, if requested.
        # Unknown types are kept in their current (Cassandra) form.
        timer_start = datetime.now()
        if datetimes and convert_values_to is not None and 'value' in data_flat:
            dtype_map = {
                'float': np.float32,
                'integer': np.int32
            }
            dtype = dtype_map.get(convert_values_to)
            if dtype is not None:
                # There's a bug in numpy.genfromtxt causing a list of length 1
                # to produce a 0-dim array. Appending and removing a dummy
                # value resolves this.
                length = len(data_flat['value'])
                data_flat['value'].append('dummy')
                data_flat['value'] = \
                    np.genfromtxt(data_flat['value'], dtype=dtype)[:length]
        logger.debug("converted in %s", datetime.now() - timer_start)

        # And create the Pandas DataFrame.
        timer_start = datetime.now()
        result = pd.DataFrame(data=data_flat, index=datetimes)
        logger.debug("pandafied in %s", datetime.now() - timer_start)
        if len(datetimes) > 0:
            result.tz_localize(INTERNAL_TIMEZONE, copy=False)
        return result

    def write_row(self, column_family, sensor_id, timestamp, row):
        ts_int = timestamp.astimezone(INTERNAL_TIMEZONE)
        key = ts_int.strftime(sensor_id + ':' + bucket_format(sensor_id))
        stamp = ts_int.strftime(COLNAME_FORMAT_MS)
        self._get_batch(column_family).insert(key, dict(
            ("%s%s%s" % (stamp, COLNAME_SEPERATOR, k), str(v))
            for k, v in row.iteritems()
        ))

    def truncate(self, column_family):
        self._get_column_family(column_family).truncate()

    def commit(self, column_family):
        if column_family in self._batches.keys():
            self._batches[column_family].send()
            del self._batches[column_family]
<|endoftext|>


<commit_before>import sys
import functools
import multiprocessing

import ce.logger as logger
from ce.logger import levels
from ce.common import cached
from ce.expr.common import is_expr
from ce.expr import Expr

RECURSION_LIMIT = sys.getrecursionlimit()

def item_to_list(f):
    def wrapper(t):
        v = f(t)
        return [v] if not v is None else []
    return functools.wraps(f)(wrapper)

def none_to_list(f):
    def wrapper(t):
        v = f(t)
        return v if not v is None else []
    return functools.wraps(f)(wrapper)

class ValidationError(Exception):
    """Failed to find equivalence."""

class TreeTransformer(object):

    transform_methods = None
    reduction_methods = None

    def __init__(self, tree_or_trees,
                 validate=False, depth=None,
                 step_plugin=None, reduce_plugin=None,
                 multiprocessing=True):
        try:
            self._t = [Expr(tree_or_trees)]
        except TypeError:
            self._t = tree_or_trees
        self._v = validate
        self._m = multiprocessing
        self._d = depth or RECURSION_LIMIT
        self._n = {}
        self._sp = step_plugin
        self._rp = reduce_plugin
        self.transform_methods = list(self.__class__.transform_methods or [])
        self.reduction_methods = list(self.__class__.reduction_methods or [])
        super().__init__()

    def _harvest(self, trees):
        if self._d >= RECURSION_LIMIT:
            return trees
        logger.debug('Cropping trees')
        cropped = []
        for t in trees:
            try:
                t, e = t.crop(self._d)
            except AttributeError:
                t, e = t, {}
            cropped.append(t)
            self._n.update(e)
        return cropped

    def _seed(self, trees):
        if not self._n:
            return trees
        logger.debug('Stitching back trees')
        seeded = set()
        for t in trees:
            try:
                t = t.stitch(self._n)
            except AttributeError:
                pass
            seeded.add(t)
        return seeded

    def _plugin(self, trees, plugin):
        if not plugin:
            return trees
        trees = self._seed(trees)
        trees = plugin(trees)
        trees = set(self._harvest(trees))
        return trees

    def _closure_r(self, trees, reduced=False):
        v = self._validate if self._v else None
        done_trees = set()
        todo_trees = set(trees)
        i = 0
        try:
            while todo_trees:
                # print set size
                i += 1
                logger.persistent(
                    'Iteration' if not reduced else 'Reduction', i,
                    l=levels.debug)
                logger.persistent('Trees', len(done_trees), l=levels.debug)
                logger.persistent('Todo', len(todo_trees), l=levels.debug)
                if not reduced:
                    f = self.transform_methods
                    _, step_trees = \
                        _step(todo_trees, f, v, not reduced, self._d, self._m)
                    step_trees -= done_trees
                    step_trees = self._closure_r(step_trees, True)
                    step_trees = self._plugin(step_trees, self._sp)
                    done_trees |= todo_trees
                    todo_trees = step_trees - done_trees
                else:
                    f = self.reduction_methods
                    nore_trees, step_trees = \
                        _step(todo_trees, f, v, not reduced, None, self._m)
                    step_trees = self._plugin(step_trees, self._rp)
                    done_trees |= nore_trees
                    todo_trees = step_trees - nore_trees
        except KeyboardInterrupt:
            if reduced:
                raise
        logger.unpersistent('Iteration', 'Reduction', 'Trees', 'Todo')
        return done_trees

    def closure(self):
        """Perform transforms until transitive closure is reached.

        Returns:
            A set of trees after transform.
        """
        s = self._harvest(self._t)
        logger.debug('Computing equivalent expressions.')
        s = self._closure_r(s)
        s = self._seed(s)
        return s

    def validate(self, t, tn):
        """Perform validation of tree.

        Args:
            t: An original tree.
            tn: A transformed tree.

        Raises:
            ValidationError: If failed to find equivalences between t and tn.
        """
        pass

    def _validate(self, t, tn):
        if t == tn:
            return
        if self._v:
            self.validate(t, tn)

def _walk(t_fs_v_c_d):
    t, fs, v, c, d = t_fs_v_c_d
    s = set()
    for f in fs:
        s |= _walk_r(t, f, v, d if c else RECURSION_LIMIT)
    if c:
        s.add(t)
    return True if not c and not s else False, s

@cached
def _walk_r(t, f, v, d):
    s = set()
    if not is_expr(t):
        return s
    for e in f(t):
        s.add(e)
    for a, b in (t.args, t.args[::-1]):
        for e in _walk_r(a, f, v, d):
            s.add(Expr(t.op, e, b))
    if not v:
        return s
    try:
        for tn in s:
            v(t, tn)
    except ValidationError:
        logger.error('Violating transformation:', f.__name__)
        raise
    return s

def par_union(sl):
    return functools.reduce(lambda s, t: s | t, sl)

_pool = None

def pool():
    global _pool
    if _pool is None:
        _pool = multiprocessing.Pool()
    return _pool

def _iunion(sl, no_processes):
    def chunks(l, n):
        for i in range(0, len(l), n):
            yield l[i:i+n]
    chunksize = int(len(sl) / no_processes) + 2
    while len(sl) > 1:
        sys.stdout.write('\rUnion: %d.' % len(sl))
        sys.stdout.flush()
        sl = list(pool().imap_unordered(par_union, chunks(sl, chunksize)))
    return sl.pop()

def _step(s, fs, v=None, c=False, d=None, m=True):
    """Find the set of trees related by the function f.

    Args:
        s: A set of trees.
        fs: A set of functions which transforms the trees. Each function has
            one argument, the tree, and returns a set of trees after transform.
        v: A function which validates the transform.
        c: If set, it will include everything in s. Otherwise only derived
            trees.
        d: Depth of node traversal in trees.
        m: Whether multiprocessing is used.

    Returns:
        A set of trees related by f.
    """
    if m:
        cpu_count = multiprocessing.cpu_count()
        chunksize = int(len(s) / cpu_count) + 1
        map = pool().imap
    else:
        cpu_count = chunksize = 1
        map = lambda f, l, _: [f(a) for a in l]
    union = lambda s, _: functools.reduce(lambda x, y: x | y, s)
    r = [(t, fs, v, c, d) for i, t in enumerate(s)]
    b, r = list(zip(*map(_walk, r, chunksize)))
    s = {t for i, t in enumerate(s) if b[i]}
    r = union(r, cpu_count)
    return s, r
<commit_msg>Disallow keyboarditerrupt for transformer<commit_after>import sys
import functools
import multiprocessing

import ce.logger as logger
from ce.logger import levels
from ce.common import cached
from ce.expr.common import is_expr
from ce.expr import Expr

RECURSION_LIMIT = sys.getrecursionlimit()

def item_to_list(f):
    def wrapper(t):
        v = f(t)
        return [v] if not v is None else []
    return functools.wraps(f)(wrapper)

def none_to_list(f):
    def wrapper(t):
        v = f(t)
        return v if not v is None else []
    return functools.wraps(f)(wrapper)

class ValidationError(Exception):
    """Failed to find equivalence."""

class TreeTransformer(object):

    transform_methods = None
    reduction_methods = None

    def __init__(self, tree_or_trees,
                 validate=False, depth=None,
                 step_plugin=None, reduce_plugin=None,
                 multiprocessing=True):
        try:
            self._t = [Expr(tree_or_trees)]
        except TypeError:
            self._t = tree_or_trees
        self._v = validate
        self._m = multiprocessing
        self._d = depth or RECURSION_LIMIT
        self._n = {}
        self._sp = step_plugin
        self._rp = reduce_plugin
        self.transform_methods = list(self.__class__.transform_methods or [])
        self.reduction_methods = list(self.__class__.reduction_methods or [])
        super().__init__()

    def _harvest(self, trees):
        if self._d >= RECURSION_LIMIT:
            return trees
        logger.debug('Cropping trees')
        cropped = []
        for t in trees:
            try:
                t, e = t.crop(self._d)
            except AttributeError:
                t, e = t, {}
            cropped.append(t)
            self._n.update(e)
        return cropped

    def _seed(self, trees):
        if not self._n:
            return trees
        logger.debug('Stitching back trees')
        seeded = set()
        for t in trees:
            try:
                t = t.stitch(self._n)
            except AttributeError:
                pass
            seeded.add(t)
        return seeded

    def _plugin(self, trees, plugin):
        if not plugin:
            return trees
        trees = self._seed(trees)
        trees = plugin(trees)
        trees = set(self._harvest(trees))
        return trees

    def _closure_r(self, trees, reduced=False):
        v = self._validate if self._v else None
        done_trees = set()
        todo_trees = set(trees)
        i = 0
        while todo_trees:
            # print set size
            i += 1
            logger.persistent(
                'Iteration' if not reduced else 'Reduction', i,
                l=levels.debug)
            logger.persistent('Trees', len(done_trees), l=levels.debug)
            logger.persistent('Todo', len(todo_trees), l=levels.debug)
            if not reduced:
                f = self.transform_methods
                _, step_trees = \
                    _step(todo_trees, f, v, not reduced, self._d, self._m)
                step_trees -= done_trees
                step_trees = self._closure_r(step_trees, True)
                step_trees = self._plugin(step_trees, self._sp)
                done_trees |= todo_trees
                todo_trees = step_trees - done_trees
            else:
                f = self.reduction_methods
                nore_trees, step_trees = \
                    _step(todo_trees, f, v, not reduced, None, self._m)
                step_trees = self._plugin(step_trees, self._rp)
                done_trees |= nore_trees
                todo_trees = step_trees - nore_trees
        logger.unpersistent('Iteration', 'Reduction', 'Trees', 'Todo')
        return done_trees

    def closure(self):
        """Perform transforms until transitive closure is reached.

        Returns:
            A set of trees after transform.
        """
        s = self._harvest(self._t)
        logger.debug('Computing equivalent expressions.')
        s = self._closure_r(s)
        s = self._seed(s)
        return s

    def validate(self, t, tn):
        """Perform validation of tree.

        Args:
            t: An original tree.
            tn: A transformed tree.

        Raises:
            ValidationError: If failed to find equivalences between t and tn.
        """
        pass

    def _validate(self, t, tn):
        if t == tn:
            return
        if self._v:
            self.validate(t, tn)

def _walk(t_fs_v_c_d):
    t, fs, v, c, d = t_fs_v_c_d
    s = set()
    for f in fs:
        s |= _walk_r(t, f, v, d if c else RECURSION_LIMIT)
    if c:
        s.add(t)
    return True if not c and not s else False, s

@cached
def _walk_r(t, f, v, d):
    s = set()
    if not is_expr(t):
        return s
    for e in f(t):
        s.add(e)
    for a, b in (t.args, t.args[::-1]):
        for e in _walk_r(a, f, v, d):
            s.add(Expr(t.op, e, b))
    if not v:
        return s
    try:
        for tn in s:
            v(t, tn)
    except ValidationError:
        logger.error('Violating transformation:', f.__name__)
        raise
    return s

def par_union(sl):
    return functools.reduce(lambda s, t: s | t, sl)

_pool = None

def pool():
    global _pool
    if _pool is None:
        _pool = multiprocessing.Pool()
    return _pool

def _iunion(sl, no_processes):
    def chunks(l, n):
        for i in range(0, len(l), n):
            yield l[i:i+n]
    chunksize = int(len(sl) / no_processes) + 2
    while len(sl) > 1:
        sys.stdout.write('\rUnion: %d.' % len(sl))
        sys.stdout.flush()
        sl = list(pool().imap_unordered(par_union, chunks(sl, chunksize)))
    return sl.pop()

def _step(s, fs, v=None, c=False, d=None, m=True):
    """Find the set of trees related by the function f.

    Args:
        s: A set of trees.
        fs: A set of functions which transforms the trees. Each function has
            one argument, the tree, and returns a set of trees after transform.
        v: A function which validates the transform.
        c: If set, it will include everything in s. Otherwise only derived
            trees.
        d: Depth of node traversal in trees.
        m: Whether multiprocessing is used.

    Returns:
        A set of trees related by f.
    """
    if m:
        cpu_count = multiprocessing.cpu_count()
        chunksize = int(len(s) / cpu_count) + 1
        map = pool().imap
    else:
        cpu_count = chunksize = 1
        map = lambda f, l, _: [f(a) for a in l]
    union = lambda s, _: functools.reduce(lambda x, y: x | y, s)
    r = [(t, fs, v, c, d) for i, t in enumerate(s)]
    b, r = list(zip(*map(_walk, r, chunksize)))
    s = {t for i, t in enumerate(s) if b[i]}
    r = union(r, cpu_count)
    return s, r
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    celery.utils.compat
    ~~~~~~~~~~~~~~~~~~~

    Backward compatible implementations of features
    only available in newer Python versions.

    :copyright: (c) 2009 - 2012 by Ask Solem.
    :license: BSD, see LICENSE for more details.

"""
from __future__ import absolute_import

############## py3k #########################################################
import sys
is_py3k = sys.version_info[0] == 3

try:
    reload = reload                     # noqa
except NameError:
    from imp import reload              # noqa

try:
    from UserList import UserList       # noqa
except ImportError:
    from collections import UserList    # noqa

try:
    from UserDict import UserDict       # noqa
except ImportError:
    from collections import UserDict    # noqa

if is_py3k:
    from io import StringIO, BytesIO
    from .encoding import bytes_to_str

    class WhateverIO(StringIO):

        def write(self, data):
            StringIO.write(self, bytes_to_str(data))
else:
    try:
        from cStringIO import StringIO  # noqa
    except ImportError:
        from StringIO import StringIO   # noqa
    BytesIO = WhateverIO = StringIO     # noqa

# im_func is no longer available in Py3.
# instead the unbound method itself can be used.
if is_py3k:
    def fun_of_method(method):
        return method
else:
    def fun_of_method(method):  # noqa
        return method.im_func

############## collections.OrderedDict ######################################
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict  # noqa

############## logging.LoggerAdapter ########################################
import logging
try:
    import multiprocessing
except ImportError:
    multiprocessing = None  # noqa
import sys

def _checkLevel(level):
    if isinstance(level, int):
        rv = level
    elif str(level) == level:
        if level not in logging._levelNames:
            raise ValueError("Unknown level: %r" % level)
        rv = logging._levelNames[level]
    else:
        raise TypeError("Level not an integer or a valid string: %r" % level)
    return rv

class _CompatLoggerAdapter(object):

    def __init__(self, logger, extra):
        self.logger = logger
        self.extra = extra

    def setLevel(self, level):
        self.logger.level = _checkLevel(level)

    def process(self, msg, kwargs):
        kwargs["extra"] = self.extra
        return msg, kwargs

    def debug(self, msg, *args, **kwargs):
        self.log(logging.DEBUG, msg, *args, **kwargs)

    def info(self, msg, *args, **kwargs):
        self.log(logging.INFO, msg, *args, **kwargs)

    def warning(self, msg, *args, **kwargs):
        self.log(logging.WARNING, msg, *args, **kwargs)
    warn = warning

    def error(self, msg, *args, **kwargs):
        self.log(logging.ERROR, msg, *args, **kwargs)

    def exception(self, msg, *args, **kwargs):
        kwargs.setdefault("exc_info", 1)
        self.error(msg, *args, **kwargs)

    def critical(self, msg, *args, **kwargs):
        self.log(logging.CRITICAL, msg, *args, **kwargs)
    fatal = critical

    def log(self, level, msg, *args, **kwargs):
        if self.logger.isEnabledFor(level):
            msg, kwargs = self.process(msg, kwargs)
            self._log(level, msg, args, **kwargs)

    def makeRecord(self, name, level, fn, lno, msg, args, exc_info,
            func=None, extra=None):
        rv = logging.LogRecord(name, level, fn, lno, msg, args, exc_info, func)
        if extra is not None:
            for key, value in extra.items():
                if key in ("message", "asctime") or key in rv.__dict__:
                    raise KeyError(
                            "Attempt to override %r in LogRecord" % key)
                rv.__dict__[key] = value
        if multiprocessing is not None:
            rv.processName = multiprocessing.current_process()._name
        else:
            rv.processName = ""
        return rv

    def _log(self, level, msg, args, exc_info=None, extra=None):
        defcaller = "(unknown file)", 0, "(unknown function)"
        if logging._srcfile:
            # IronPython doesn't track Python frames, so findCaller
            # throws an exception on some versions of IronPython.
            # We trap it here so that IronPython can use logging.
            try:
                fn, lno, func = self.logger.findCaller()
            except ValueError:
                fn, lno, func = defcaller
        else:
            fn, lno, func = defcaller
        if exc_info:
            if not isinstance(exc_info, tuple):
                exc_info = sys.exc_info()
        record = self.makeRecord(self.logger.name, level, fn, lno, msg,
                                 args, exc_info, func, extra)
        self.logger.handle(record)

    def isEnabledFor(self, level):
        return self.logger.isEnabledFor(level)

    def addHandler(self, hdlr):
        self.logger.addHandler(hdlr)

    def removeHandler(self, hdlr):
        self.logger.removeHandler(hdlr)

    @property
    def level(self):
        return self.logger.level

try:
    from logging import LoggerAdapter
except ImportError:
    LoggerAdapter = _CompatLoggerAdapter  # noqa

############## itertools.zip_longest #######################################

try:
    from itertools import izip_longest as zip_longest
except ImportError:
    import itertools

    def zip_longest(*args, **kwds):  # noqa
        fillvalue = kwds.get("fillvalue")

        def sentinel(counter=([fillvalue] * (len(args) - 1)).pop):
            yield counter()     # yields the fillvalue, or raises IndexError

        fillers = itertools.repeat(fillvalue)
        iters = [itertools.chain(it, sentinel(), fillers)
                    for it in args]
        try:
            for tup in itertools.izip(*iters):
                yield tup
        except IndexError:
            pass

############## itertools.chain.from_iterable ################################
from itertools import chain

def _compat_chain_from_iterable(iterables):
    for it in iterables:
        for element in it:
            yield element

try:
    chain_from_iterable = getattr(chain, "from_iterable")
except AttributeError:
    chain_from_iterable = _compat_chain_from_iterable

############## logging.handlers.WatchedFileHandler ##########################
import os
from stat import ST_DEV, ST_INO
import platform as _platform

if _platform.system() == "Windows":
    #since windows doesn't go with WatchedFileHandler use FileHandler instead
    WatchedFileHandler = logging.FileHandler
else:
    try:
        from logging.handlers import WatchedFileHandler
    except ImportError:
        class WatchedFileHandler(logging.FileHandler):  # noqa
            """
            A handler for logging to a file, which watches the file
            to see if it has changed while in use. This can happen because of
            usage of programs such as newsyslog and logrotate which perform
            log file rotation. This handler, intended for use under Unix,
            watches the file to see if it has changed since the last emit.
            (A file has changed if its device or inode have changed.)
            If it has changed, the old file stream is closed, and the file
            opened to get a new stream.

            This handler is not appropriate for use under Windows, because
            under Windows open files cannot be moved or renamed - logging
            opens the files with exclusive locks - and so there is no need
            for such a handler. Furthermore, ST_INO is not supported under
            Windows; stat always returns zero for this value.

            This handler is based on a suggestion and patch by Chad J.
            Schroeder.
            """
            def __init__(self, *args, **kwargs):
                logging.FileHandler.__init__(self, *args, **kwargs)

                if not os.path.exists(self.baseFilename):
                    self.dev, self.ino = -1, -1
                else:
                    stat = os.stat(self.baseFilename)
                    self.dev, self.ino = stat[ST_DEV], stat[ST_INO]

            def emit(self, record):
                """
                Emit a record.

                First check if the underlying file has changed, and if it
                has, close the old stream and reopen the file to get the
                current stream.
                """
                if not os.path.exists(self.baseFilename):
                    stat = None
                    changed = 1
                else:
                    stat = os.stat(self.baseFilename)
                    changed = ((stat[ST_DEV] != self.dev) or
                               (stat[ST_INO] != self.ino))
                if changed and self.stream is not None:
                    self.stream.flush()
                    self.stream.close()
                    self.stream = self._open()
                    if stat is None:
                        stat = os.stat(self.baseFilename)
                    self.dev, self.ino = stat[ST_DEV], stat[ST_INO]
                logging.FileHandler.emit(self, record)
<commit_msg>fun_of_method moved to __compat__<commit_after># -*- coding: utf-8 -*-
"""
    celery.utils.compat
    ~~~~~~~~~~~~~~~~~~~

    Backward compatible implementations of features
    only available in newer Python versions.

    :copyright: (c) 2009 - 2012 by Ask Solem.
    :license: BSD, see LICENSE for more details.

"""
from __future__ import absolute_import

############## py3k #########################################################
import sys
is_py3k = sys.version_info[0] == 3

try:
    reload = reload                     # noqa
except NameError:
    from imp import reload              # noqa

try:
    from UserList import UserList       # noqa
except ImportError:
    from collections import UserList    # noqa

try:
    from UserDict import UserDict       # noqa
except ImportError:
    from collections import UserDict    # noqa

if is_py3k:
    from io import StringIO, BytesIO
    from .encoding import bytes_to_str

    class WhateverIO(StringIO):

        def write(self, data):
            StringIO.write(self, bytes_to_str(data))
else:
    try:
        from cStringIO import StringIO  # noqa
    except ImportError:
        from StringIO import StringIO   # noqa
    BytesIO = WhateverIO = StringIO     # noqa

############## collections.OrderedDict ######################################
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict  # noqa

############## logging.LoggerAdapter ########################################
import logging
try:
    import multiprocessing
except ImportError:
    multiprocessing = None  # noqa
import sys

def _checkLevel(level):
    if isinstance(level, int):
        rv = level
    elif str(level) == level:
        if level not in logging._levelNames:
            raise ValueError("Unknown level: %r" % level)
        rv = logging._levelNames[level]
    else:
        raise TypeError("Level not an integer or a valid string: %r" % level)
    return rv

class _CompatLoggerAdapter(object):

    def __init__(self, logger, extra):
        self.logger = logger
        self.extra = extra

    def setLevel(self, level):
        self.logger.level = _checkLevel(level)

    def process(self, msg, kwargs):
        kwargs["extra"] = self.extra
        return msg, kwargs

    def debug(self, msg, *args, **kwargs):
        self.log(logging.DEBUG, msg, *args, **kwargs)

    def info(self, msg, *args, **kwargs):
        self.log(logging.INFO, msg, *args, **kwargs)

    def warning(self, msg, *args, **kwargs):
        self.log(logging.WARNING, msg, *args, **kwargs)
    warn = warning

    def error(self, msg, *args, **kwargs):
        self.log(logging.ERROR, msg, *args, **kwargs)

    def exception(self, msg, *args, **kwargs):
        kwargs.setdefault("exc_info", 1)
        self.error(msg, *args, **kwargs)

    def critical(self, msg, *args, **kwargs):
        self.log(logging.CRITICAL, msg, *args, **kwargs)
    fatal = critical

    def log(self, level, msg, *args, **kwargs):
        if self.logger.isEnabledFor(level):
            msg, kwargs = self.process(msg, kwargs)
            self._log(level, msg, args, **kwargs)

    def makeRecord(self, name, level, fn, lno, msg, args, exc_info,
            func=None, extra=None):
        rv = logging.LogRecord(name, level, fn, lno, msg, args, exc_info, func)
        if extra is not None:
            for key, value in extra.items():
                if key in ("message", "asctime") or key in rv.__dict__:
                    raise KeyError(
                            "Attempt to override %r in LogRecord" % key)
                rv.__dict__[key] = value
        if multiprocessing is not None:
            rv.processName = multiprocessing.current_process()._name
        else:
            rv.processName = ""
        return rv

    def _log(self, level, msg, args, exc_info=None, extra=None):
        defcaller = "(unknown file)", 0, "(unknown function)"
        if logging._srcfile:
            # IronPython doesn't track Python frames, so findCaller
            # throws an exception on some versions of IronPython.
            # We trap it here so that IronPython can use logging.
            try:
                fn, lno, func = self.logger.findCaller()
            except ValueError:
                fn, lno, func = defcaller
        else:
            fn, lno, func = defcaller
        if exc_info:
            if not isinstance(exc_info, tuple):
                exc_info = sys.exc_info()
        record = self.makeRecord(self.logger.name, level, fn, lno, msg,
                                 args, exc_info, func, extra)
        self.logger.handle(record)

    def isEnabledFor(self, level):
        return self.logger.isEnabledFor(level)

    def addHandler(self, hdlr):
        self.logger.addHandler(hdlr)

    def removeHandler(self, hdlr):
        self.logger.removeHandler(hdlr)

    @property
    def level(self):
        return self.logger.level

try:
    from logging import LoggerAdapter
except ImportError:
    LoggerAdapter = _CompatLoggerAdapter  # noqa

############## itertools.zip_longest #######################################

try:
    from itertools import izip_longest as zip_longest
except ImportError:
    import itertools

    def zip_longest(*args, **kwds):  # noqa
        fillvalue = kwds.get("fillvalue")

        def sentinel(counter=([fillvalue] * (len(args) - 1)).pop):
            yield counter()     # yields the fillvalue, or raises IndexError

        fillers = itertools.repeat(fillvalue)
        iters = [itertools.chain(it, sentinel(), fillers)
                    for it in args]
        try:
            for tup in itertools.izip(*iters):
                yield tup
        except IndexError:
            pass

############## itertools.chain.from_iterable ################################
from itertools import chain

def _compat_chain_from_iterable(iterables):
    for it in iterables:
        for element in it:
            yield element

try:
    chain_from_iterable = getattr(chain, "from_iterable")
except AttributeError:
    chain_from_iterable = _compat_chain_from_iterable

############## logging.handlers.WatchedFileHandler ##########################
import os
from stat import ST_DEV, ST_INO
import platform as _platform

if _platform.system() == "Windows":
    #since windows doesn't go with WatchedFileHandler use FileHandler instead
    WatchedFileHandler = logging.FileHandler
else:
    try:
        from logging.handlers import WatchedFileHandler
    except ImportError:
        class WatchedFileHandler(logging.FileHandler):  # noqa
            """
            A handler for logging to a file, which watches the file
            to see if it has changed while in use. This can happen because of
            usage of programs such as newsyslog and logrotate which perform
            log file rotation. This handler, intended for use under Unix,
            watches the file to see if it has changed since the last emit.
            (A file has changed if its device or inode have changed.)
            If it has changed, the old file stream is closed, and the file
            opened to get a new stream.

            This handler is not appropriate for use under Windows, because
            under Windows open files cannot be moved or renamed - logging
            opens the files with exclusive locks - and so there is no need
            for such a handler. Furthermore, ST_INO is not supported under
            Windows; stat always returns zero for this value.

            This handler is based on a suggestion and patch by Chad J.
            Schroeder.
            """
            def __init__(self, *args, **kwargs):
                logging.FileHandler.__init__(self, *args, **kwargs)

                if not os.path.exists(self.baseFilename):
                    self.dev, self.ino = -1, -1
                else:
                    stat = os.stat(self.baseFilename)
                    self.dev, self.ino = stat[ST_DEV], stat[ST_INO]

            def emit(self, record):
                """
                Emit a record.

                First check if the underlying file has changed, and if it
                has, close the old stream and reopen the file to get the
                current stream.
                """
                if not os.path.exists(self.baseFilename):
                    stat = None
                    changed = 1
                else:
                    stat = os.stat(self.baseFilename)
                    changed = ((stat[ST_DEV] != self.dev) or
                               (stat[ST_INO] != self.ino))
                if changed and self.stream is not None:
                    self.stream.flush()
                    self.stream.close()
                    self.stream = self._open()
                    if stat is None:
                        stat = os.stat(self.baseFilename)
                    self.dev, self.ino = stat[ST_DEV], stat[ST_INO]
                logging.FileHandler.emit(self, record)
<|endoftext|>


<commit_before>from __future__ import absolute_import
import json
import os
import logging
import time
import tempfile
from celery import Celery
from celery.signals import after_setup_task_logger

class myFileHandler(logging.FileHandler):

    def __init__(self, logfile, mode):
        self.logfile = logfile
        super(myFileHandler,self).__init__(self.logfile,mode)

    def getlogfile(self):
        return self.logfile

celery = Celery()
celery.config_from_object('celeryconfig')

# Return a filename of the form imagebuild_.log
def getfilename():
    time_now = str(time.time()).split('.')
    logfile = tempfile.gettempdir() + '/imagebuild_{0:s}.log'.format(time_now[0]+time_now[1])
    return logfile

@after_setup_task_logger.connect
def augment_celery_log(**kwargs):
    logger = logging.getLogger('imagebuilder')
    logfile = getfilename()
    handler = myFileHandler(logfile,'w')
    formatter = logging.Formatter('%(asctime)s - %(message)s')

    if not logger.handlers:
        formatter = logging.Formatter(logging.BASIC_FORMAT)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.propagate = 0

        logger.setLevel(logging.DEBUG)

import mylib.fire as fire

@celery.task
def build(buildconfig, kickstart):

    logger = logging.getLogger('imagebuilder')
    logfile = getfilename()
    handler = myFileHandler(logfile,'w')
    formatter = logging.Formatter('%(asctime)s - %(message)s')
    handler.setFormatter(formatter)

    # replace the handler
    logger.handlers[0] = handler
    logger.info("this is a task %s" % kickstart)
    #import dna.aligner as aligner
    fire.burn(">>>> %s" % kickstart)
<commit_msg>Update main.py<commit_after>from __future__ import absolute_import
import json
import os
import logging
import time
import tempfile
from celery import Celery
from celery.signals import after_setup_task_logger

class myFileHandler(logging.FileHandler):

    def __init__(self, logfile, mode):
        self.logfile = logfile
        super(myFileHandler,self).__init__(self.logfile,mode)

    def getlogfile(self):
        return self.logfile

celery = Celery()
celery.config_from_object('celeryconfig')

# Return a filename of the form imagebuild_.log
def getfilename():
    time_now = str(time.time()).split('.')
    logfile = tempfile.gettempdir() + '/imagebuild_{0:s}.log'.format(time_now[0]+time_now[1])
    return logfile

@after_setup_task_logger.connect
def augment_celery_log(**kwargs):
    logger = logging.getLogger('imagebuilder')
    logfile = getfilename()
    handler = myFileHandler(logfile,'w')
    formatter = logging.Formatter('%(asctime)s - %(message)s')

    if not logger.handlers:
        formatter = logging.Formatter(logging.BASIC_FORMAT)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.propagate = 0

        logger.setLevel(logging.DEBUG)

import mylib.fire as fire

@celery.task
def build(buildconfig, kickstart):

    logger = logging.getLogger('imagebuilder')
    logfile = getfilename()
    handler = myFileHandler(logfile,'w')
    formatter = logging.Formatter('%(asctime)s - %(message)s')
    handler.setFormatter(formatter)

    # replace the handler
    logger.handlers[0] = handler
    logger.info("this is a task %s" % kickstart)
    fire.burn(kickstart)
<|endoftext|>