<commit_before>#!/usr/bin/env python

"""
Produces slice images for Mango volume data stored in NetCDF files.

Currently supported data types are tomo, segmented and labels. If the
data is split into multiple files, those must sit in a common directory
and that directory given as the path argument.

Slices are taken through the center of the volume. If either of the
dimensions for a slice image would be less than 10, that slice is
suppressed. Image output is in .png format.

Typical usage:
    for (data, name, action) in slices(path):
        fp = file(name, 'wb')
        fp.write(data)
        fp.close()

(Requires Python 2.6 or higher.)
"""

import os, os.path, re, struct, sys
import numpy
import bz2

from logger import Logger, LOGGER_INFO, LOGGER_WARNING
import make_image
from nc3header import NC3File

class Histogram:
    """
    Maintains a frequency count for a series of numpy arrays. Entries
    equal to <mask_value> are counted separately. Negative values are
    ignored.
    """

    def __init__(self, mask_value = 0, minval = None, maxval = None):
        # -- save the mask value
        self.mask_value = mask_value
        # -- initialize the frequency count
        self.counts = numpy.array([], dtype = 'uint64')
        # -- initialize the count of masked and total non-negative entries
        self.total = 0
        self.masked = 0

        if maxval is not None:
            self.offset = minval
            self.binsize = (maxval - minval) * (1.0 - 1.0e-12) / 0x10000;
        else:
            self.offset = 0
            self.binsize = 1

    def update(self, slice):
        """
        Updates the frequency count with the data from the numpy array
        <slice>.
        """

        # -- process the new data
        tmp = (slice.flatten() - self.offset) / self.binsize
        flat = numpy.array(tmp, dtype = 'uint16')
        mask = (flat == self.mask_value) | (flat < 0) | (flat > 0xffff)
        new_masked = flat[mask].size
        new_counts = numpy.bincount(numpy.where(mask, 0, flat))

        # -- update the frequency count, resizing if necessary
        s = max(self.counts.size, new_counts.size)
        if s > self.counts.size:
            self.counts.resize(s)
        if s > new_counts.size:
            new_counts.resize(s)
        self.counts += new_counts

        # -- update the count of masked and total non-negative entries
        self.masked += new_masked
        self.total  += new_masked + int(new_counts.sum())

class Slice:
    def __init__(self, size, type, axis, pos):
        self.axis = axis.lower()
        self.pos  = pos
        self.slice_dims = {'x': (size[2], size[1]),
                           'y': (size[2], size[0]),
                           'z': (size[1], size[0]) }[self.axis]
        self.content = numpy.zeros(self.slice_dims, type)

    def update(self, z_slice, z_pos):
        """
        Updates this slices with data from the array <z_slice>, which
        is taken to be at z = <z_pos>.
        """

        if self.axis == 'x':
            self.content[z_pos, :] = z_slice[:, self.pos]
        elif self.axis == 'y':
            self.content[z_pos, :] = z_slice[self.pos, :]
        elif self.axis == 'z':
            if z_pos == self.pos:
                self.content[:, :] = z_slice[:, :]

def get_attribute(file, var, name):
    """
    Looks in the open NC3File <file> for the attribute <name>,
    first in the given variable <var> and then in the global
    attributes. Return the attribute's value if found, or None
    otherwise.
    """

    for attr in var.attributes:
        if attr.name == name:
            return attr.value
    for attr in file.attributes:
        if attr.name == name:
            return attr.value
    return None

def volume_variable(file, var):
    """
    Extracts information pertaining to the volume variable <var> within the
    NC3File <file>. The name, shape, origin and data type - mapped to the
    corresponding numpy type - are stored.
    """

    # -- remember the name
    name = var.name

    # -- extract the variable's shape
    dims = var.dimensions
    size = (dims[2].value, dims[1].value, dims[0].value)
    zdim_total = get_attribute(file, var, 'zdim_total')
    if zdim_total is not None:
        size = (size[0], size[1], zdim_total[0])

    # -- determine the origin
    origin = get_attribute(file, var, 'coordinate_origin_xyz') or (0, 0, 0)

    # -- determine the data type
    (dtype, big_endian_type) = {
      'b': (numpy.uint8,  numpy.dtype(">u1")),
      'h': (numpy.uint16, numpy.dtype(">u2")),
      'i': (numpy.int32,  numpy.dtype(">i4")),
      'f': (numpy.float32, numpy.dtype(">f4")) }[var.python_type_code]

    return { 'name'  : name,
             'size'  : size,
             'origin': origin,
             'dtype' : dtype,
             'big_endian_type': big_endian_type }

def find_variable(path):
    """
    Looks for a volume variable to extract slice images from in the
    NetCDF file at location <path>. Returns an descriptor object of
    class VolumeVariable if something appropriate is found, or None
    otherwise.
    """

    # -- the default return value
    var = None

    # -- open the NetCDF file
    file = NC3File(path)

    # -- loop through all the variables in the file
    for v in file.variables:
        if (len(v.dimensions) == 3 and v.dimensions[0].value > 1 and
            v.python_type_code in 'bhif'
            ):
            return volume_variable(file, v)

def z_slices(variable, path):
    """
    A generator method that yields constant z slices corresponding
    to the variable <var> from the NetCDF file located at <path>.

    Each value produced is a pair containing the z coordinate of
    the slice and a two-dimensional numpy array containing the
    extracted data.

    Usage example:

        for (z, data) in z_slices(var, path):
            ... # do something with data
    """

    file = NC3File(path)

    for var in file.variables:
        if var.name == variable['name']:
            break
    else:
        return

    if volume_variable(file, var) != variable:
        raise RuntimeError("variable mismatch between files")

    (x, y, z) = variable['size']

    z_range = get_attribute(file, var, 'zdim_range')
    if z_range is None:
        z_range = range(0, z)
    else:
        z_range = range(z_range[0], z_range[1] + 1)

    bytes_per_slice = x * y * var.element_size
    offset = var.data_start

    if path.endswith('.bz2'):
        fp = bz2.BZ2File(path, 'r', 1024 * 1024)
    else:
        fp = open(path, "rb")
    fp.seek(var.data_start)
    for z in z_range:
        buffer = fp.read(bytes_per_slice)
        if len(buffer) < bytes_per_slice:
            yield (z, None, "insufficient data")
            break
        data = numpy.fromstring(buffer, variable['big_endian_type'])
        data.shape = (y, x)
        yield (z, data)
    fp.close()

def bottom_percentile(histogram, p):
    """
    Returns the smallest number i such that at least <p> percent of
    the non-masked entries counted so far have value i or less.
    """
    threshold = p * (histogram.total - histogram.masked) / 100.0
    count = 0
    for i in xrange(histogram.counts.size):
        count += int(histogram.counts[i])
        if count >= threshold:
            return histogram.offset + i * histogram.binsize

def top_percentile(histogram, p):
    """
    Returns the largest number i such that at least <p> percent of
    the non-masked entries counted so far have value i or more.
    """
    threshold = p * (histogram.total - histogram.masked) / 100.0
    count = 0
    for i in xrange(histogram.counts.size-1, -1, -1):
        count += int(histogram.counts[i])
        if count >= threshold:
            return histogram.offset + i * histogram.binsize

def default_slice_set(var, delta, basename):
    """
    Creates a default list of empty slice instances, paired with
    associated file names, based on the shape of the volume variable
    <var>. For each axis, a slice centered at that axis is created,
    provided that the extend of the slice in both directions would be at
    least <delta>.
    """
    pos = list((x - 1) / 2 for x in var['size'])
    size = var['size']
    dtype = var['dtype']
    origin = var['origin']

    def slice_and_name(axis, pos, offset):
        s = Slice(size, dtype, axis, pos)
        n = "slice%c%d_%s.png" % (axis.upper(), pos + offset, basename)
        return (s, n)

    slices = []
    if size[1] > delta and size[2] > delta:
        slices.append(slice_and_name('x', pos[0], origin[0]))
    if size[0] > delta and size[2] > delta:
        slices.append(slice_and_name('y', pos[1], origin[1]))
    if size[0] > delta and size[1] > delta:
        slices.append(slice_and_name('z', pos[2], origin[2]))

    return slices

def image_data(slice, lo, hi, mask_val, info, thumb_size = None):
    """
    Converts a slice into an image and returns it as a PNG encoded
    string. The values <lo> and <hi> specify the range of relevant
    data. Entries equal to <mask_val> are treated as masked.

    An appropriate encoding mode is determined from the data type
    and value range. Details of the actual encoding are deferred to
    the function 'make_image' in the package of the same name.
    """

    # -- determine the encoding mode
    content = slice.content
    if content.dtype == numpy.uint8:
        if hi <= 1:
            img_mode = make_image.BLACK_AND_WHITE
        else:
            img_mode = make_image.COLOR_CODED_FIXED
    elif content.dtype == numpy.uint16:
        img_mode = make_image.GRAYSCALE
    elif content.dtype == numpy.int32:
        img_mode = make_image.COLOR_CODED
    elif content.dtype == numpy.float32:
        img_mode = make_image.GRAYSCALE
    else:
        raise Exception("unexpected array type: %s" % content.dtype)

    # -- generate slice-specific info
    myinfo = info.copy()
    myinfo.update({ 'slice-axis': slice.axis, 'slice-pos': slice.pos })

    # -- generate and return the data
    return make_image.make_image(content, lo, hi, mask_val, img_mode,
                                 thumb_size, myinfo)

def data_range(var, entries):
    minval = maxval = None

    for filename in entries:
        for tmp in z_slices(var, filename):
            z, data = tmp[:2]
            if data is not None:
                lo = numpy.min(data)
                hi = numpy.max(data)
                if minval is None or lo < minval:
                    minval = lo
                if maxval is None or hi > maxval:
                    maxval = hi

    return (minval, maxval)

def slices(path,
           existing = [],
           replace = False,
           dry_run = False,
           sizes = (None,),
           info = {}):
    """
    A generator which extracts slice images from a Mango volume data set
    stored in a collection of NetCDF files.

    The parameter <path> specifies a single NetCDF file or a directory
    containing a single volume split into several NetCDF files.

    Basic usage:
        for (data, name, action) in slices(path):
            fp = file(name, 'wb')
            fp.write(data)
            fp.close()
    """

    path      = re.sub("/$", "", path)
    name      = re.sub("[._]nc$", "", os.path.basename(path))
    basename  = re.sub("^tomo", "tom", re.sub("^segmented", "seg", name))
    img_mode  = None
    log       = Logger()

    # -- collect the list of files to be processed
    if os.path.isdir(path):
        entries = list(os.path.join(root, f)
                       for (root, dirs, files) in os.walk(path)
                       for f in files)
        entries.sort()
    else:
        entries = [ path ]
    if not entries:
        return

    # -- find a useable volume variable and analyse it
    filename = entries[0]
    log.writeln("slices(): looking for a volume variable in %s..." %
                os.path.basename(filename))
    var = find_variable(filename)
    if var is None :
        log.writeln("No appropriate volume data found.")
        return

    # -- determine the appropriate mask value
    mask_value = { numpy.uint8:  0xff,
                   numpy.uint16: 0xffff,
                   numpy.int32:  0x7fffffff,
                   numpy.float32: 1.0e30 }[var['dtype']]

    # -- initialize the slice set to be created
    slices = default_slice_set(var, 10, basename)
    r_or_s = 'REPLACE' if replace else 'SKIP'
    actions = list((r_or_s if n in existing else 'ADD')
                   for (s, n) in slices)
    slices = list(slices[i] + (actions[i],)
                  for i in range(len(slices))
                  if actions[i] != 'SKIP')

    if len(slices) == 0:
        log.writeln("No slices are to be made.")
        return

    if not dry_run:
        # -- initialize the histogram
        if var['dtype'] == numpy.float32:
            log.writeln("Determining the data range...")
            (minval, maxval) = data_range(var, entries)
            hist = Histogram(mask_value, minval, maxval)
        else:
            hist = Histogram(mask_value)

        # -- loop through files and copy data into slice arrays
        for filename in entries:
            log.writeln("Processing %s..." % os.path.basename(filename))
            for tmp in z_slices(var, filename):
                z, data = tmp[:2]
                if data is None:
                    log.writeln(tmp[2] + " at z = %d" % z, LOGGER_WARNING)
                else:
                    hist.update(data)
                    for (s, n, a) in slices:
                        s.update(data, z)

        # -- analyse histogram to determine 'lo' and 'hi' values
        log.writeln("Analysing the histogram...")
        if name.startswith("tom"):
            # -- determine 0.1 and 99.9 percentile for contrast stretching
            lo = bottom_percentile(hist, 0.1)
            hi = top_percentile(hist, 0.1)
        else:
            lo = 0
            hi = hist.counts.size - 1

    # -- encode slices as PNG images
    log.writeln("Making the images...")
    for (s, n, a) in slices:
        for sz in sizes:
            if dry_run:
                data = make_image.make_dummy(n, 256, 256, sz)
            else:
                data = image_data(s, lo, hi, mask_value, info, sz)
            prefix = ("__%sx%s__" % sz) if sz else ""
            yield (data, prefix + n, a)

    # -- report success
    log.writeln("Slice image generation finished.")

if __name__ == "__main__":
    Logger().priority = LOGGER_INFO
    source = sys.argv[1]
    s = slices(source, info = { 'source': source }, sizes = (None, (80,80)))

    for (data, name, action) in s:
        fp = file(name, 'wb')
        fp.write(data)
        fp.close()
<commit_msg>adjusted make_slices.py to change in NC3File class<commit_after>#!/usr/bin/env python

"""
Produces slice images for Mango volume data stored in NetCDF files.

Currently supported data types are tomo, segmented and labels. If the
data is split into multiple files, those must sit in a common directory
and that directory given as the path argument.

Slices are taken through the center of the volume. If either of the
dimensions for a slice image would be less than 10, that slice is
suppressed. Image output is in .png format.

Typical usage:
    for (data, name, action) in slices(path):
        fp = file(name, 'wb')
        fp.write(data)
        fp.close()

(Requires Python 2.6 or higher.)
"""

import os, os.path, re, struct, sys
import numpy
import bz2

from logger import Logger, LOGGER_INFO, LOGGER_WARNING
import make_image
from nc3header import NC3File

class Histogram:
    """
    Maintains a frequency count for a series of numpy arrays. Entries
    equal to <mask_value> are counted separately. Negative values are
    ignored.
    """

    def __init__(self, mask_value = 0, minval = None, maxval = None):
        # -- save the mask value
        self.mask_value = mask_value
        # -- initialize the frequency count
        self.counts = numpy.array([], dtype = 'uint64')
        # -- initialize the count of masked and total non-negative entries
        self.total = 0
        self.masked = 0

        if maxval is not None:
            self.offset = minval
            self.binsize = (maxval - minval) * (1.0 - 1.0e-12) / 0x10000;
        else:
            self.offset = 0
            self.binsize = 1

    def update(self, slice):
        """
        Updates the frequency count with the data from the numpy array
        <slice>.
        """

        # -- process the new data
        tmp = (slice.flatten() - self.offset) / self.binsize
        flat = numpy.array(tmp, dtype = 'uint16')
        mask = (flat == self.mask_value) | (flat < 0) | (flat > 0xffff)
        new_masked = flat[mask].size
        new_counts = numpy.bincount(numpy.where(mask, 0, flat))

        # -- update the frequency count, resizing if necessary
        s = max(self.counts.size, new_counts.size)
        if s > self.counts.size:
            self.counts.resize(s)
        if s > new_counts.size:
            new_counts.resize(s)
        self.counts += new_counts

        # -- update the count of masked and total non-negative entries
        self.masked += new_masked
        self.total  += new_masked + int(new_counts.sum())

class Slice:
    def __init__(self, size, type, axis, pos):
        self.axis = axis.lower()
        self.pos  = pos
        self.slice_dims = {'x': (size[2], size[1]),
                           'y': (size[2], size[0]),
                           'z': (size[1], size[0]) }[self.axis]
        self.content = numpy.zeros(self.slice_dims, type)

    def update(self, z_slice, z_pos):
        """
        Updates this slices with data from the array <z_slice>, which
        is taken to be at z = <z_pos>.
        """

        if self.axis == 'x':
            self.content[z_pos, :] = z_slice[:, self.pos]
        elif self.axis == 'y':
            self.content[z_pos, :] = z_slice[self.pos, :]
        elif self.axis == 'z':
            if z_pos == self.pos:
                self.content[:, :] = z_slice[:, :]

def get_attribute(info, var, name):
    """
    Looks in the open NC3File <info> for the attribute <name>,
    first in the given variable <var> and then in the global
    attributes. Return the attribute's value if found, or None
    otherwise.
    """

    for attr in var.attributes:
        if attr.name == name:
            return attr.value
    for attr in info.attributes:
        if attr.name == name:
            return attr.value
    return None

def volume_variable(info, var):
    """
    Extracts information pertaining to the volume variable <var> within the
    NC3File <info>. The name, shape, origin and data type - mapped to the
    corresponding numpy type - are stored.
    """

    # -- remember the name
    name = var.name

    # -- extract the variable's shape
    dims = var.dimensions
    size = (dims[2].value, dims[1].value, dims[0].value)
    zdim_total = get_attribute(info, var, 'zdim_total')
    if zdim_total is not None:
        size = (size[0], size[1], zdim_total[0])

    # -- determine the origin
    origin = get_attribute(info, var, 'coordinate_origin_xyz') or (0, 0, 0)

    # -- determine the data type
    (dtype, big_endian_type) = {
      'b': (numpy.uint8,  numpy.dtype(">u1")),
      'h': (numpy.uint16, numpy.dtype(">u2")),
      'i': (numpy.int32,  numpy.dtype(">i4")),
      'f': (numpy.float32, numpy.dtype(">f4")) }[var.python_type_code]

    return { 'name'  : name,
             'size'  : size,
             'origin': origin,
             'dtype' : dtype,
             'big_endian_type': big_endian_type }

def find_variable(path):
    """
    Looks for a volume variable to extract slice images from in the
    NetCDF info at location <path>. Returns an descriptor object of
    class VolumeVariable if something appropriate is found, or None
    otherwise.
    """

    # -- the default return value
    var = None

    # -- open the NetCDF file
    fp = open(path)
    try:
        info = NC3File(fp)
    finally:
        fp.close()

    # -- loop through all the variables in the file
    for v in info.variables:
        if (len(v.dimensions) == 3 and v.dimensions[0].value > 1 and
            v.python_type_code in 'bhif'
            ):
            return volume_variable(info, v)

def z_slices(variable, path):
    """
    A generator method that yields constant z slices corresponding
    to the variable <var> from the NetCDF file located at <path>.

    Each value produced is a pair containing the z coordinate of
    the slice and a two-dimensional numpy array containing the
    extracted data.

    Usage example:

        for (z, data) in z_slices(var, path):
            ... # do something with data
    """

    fp = open(path)
    try:
        info = NC3File(fp)
    finally:
        fp.close()

    for var in info.variables:
        if var.name == variable['name']:
            break
    else:
        return

    if volume_variable(info, var) != variable:
        raise RuntimeError("variable mismatch between files")

    (x, y, z) = variable['size']

    z_range = get_attribute(info, var, 'zdim_range')
    if z_range is None:
        z_range = range(0, z)
    else:
        z_range = range(z_range[0], z_range[1] + 1)

    bytes_per_slice = x * y * var.element_size
    offset = var.data_start

    if path.endswith('.bz2'):
        fp = bz2.BZ2File(path, 'r', 1024 * 1024)
    else:
        fp = open(path, "rb")
    fp.seek(var.data_start)
    for z in z_range:
        buffer = fp.read(bytes_per_slice)
        if len(buffer) < bytes_per_slice:
            yield (z, None, "insufficient data")
            break
        data = numpy.fromstring(buffer, variable['big_endian_type'])
        data.shape = (y, x)
        yield (z, data)
    fp.close()

def bottom_percentile(histogram, p):
    """
    Returns the smallest number i such that at least <p> percent of
    the non-masked entries counted so far have value i or less.
    """
    threshold = p * (histogram.total - histogram.masked) / 100.0
    count = 0
    for i in xrange(histogram.counts.size):
        count += int(histogram.counts[i])
        if count >= threshold:
            return histogram.offset + i * histogram.binsize

def top_percentile(histogram, p):
    """
    Returns the largest number i such that at least <p> percent of
    the non-masked entries counted so far have value i or more.
    """
    threshold = p * (histogram.total - histogram.masked) / 100.0
    count = 0
    for i in xrange(histogram.counts.size-1, -1, -1):
        count += int(histogram.counts[i])
        if count >= threshold:
            return histogram.offset + i * histogram.binsize

def default_slice_set(var, delta, basename):
    """
    Creates a default list of empty slice instances, paired with
    associated file names, based on the shape of the volume variable
    <var>. For each axis, a slice centered at that axis is created,
    provided that the extend of the slice in both directions would be at
    least <delta>.
    """
    pos = list((x - 1) / 2 for x in var['size'])
    size = var['size']
    dtype = var['dtype']
    origin = var['origin']

    def slice_and_name(axis, pos, offset):
        s = Slice(size, dtype, axis, pos)
        n = "slice%c%d_%s.png" % (axis.upper(), pos + offset, basename)
        return (s, n)

    slices = []
    if size[1] > delta and size[2] > delta:
        slices.append(slice_and_name('x', pos[0], origin[0]))
    if size[0] > delta and size[2] > delta:
        slices.append(slice_and_name('y', pos[1], origin[1]))
    if size[0] > delta and size[1] > delta:
        slices.append(slice_and_name('z', pos[2], origin[2]))

    return slices

def image_data(slice, lo, hi, mask_val, info, thumb_size = None):
    """
    Converts a slice into an image and returns it as a PNG encoded
    string. The values <lo> and <hi> specify the range of relevant
    data. Entries equal to <mask_val> are treated as masked.

    An appropriate encoding mode is determined from the data type
    and value range. Details of the actual encoding are deferred to
    the function 'make_image' in the package of the same name.
    """

    # -- determine the encoding mode
    content = slice.content
    if content.dtype == numpy.uint8:
        if hi <= 1:
            img_mode = make_image.BLACK_AND_WHITE
        else:
            img_mode = make_image.COLOR_CODED_FIXED
    elif content.dtype == numpy.uint16:
        img_mode = make_image.GRAYSCALE
    elif content.dtype == numpy.int32:
        img_mode = make_image.COLOR_CODED
    elif content.dtype == numpy.float32:
        img_mode = make_image.GRAYSCALE
    else:
        raise Exception("unexpected array type: %s" % content.dtype)

    # -- generate slice-specific info
    myinfo = info.copy()
    myinfo.update({ 'slice-axis': slice.axis, 'slice-pos': slice.pos })

    # -- generate and return the data
    return make_image.make_image(content, lo, hi, mask_val, img_mode,
                                 thumb_size, myinfo)

def data_range(var, entries):
    minval = maxval = None

    for filename in entries:
        for tmp in z_slices(var, filename):
            z, data = tmp[:2]
            if data is not None:
                lo = numpy.min(data)
                hi = numpy.max(data)
                if minval is None or lo < minval:
                    minval = lo
                if maxval is None or hi > maxval:
                    maxval = hi

    return (minval, maxval)

def slices(path,
           existing = [],
           replace = False,
           dry_run = False,
           sizes = (None,),
           info = {}):
    """
    A generator which extracts slice images from a Mango volume data set
    stored in a collection of NetCDF files.

    The parameter <path> specifies a single NetCDF file or a directory
    containing a single volume split into several NetCDF files.

    Basic usage:
        for (data, name, action) in slices(path):
            fp = file(name, 'wb')
            fp.write(data)
            fp.close()
    """

    path      = re.sub("/$", "", path)
    name      = re.sub("[._]nc$", "", os.path.basename(path))
    basename  = re.sub("^tomo", "tom", re.sub("^segmented", "seg", name))
    img_mode  = None
    log       = Logger()

    # -- collect the list of files to be processed
    if os.path.isdir(path):
        entries = list(os.path.join(root, f)
                       for (root, dirs, files) in os.walk(path)
                       for f in files)
        entries.sort()
    else:
        entries = [ path ]
    if not entries:
        return

    # -- find a useable volume variable and analyse it
    filename = entries[0]
    log.writeln("slices(): looking for a volume variable in %s..." %
                os.path.basename(filename))
    var = find_variable(filename)
    if var is None :
        log.writeln("No appropriate volume data found.")
        return

    # -- determine the appropriate mask value
    mask_value = { numpy.uint8:  0xff,
                   numpy.uint16: 0xffff,
                   numpy.int32:  0x7fffffff,
                   numpy.float32: 1.0e30 }[var['dtype']]

    # -- initialize the slice set to be created
    slices = default_slice_set(var, 10, basename)
    r_or_s = 'REPLACE' if replace else 'SKIP'
    actions = list((r_or_s if n in existing else 'ADD')
                   for (s, n) in slices)
    slices = list(slices[i] + (actions[i],)
                  for i in range(len(slices))
                  if actions[i] != 'SKIP')

    if len(slices) == 0:
        log.writeln("No slices are to be made.")
        return

    if not dry_run:
        # -- initialize the histogram
        if var['dtype'] == numpy.float32:
            log.writeln("Determining the data range...")
            (minval, maxval) = data_range(var, entries)
            hist = Histogram(mask_value, minval, maxval)
        else:
            hist = Histogram(mask_value)

        # -- loop through files and copy data into slice arrays
        for filename in entries:
            log.writeln("Processing %s..." % os.path.basename(filename))
            for tmp in z_slices(var, filename):
                z, data = tmp[:2]
                if data is None:
                    log.writeln(tmp[2] + " at z = %d" % z, LOGGER_WARNING)
                else:
                    hist.update(data)
                    for (s, n, a) in slices:
                        s.update(data, z)

        # -- analyse histogram to determine 'lo' and 'hi' values
        log.writeln("Analysing the histogram...")
        if name.startswith("tom"):
            # -- determine 0.1 and 99.9 percentile for contrast stretching
            lo = bottom_percentile(hist, 0.1)
            hi = top_percentile(hist, 0.1)
        else:
            lo = 0
            hi = hist.counts.size - 1

    # -- encode slices as PNG images
    log.writeln("Making the images...")
    for (s, n, a) in slices:
        for sz in sizes:
            if dry_run:
                data = make_image.make_dummy(n, 256, 256, sz)
            else:
                data = image_data(s, lo, hi, mask_value, info, sz)
            prefix = ("__%sx%s__" % sz) if sz else ""
            yield (data, prefix + n, a)

    # -- report success
    log.writeln("Slice image generation finished.")

if __name__ == "__main__":
    Logger().priority = LOGGER_INFO
    source = sys.argv[1]
    s = slices(source, info = { 'source': source }, sizes = (None, (80,80)))

    for (data, name, action) in s:
        fp = file(name, 'wb')
        fp.write(data)
        fp.close()
<|endoftext|>


<commit_before>"""
Install Python and Node prerequisites.
"""

from distutils import sysconfig
import hashlib
import os
import re
import sys

from paver.easy import sh, task, BuildFailure

from .utils.envs import Env
from .utils.timer import timed

PREREQS_STATE_DIR = os.getenv('PREREQ_CACHE_DIR', Env.REPO_ROOT / '.prereqs_cache')
NPM_REGISTRY = "https://registry.npmjs.org/"
NO_PREREQ_MESSAGE = "NO_PREREQ_INSTALL is set, not installing prereqs"
COVERAGE_REQ_FILE = 'requirements/edx/coverage.txt'

# If you make any changes to this list you also need to make
# a corresponding change to circle.yml, which is how the python
# prerequisites are installed for builds on circleci.com
PYTHON_REQ_FILES = [
    'requirements/edx/pre.txt',
    'requirements/edx/github.txt',
    'requirements/edx/custom.txt',
    'requirements/edx/local.txt',
    'requirements/edx/base.txt',
    'requirements/edx/paver.txt',
    'requirements/edx/post.txt',
]

# Developers can have private requirements, for local copies of github repos,
# or favorite debugging tools, etc.
PRIVATE_REQS = 'requirements/private.txt'
if os.path.exists(PRIVATE_REQS):
    PYTHON_REQ_FILES.append(PRIVATE_REQS)

def no_prereq_install():
    """
    Determine if NO_PREREQ_INSTALL should be truthy or falsy.
    """
    vals = {
        '0': False,
        '1': True,
        'true': True,
        'false': False,
    }

    val = os.environ.get("NO_PREREQ_INSTALL", 'False').lower()

    try:
        return vals[val]
    except KeyError:
        return False

def create_prereqs_cache_dir():
    """Create the directory for storing the hashes, if it doesn't exist already."""
    try:
        os.makedirs(PREREQS_STATE_DIR)
    except OSError:
        if not os.path.isdir(PREREQS_STATE_DIR):
            raise

def compute_fingerprint(path_list):
    """
    Hash the contents of all the files and directories in `path_list`.
    Returns the hex digest.
    """

    hasher = hashlib.sha1()

    for path_item in path_list:

        # For directories, create a hash based on the modification times
        # of first-level subdirectories
        if os.path.isdir(path_item):
            for dirname in sorted(os.listdir(path_item)):
                path_name = os.path.join(path_item, dirname)
                if os.path.isdir(path_name):
                    hasher.update(str(os.stat(path_name).st_mtime))

        # For files, hash the contents of the file
        if os.path.isfile(path_item):
            with open(path_item, "rb") as file_handle:
                hasher.update(file_handle.read())

    return hasher.hexdigest()

def prereq_cache(cache_name, paths, install_func):
    """
    Conditionally execute `install_func()` only if the files/directories
    specified by `paths` have changed.

    If the code executes successfully (no exceptions are thrown), the cache
    is updated with the new hash.
    """
    # Retrieve the old hash
    cache_filename = cache_name.replace(" ", "_")
    cache_file_path = os.path.join(PREREQS_STATE_DIR, "{}.sha1".format(cache_filename))
    old_hash = None
    if os.path.isfile(cache_file_path):
        with open(cache_file_path) as cache_file:
            old_hash = cache_file.read()

    # Compare the old hash to the new hash
    # If they do not match (either the cache hasn't been created, or the files have changed),
    # then execute the code within the block.
    new_hash = compute_fingerprint(paths)
    if new_hash != old_hash:
        install_func()

        # Update the cache with the new hash
        # If the code executed within the context fails (throws an exception),
        # then this step won't get executed.
        create_prereqs_cache_dir()
        with open(cache_file_path, "w") as cache_file:
            # Since the pip requirement files are modified during the install
            # process, we need to store the hash generated AFTER the installation
            post_install_hash = compute_fingerprint(paths)
            cache_file.write(post_install_hash)
    else:
        print '{cache} unchanged, skipping...'.format(cache=cache_name)

def node_prereqs_installation():
    """
    Configures npm and installs Node prerequisites
    """
    cb_error_text = "Subprocess return code: 1"
    sh("test `npm config get registry` = \"{reg}\" || "
       "(echo setting registry; npm config set registry"
       " {reg})".format(reg=NPM_REGISTRY))

    # Error handling around a race condition that produces "cb() never called" error. This
    # evinces itself as `cb_error_text` and it ought to disappear when we upgrade
    # npm to 3 or higher. TODO: clean this up when we do that.
    try:
        sh('npm install')
    except BuildFailure, error_text:
        if cb_error_text in error_text:
            print "npm install error detected. Retrying..."
            sh('npm install')
        else:
            raise BuildFailure(error_text)

def python_prereqs_installation():
    """
    Installs Python prerequisites
    """
    for req_file in PYTHON_REQ_FILES:
        pip_install_req_file(req_file)

def pip_install_req_file(req_file):
    """Pip install the requirements file."""
    pip_cmd = 'pip install -q --disable-pip-version-check --exists-action w'
    sh("{pip_cmd} -r {req_file}".format(pip_cmd=pip_cmd, req_file=req_file))

@task
@timed
def install_node_prereqs():
    """
    Installs Node prerequisites
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    prereq_cache("Node prereqs", ["package.json"], node_prereqs_installation)

# To add a package to the uninstall list, just add it to this list! No need
# to touch any other part of this file.
PACKAGES_TO_UNINSTALL = [
    "South",                        # Because it interferes with Django 1.8 migrations.
    "edxval",                       # Because it was bork-installed somehow.
    "django-storages",
    "django-oauth2-provider",       # Because now it's called edx-django-oauth2-provider.
    "edx-oauth2-provider",          # Because it moved from github to pypi
    "i18n-tools",                   # Because now it's called edx-i18n-tools
]

@task
@timed
def uninstall_python_packages():
    """
    Uninstall Python packages that need explicit uninstallation.

    Some Python packages that we no longer want need to be explicitly
    uninstalled, notably, South.  Some other packages were once installed in
    ways that were resistant to being upgraded, like edxval.  Also uninstall
    them.

    """
    # So that we don't constantly uninstall things, use a hash of the packages
    # to be uninstalled.  Check it, and skip this if we're up to date.
    hasher = hashlib.sha1()
    hasher.update(repr(PACKAGES_TO_UNINSTALL))
    expected_version = hasher.hexdigest()
    state_file_path = os.path.join(PREREQS_STATE_DIR, "Python_uninstall.sha1")
    create_prereqs_cache_dir()

    if os.path.isfile(state_file_path):
        with open(state_file_path) as state_file:
            version = state_file.read()
        if version == expected_version:
            print 'Python uninstalls unchanged, skipping...'
            return

    # Run pip to find the packages we need to get rid of.  Believe it or not,
    # edx-val is installed in a way that it is present twice, so we have a loop
    # to really really get rid of it.
    for _ in range(3):
        uninstalled = False
        frozen = sh("pip freeze", capture=True)

        for package_name in PACKAGES_TO_UNINSTALL:
            if package_in_frozen(package_name, frozen):
                # Uninstall the pacakge
                sh("pip uninstall --disable-pip-version-check -y {}".format(package_name))
                uninstalled = True
        if not uninstalled:
            break
    else:
        # We tried three times and didn't manage to get rid of the pests.
        print "Couldn't uninstall unwanted Python packages!"
        return

    # Write our version.
    with open(state_file_path, "w") as state_file:
        state_file.write(expected_version)

def package_in_frozen(package_name, frozen_output):
    """Is this package in the output of 'pip freeze'?"""
    # Look for either:
    #
    #   PACKAGE-NAME==
    #
    # or:
    #
    #   blah_blah#egg=package_name-version
    #
    pattern = r"(?mi)^{pkg}==|#egg={pkg_under}-".format(
        pkg=re.escape(package_name),
        pkg_under=re.escape(package_name.replace("-", "_")),
    )
    return bool(re.search(pattern, frozen_output))

@task
@timed
def install_coverage_prereqs():
    """ Install python prereqs for measuring coverage. """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return
    pip_install_req_file(COVERAGE_REQ_FILE)

@task
@timed
def install_python_prereqs():
    """
    Installs Python prerequisites.
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    uninstall_python_packages()

    # Include all of the requirements files in the fingerprint.
    files_to_fingerprint = list(PYTHON_REQ_FILES)

    # Also fingerprint the directories where packages get installed:
    # ("/edx/app/edxapp/venvs/edxapp/lib/python2.7/site-packages")
    files_to_fingerprint.append(sysconfig.get_python_lib())

    # In a virtualenv, "-e installs" get put in a src directory.
    src_dir = os.path.join(sys.prefix, "src")
    if os.path.isdir(src_dir):
        files_to_fingerprint.append(src_dir)

    # Also fingerprint this source file, so that if the logic for installations
    # changes, we will redo the installation.
    this_file = __file__
    if this_file.endswith(".pyc"):
        this_file = this_file[:-1]      # use the .py file instead of the .pyc
    files_to_fingerprint.append(this_file)

    prereq_cache("Python prereqs", files_to_fingerprint, python_prereqs_installation)

@task
@timed
def install_prereqs():
    """
    Installs Node and Python prerequisites
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    install_node_prereqs()
    install_python_prereqs()
    log_installed_python_prereqs()

def log_installed_python_prereqs():
    """  Logs output of pip freeze for debugging. """
    sh("pip freeze > {}".format(Env.GEN_LOG_DIR + "/pip_freeze.log"))
    return
<commit_msg>rebase-20140926: Fix errors from conflicts resolution<commit_after>"""
Install Python and Node prerequisites.
"""

from distutils import sysconfig
import hashlib
import os
import re
import sys

from paver.easy import sh, task, BuildFailure

from .utils.envs import Env
from .utils.timer import timed

PREREQS_STATE_DIR = os.getenv('PREREQ_CACHE_DIR', Env.REPO_ROOT / '.prereqs_cache')
NPM_REGISTRY = "https://registry.npmjs.org/"
NO_PREREQ_MESSAGE = "NO_PREREQ_INSTALL is set, not installing prereqs"
COVERAGE_REQ_FILE = 'requirements/edx/coverage.txt'

# If you make any changes to this list you also need to make
# a corresponding change to circle.yml, which is how the python
# prerequisites are installed for builds on circleci.com
PYTHON_REQ_FILES = [
    'requirements/edx/pre.txt',
    'requirements/edx/github.txt',
    'requirements/edx/custom.txt',
    'requirements/edx/local.txt',
    'requirements/edx/base.txt',
    'requirements/edx/paver.txt',
    'requirements/edx/post.txt',
]

# Developers can have private requirements, for local copies of github repos,
# or favorite debugging tools, etc.
PRIVATE_REQS = 'requirements/private.txt'
if os.path.exists(PRIVATE_REQS):
    PYTHON_REQ_FILES.append(PRIVATE_REQS)

def no_prereq_install():
    """
    Determine if NO_PREREQ_INSTALL should be truthy or falsy.
    """
    vals = {
        '0': False,
        '1': True,
        'true': True,
        'false': False,
    }

    val = os.environ.get("NO_PREREQ_INSTALL", 'False').lower()

    try:
        return vals[val]
    except KeyError:
        return False

def create_prereqs_cache_dir():
    """Create the directory for storing the hashes, if it doesn't exist already."""
    try:
        os.makedirs(PREREQS_STATE_DIR)
    except OSError:
        if not os.path.isdir(PREREQS_STATE_DIR):
            raise

def compute_fingerprint(path_list):
    """
    Hash the contents of all the files and directories in `path_list`.
    Returns the hex digest.
    """

    hasher = hashlib.sha1()

    for path_item in path_list:

        # For directories, create a hash based on the modification times
        # of first-level subdirectories
        if os.path.isdir(path_item):
            for dirname in sorted(os.listdir(path_item)):
                path_name = os.path.join(path_item, dirname)
                if os.path.isdir(path_name):
                    hasher.update(str(os.stat(path_name).st_mtime))

        # For files, hash the contents of the file
        if os.path.isfile(path_item):
            with open(path_item, "rb") as file_handle:
                hasher.update(file_handle.read())

    return hasher.hexdigest()

def prereq_cache(cache_name, paths, install_func):
    """
    Conditionally execute `install_func()` only if the files/directories
    specified by `paths` have changed.

    If the code executes successfully (no exceptions are thrown), the cache
    is updated with the new hash.
    """
    # Retrieve the old hash
    cache_filename = cache_name.replace(" ", "_")
    cache_file_path = os.path.join(PREREQS_STATE_DIR, "{}.sha1".format(cache_filename))
    old_hash = None
    if os.path.isfile(cache_file_path):
        with open(cache_file_path) as cache_file:
            old_hash = cache_file.read()

    # Compare the old hash to the new hash
    # If they do not match (either the cache hasn't been created, or the files have changed),
    # then execute the code within the block.
    new_hash = compute_fingerprint(paths)
    if new_hash != old_hash:
        install_func()

        # Update the cache with the new hash
        # If the code executed within the context fails (throws an exception),
        # then this step won't get executed.
        create_prereqs_cache_dir()
        with open(cache_file_path, "w") as cache_file:
            # Since the pip requirement files are modified during the install
            # process, we need to store the hash generated AFTER the installation
            post_install_hash = compute_fingerprint(paths)
            cache_file.write(post_install_hash)
    else:
        print '{cache} unchanged, skipping...'.format(cache=cache_name)

def node_prereqs_installation():
    """
    Configures npm and installs Node prerequisites
    """
    cb_error_text = "Subprocess return code: 1"
    sh("test `npm config get registry` = \"{reg}\" || "
       "(echo setting registry; npm config set registry"
       " {reg})".format(reg=NPM_REGISTRY))

    # Error handling around a race condition that produces "cb() never called" error. This
    # evinces itself as `cb_error_text` and it ought to disappear when we upgrade
    # npm to 3 or higher. TODO: clean this up when we do that.
    try:
        sh('npm install')
    except BuildFailure, error_text:
        if cb_error_text in error_text:
            print "npm install error detected. Retrying..."
            sh('npm install')
        else:
            raise BuildFailure(error_text)

def python_prereqs_installation():
    """
    Installs Python prerequisites
    """
    for req_file in PYTHON_REQ_FILES:
        pip_install_req_file(req_file)

def pip_install_req_file(req_file):
    """Pip install the requirements file."""
    pip_cmd = 'pip install -q --disable-pip-version-check --exists-action w'
    sh("rm -rf /edx/app/edxapp/venvs/edxapp/build/*")
    sh("{pip_cmd} -r {req_file}".format(pip_cmd=pip_cmd, req_file=req_file))

@task
@timed
def install_node_prereqs():
    """
    Installs Node prerequisites
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    prereq_cache("Node prereqs", ["package.json"], node_prereqs_installation)

# To add a package to the uninstall list, just add it to this list! No need
# to touch any other part of this file.
PACKAGES_TO_UNINSTALL = [
    "South",                        # Because it interferes with Django 1.8 migrations.
    "edxval",                       # Because it was bork-installed somehow.
    "django-storages",
    "django-oauth2-provider",       # Because now it's called edx-django-oauth2-provider.
    "edx-oauth2-provider",          # Because it moved from github to pypi
    "i18n-tools",                   # Because now it's called edx-i18n-tools
]

@task
@timed
def uninstall_python_packages():
    """
    Uninstall Python packages that need explicit uninstallation.

    Some Python packages that we no longer want need to be explicitly
    uninstalled, notably, South.  Some other packages were once installed in
    ways that were resistant to being upgraded, like edxval.  Also uninstall
    them.

    """
    # So that we don't constantly uninstall things, use a hash of the packages
    # to be uninstalled.  Check it, and skip this if we're up to date.
    hasher = hashlib.sha1()
    hasher.update(repr(PACKAGES_TO_UNINSTALL))
    expected_version = hasher.hexdigest()
    state_file_path = os.path.join(PREREQS_STATE_DIR, "Python_uninstall.sha1")
    create_prereqs_cache_dir()

    if os.path.isfile(state_file_path):
        with open(state_file_path) as state_file:
            version = state_file.read()
        if version == expected_version:
            print 'Python uninstalls unchanged, skipping...'
            return

    # Run pip to find the packages we need to get rid of.  Believe it or not,
    # edx-val is installed in a way that it is present twice, so we have a loop
    # to really really get rid of it.
    for _ in range(3):
        uninstalled = False
        frozen = sh("pip freeze", capture=True)

        for package_name in PACKAGES_TO_UNINSTALL:
            if package_in_frozen(package_name, frozen):
                # Uninstall the pacakge
                sh("pip uninstall --disable-pip-version-check -y {}".format(package_name))
                uninstalled = True
        if not uninstalled:
            break
    else:
        # We tried three times and didn't manage to get rid of the pests.
        print "Couldn't uninstall unwanted Python packages!"
        return

    # Write our version.
    with open(state_file_path, "w") as state_file:
        state_file.write(expected_version)

def package_in_frozen(package_name, frozen_output):
    """Is this package in the output of 'pip freeze'?"""
    # Look for either:
    #
    #   PACKAGE-NAME==
    #
    # or:
    #
    #   blah_blah#egg=package_name-version
    #
    pattern = r"(?mi)^{pkg}==|#egg={pkg_under}-".format(
        pkg=re.escape(package_name),
        pkg_under=re.escape(package_name.replace("-", "_")),
    )
    return bool(re.search(pattern, frozen_output))

@task
@timed
def install_coverage_prereqs():
    """ Install python prereqs for measuring coverage. """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return
    pip_install_req_file(COVERAGE_REQ_FILE)

@task
@timed
def install_python_prereqs():
    """
    Installs Python prerequisites.
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    uninstall_python_packages()

    # Include all of the requirements files in the fingerprint.
    files_to_fingerprint = list(PYTHON_REQ_FILES)

    # Also fingerprint the directories where packages get installed:
    # ("/edx/app/edxapp/venvs/edxapp/lib/python2.7/site-packages")
    files_to_fingerprint.append(sysconfig.get_python_lib())

    # In a virtualenv, "-e installs" get put in a src directory.
    src_dir = os.path.join(sys.prefix, "src")
    if os.path.isdir(src_dir):
        files_to_fingerprint.append(src_dir)

    # Also fingerprint this source file, so that if the logic for installations
    # changes, we will redo the installation.
    this_file = __file__
    if this_file.endswith(".pyc"):
        this_file = this_file[:-1]      # use the .py file instead of the .pyc
    files_to_fingerprint.append(this_file)

    prereq_cache("Python prereqs", files_to_fingerprint, python_prereqs_installation)

@task
@timed
def install_prereqs():
    """
    Installs Node and Python prerequisites
    """
    if no_prereq_install():
        print NO_PREREQ_MESSAGE
        return

    install_node_prereqs()
    install_python_prereqs()
    log_installed_python_prereqs()

def log_installed_python_prereqs():
    """  Logs output of pip freeze for debugging. """
    sh("pip freeze > {}".format(Env.GEN_LOG_DIR + "/pip_freeze.log"))
    return
<|endoftext|>


<commit_before>import csv
import sqlalchemy as sa
import datetime
import pprint
import time
import os
import sqlparse

def load_csv_files_into_db(connection_string, data_dict, schema_ddl=None, indices_ddl=None, schema=None, delimiter=",",
                           lower_case_keys=True, i_print_update=10000, truncate=False, truncate_long_fields=True,
                           conditions=None, null_flag=False
                           ):

    db_engine = sa.create_engine(connection_string)
    db_connection = db_engine.connect()

    table_names = []
    if schema_ddl is not None:
        split_sql = sqlparse.split(schema_ddl)
        for sql_statement in split_sql:
            db_connection.execute(sql_statement)

    for key in data_dict:
        table_name = data_dict[key]
        if table_name not in table_names:
            if schema:
                table_name = schema + "." + table_name
            table_names += [table_name]

    if truncate:
        for table_name in table_names:
            truncate_sql = "truncate %s" % table_name
            db_connection.execute(truncate_sql)

    meta_data = sa.MetaData(db_connection, reflect=True, schema=schema)
    for data_file in data_dict:

        varchar_fields = {}
        table_name = data_dict[data_file]
        if schema:
            table_name = schema + "." + table_name

        table_obj = meta_data.tables[table_name]

        for column in table_obj.columns:
            column_name = column.name
            column_type = table_obj.c[column_name].type

            if "CHAR" in str(column_type):
                varchar_fields[column_name.lower()] = table_obj.c[column_name].type.length

        print("Loading %s" % table_name)

        db_transaction = db_connection.begin()
        try:
            with open(data_file, newline="", errors="replace") as f:
                dict_reader = csv.DictReader(f, delimiter=delimiter)
                start_time = time.time()
                elapsed_time = start_time
                i = 0
                for dict_row in dict_reader:
                    cleaned_dict = {}
                    for key in dict_row:

                        if len(dict_row[key]):

                            try:
                                if "date" in key or "DATE" in key:
                                    if "-" in dict_row[key]:
                                        if " " in dict_row[key]:
                                            cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key],
                                                                                "%Y-%m-%d %H:%M:%S")
                                        else:
                                            cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key], "%Y-%m-%d")
                                    else:
                                        cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key], "%Y%m%d")
                                else:
                                    cleaned_dict[key.lower()] = dict_row[key]
                            except ValueError:

                                cleaned_dict[key.lower()] = datetime.datetime(1900, 1, 1, 0, 0)

                    if lower_case_keys:
                        temp_cleaned_dict = {}
                        for key in cleaned_dict:
                            if truncate_long_fields:
                                if key.lower() in varchar_fields:
                                    field_length = varchar_fields[key.lower()]
                                    if len(cleaned_dict[key.lower()]) > field_length:
                                        print("Truncating: '%s'" % cleaned_dict[key.lower()])
                                        cleaned_dict[key.lower()] = cleaned_dict[key.lower()][0:field_length]
                            temp_cleaned_dict[key.lower()] = cleaned_dict[key]
                        cleaned_dict = temp_cleaned_dict

                    if conditions is None:
                        insert_data = True

                    else:
                        insert_data = False
                        for condition in conditions:
                            field_key = condition[0]
                            if field_key in cleaned_dict:
                                field_value = cleaned_dict[field_key]
                                if field_value in condition[1]:
                                    insert_data = True

                    for field_key in cleaned_dict:
                        field_value = cleaned_dict[field_key]
                        if field_value.__class__ == u"".__class__:
                            if field_value.upper() == "NULL" and null_flag:
                                cleaned_dict[field_key] = None

                    if insert_data:
                        s = table_obj.insert(cleaned_dict)
                        try:
                            db_connection.execute(s)
                        except:
                            pprint.pprint(cleaned_dict)
                            raise

                        if i > 0 and i % i_print_update == 0:
                            current_time = time.time()
                            time_difference = current_time - elapsed_time
                            print("Loaded %s total rows at %s seconds per %s rows" % (i, time_difference, i_print_update))
                            elapsed_time = time.time()
                        i += 1

                db_transaction.commit()
                current_time = time.time()
                total_time_difference = current_time - start_time
                print("Loaded %s total row in %s seconds" % (i, total_time_difference))

        except:
            db_transaction.rollback()
            raise

    if indices_ddl is not None:
        split_sql = sqlparse.split(indices_ddl)
        for sql_statement in split_sql:
            try:
                db_connection.execute(sql_statement)
            except(sa.exc.OperationalError):
                print("Skipping: '%s'" % sql_statement)

def generate_db_dict(output_directory=None, load_pairs=None):
    if load_pairs is None:
        load_pairs = [
            ("condition_occurrence", "condition_occurrence_dx_cdm.csv"),
            ("location", "location_cdm.csv"),
            ("person", "person_cdm.csv"),
            ("visit_occurrence", "visit_occurrence_cdm.csv"),
            ("visit_detail", "visit_detail_cdm.csv"),
            ("procedure_occurrence", "procedure_cdm.csv"),
            ("procedure_occurrence", "procedure_dx_cdm.csv"),
            ("measurement", "measurement_encounter_cdm.csv"),
            ("measurement", "measurement_dx_cdm.csv"),
            ("measurement", "measurement_proc_cdm.csv"),
            ("drug_exposure", "drug_exposure_cdm.csv"),
            ("drug_exposure", "drug_exposure_proc_cdm.csv"),
            ("death", "death_cdm.csv"),
            ("observation", "observation_dx_cdm.csv"),
            ("observation", "observation_measurement_encounter_cdm.csv"),
            ("observation", "observation_proc_cdm.csv"),
            ("observation_period", "observation_period_cdm.csv"),
            ("care_site", "care_site_cdm.csv"),
            ("payer_plan_period", "payer_plan_period_cdm.csv")
        ]

    data_dict = {}
    for pair in load_pairs:
        if output_directory is not None:
            data_dict[os.path.join(output_directory, pair[1])] = pair[0]
        else:
            data_dict[pair[1]] = pair[0]

    return data_dict

def generate_vocabulary_load(vocabulary_directory,  vocabularies=["CONCEPT",
                    "CONCEPT_ANCESTOR",
                    "CONCEPT_CLASS",
                    "CONCEPT_RELATIONSHIP",
                    "CONCEPT_SYNONYM",
                    "DOMAIN",
                    "DRUG_STRENGTH",
                    "RELATIONSHIP",
                    "VOCABULARY"]):

    load_pairs = []
    for vocabulary in vocabularies:
        load_pairs += [(vocabulary.lower(), os.path.join(vocabulary_directory, vocabulary + ".csv"))]

    return load_pairs

def generate_observation_period(encounter_csv_file_name, source_period_observation_csv_file_name,
                                id_field_name, start_date_field_name, end_date_field_name):

    with open(encounter_csv_file_name, newline="") as f:
        dict_reader = csv.DictReader(f)
        observation_period_dict = {}

        for row_dict in dict_reader:

            start_date_value = row_dict[start_date_field_name]
            end_date_value = row_dict[end_date_field_name]

            if len(end_date_value) == 0:
                end_date_value = start_date_value

            id_value = row_dict[id_field_name]

            if id_value in observation_period_dict:
                past_start_date_value, past_end_date_value = observation_period_dict[id_value]

                if start_date_value < past_start_date_value:
                    set_start_date_value = start_date_value
                else:
                    set_start_date_value = past_start_date_value

                if end_date_value > past_end_date_value:
                    set_end_date_value = end_date_value
                else:
                    set_end_date_value = past_end_date_value

                observation_period_dict[id_value] = (set_start_date_value, set_end_date_value)

            else:
                observation_period_dict[id_value] = (start_date_value, end_date_value)

    with open(source_period_observation_csv_file_name, "w", newline="") as fw:
        csv_writer = csv.writer(fw)

        csv_writer.writerow([id_field_name, start_date_field_name, end_date_field_name])

        for id_value in observation_period_dict:
            start_date_value, end_date_value = observation_period_dict[id_value]
            if start_date_value == "":
                start_date_value = end_date_value
            row_to_write = [id_value, start_date_value, end_date_value]
            csv_writer.writerow(row_to_write)<commit_msg>encoding error<commit_after>import csv
import sqlalchemy as sa
import datetime
import pprint
import time
import os
import sqlparse

def load_csv_files_into_db(connection_string, data_dict, schema_ddl=None, indices_ddl=None, schema=None, delimiter=",",
                           lower_case_keys=True, i_print_update=10000, truncate=False, truncate_long_fields=True,
                           conditions=None, null_flag=False
                           ):

    db_engine = sa.create_engine(connection_string)
    db_connection = db_engine.connect()

    table_names = []
    if schema_ddl is not None:
        split_sql = sqlparse.split(schema_ddl)
        for sql_statement in split_sql:
            db_connection.execute(sql_statement)

    for key in data_dict:
        table_name = data_dict[key]
        if table_name not in table_names:
            if schema:
                table_name = schema + "." + table_name
            table_names += [table_name]

    if truncate:
        for table_name in table_names:
            truncate_sql = "truncate %s" % table_name
            db_connection.execute(truncate_sql)

    meta_data = sa.MetaData(db_connection, reflect=True, schema=schema)
    for data_file in data_dict:

        varchar_fields = {}
        table_name = data_dict[data_file]
        if schema:
            table_name = schema + "." + table_name

        table_obj = meta_data.tables[table_name]

        for column in table_obj.columns:
            column_name = column.name
            column_type = table_obj.c[column_name].type

            if "CHAR" in str(column_type):
                varchar_fields[column_name.lower()] = table_obj.c[column_name].type.length

        print("Loading %s" % table_name)

        db_transaction = db_connection.begin()
        try:
            with open(data_file, newline="", errors="replace") as f:
                dict_reader = csv.DictReader(f, delimiter=delimiter)
                start_time = time.time()
                elapsed_time = start_time
                i = 0
                for dict_row in dict_reader:
                    cleaned_dict = {}
                    for key in dict_row:

                        if len(dict_row[key]):

                            try:
                                if "date" in key or "DATE" in key:
                                    if "-" in dict_row[key]:
                                        if " " in dict_row[key]:
                                            cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key],
                                                                                "%Y-%m-%d %H:%M:%S")
                                        else:
                                            cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key], "%Y-%m-%d")
                                    else:
                                        cleaned_dict[key.lower()] = datetime.datetime.strptime(dict_row[key], "%Y%m%d")
                                else:
                                    cleaned_dict[key.lower()] = dict_row[key]
                            except ValueError:

                                cleaned_dict[key.lower()] = datetime.datetime(1900, 1, 1, 0, 0)

                    if lower_case_keys:
                        temp_cleaned_dict = {}
                        for key in cleaned_dict:
                            if truncate_long_fields:
                                if key.lower() in varchar_fields:
                                    field_length = varchar_fields[key.lower()]
                                    if len(cleaned_dict[key.lower()]) > field_length:
                                        print("Truncating: '%s'" % cleaned_dict[key.lower()])
                                        cleaned_dict[key.lower()] = cleaned_dict[key.lower()][0:field_length]
                            temp_cleaned_dict[key.lower()] = cleaned_dict[key]
                        cleaned_dict = temp_cleaned_dict

                    if conditions is None:
                        insert_data = True

                    else:
                        insert_data = False
                        for condition in conditions:
                            field_key = condition[0]
                            if field_key in cleaned_dict:
                                field_value = cleaned_dict[field_key]
                                if field_value in condition[1]:
                                    insert_data = True

                    for field_key in cleaned_dict:
                        field_value = cleaned_dict[field_key]
                        if field_value.__class__ == u"".__class__:
                            if field_value.upper() == "NULL" and null_flag:
                                cleaned_dict[field_key] = None

                    if insert_data:
                        s = table_obj.insert(cleaned_dict)
                        try:
                            db_connection.execute(s)
                        except:
                            pprint.pprint(cleaned_dict)
                            raise

                        if i > 0 and i % i_print_update == 0:
                            current_time = time.time()
                            time_difference = current_time - elapsed_time
                            print("Loaded %s total rows at %s seconds per %s rows" % (i, time_difference, i_print_update))
                            elapsed_time = time.time()
                        i += 1

                db_transaction.commit()
                current_time = time.time()
                total_time_difference = current_time - start_time
                print("Loaded %s total row in %s seconds" % (i, total_time_difference))

        except:
            db_transaction.rollback()
            raise

    if indices_ddl is not None:
        split_sql = sqlparse.split(indices_ddl)
        for sql_statement in split_sql:
            try:
                db_connection.execute(sql_statement)
            except(sa.exc.OperationalError):
                print("Skipping: '%s'" % sql_statement)

def generate_db_dict(output_directory=None, load_pairs=None):
    if load_pairs is None:
        load_pairs = [
            ("condition_occurrence", "condition_occurrence_dx_cdm.csv"),
            ("location", "location_cdm.csv"),
            ("person", "person_cdm.csv"),
            ("visit_occurrence", "visit_occurrence_cdm.csv"),
            ("visit_detail", "visit_detail_cdm.csv"),
            ("procedure_occurrence", "procedure_cdm.csv"),
            ("procedure_occurrence", "procedure_dx_cdm.csv"),
            ("measurement", "measurement_encounter_cdm.csv"),
            ("measurement", "measurement_dx_cdm.csv"),
            ("measurement", "measurement_proc_cdm.csv"),
            ("drug_exposure", "drug_exposure_cdm.csv"),
            ("drug_exposure", "drug_exposure_proc_cdm.csv"),
            ("death", "death_cdm.csv"),
            ("observation", "observation_dx_cdm.csv"),
            ("observation", "observation_measurement_encounter_cdm.csv"),
            ("observation", "observation_proc_cdm.csv"),
            ("observation_period", "observation_period_cdm.csv"),
            ("care_site", "care_site_cdm.csv"),
            ("payer_plan_period", "payer_plan_period_cdm.csv")
        ]

    data_dict = {}
    for pair in load_pairs:
        if output_directory is not None:
            data_dict[os.path.join(output_directory, pair[1])] = pair[0]
        else:
            data_dict[pair[1]] = pair[0]

    return data_dict

def generate_vocabulary_load(vocabulary_directory,  vocabularies=["CONCEPT",
                    "CONCEPT_ANCESTOR",
                    "CONCEPT_CLASS",
                    "CONCEPT_RELATIONSHIP",
                    "CONCEPT_SYNONYM",
                    "DOMAIN",
                    "DRUG_STRENGTH",
                    "RELATIONSHIP",
                    "VOCABULARY"]):

    load_pairs = []
    for vocabulary in vocabularies:
        load_pairs += [(vocabulary.lower(), os.path.join(vocabulary_directory, vocabulary + ".csv"))]

    return load_pairs

def generate_observation_period(encounter_csv_file_name, source_period_observation_csv_file_name,
                                id_field_name, start_date_field_name, end_date_field_name):

    with open(encounter_csv_file_name, newline="", encoding="utf8", errors="replace") as f:
        dict_reader = csv.DictReader(f)
        observation_period_dict = {}

        for row_dict in dict_reader:

            start_date_value = row_dict[start_date_field_name]
            end_date_value = row_dict[end_date_field_name]

            if len(end_date_value) == 0:
                end_date_value = start_date_value

            id_value = row_dict[id_field_name]

            if id_value in observation_period_dict:
                past_start_date_value, past_end_date_value = observation_period_dict[id_value]

                if start_date_value < past_start_date_value:
                    set_start_date_value = start_date_value
                else:
                    set_start_date_value = past_start_date_value

                if end_date_value > past_end_date_value:
                    set_end_date_value = end_date_value
                else:
                    set_end_date_value = past_end_date_value

                observation_period_dict[id_value] = (set_start_date_value, set_end_date_value)

            else:
                observation_period_dict[id_value] = (start_date_value, end_date_value)

    with open(source_period_observation_csv_file_name, "w", newline="") as fw:
        csv_writer = csv.writer(fw)

        csv_writer.writerow([id_field_name, start_date_field_name, end_date_field_name])

        for id_value in observation_period_dict:
            start_date_value, end_date_value = observation_period_dict[id_value]
            if start_date_value == "":
                start_date_value = end_date_value
            row_to_write = [id_value, start_date_value, end_date_value]
            csv_writer.writerow(row_to_write)<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 workflow using BigQuery sources and sinks.

The workflow will read from a table that has the 'month' and 'tornado' fields as
part of the table schema (other additional fields are ignored). The 'month'
field is a number represented as a string (e.g., '23') and the 'tornado' field
is a boolean field.

The workflow will compute the number of tornadoes in each month and output
the results to a table (created if needed) with the following schema:

- month: number
- tornado_count: number

This example uses the default behavior for BigQuery source and sinks that
represents table rows as plain Python dictionaries.
"""

from __future__ import absolute_import

import argparse
import logging

import apache_beam as beam

def count_tornadoes(input_data):
  """Workflow computing the number of tornadoes for each month that had one.

  Args:
    input_data: a PCollection of dictionaries representing table rows. Each
      dictionary will have a 'month' and a 'tornado' key as described in the
      module comment.

  Returns:
    A PCollection of dictionaries containing 'month' and 'tornado_count' keys.
    Months without tornadoes are skipped.
  """

  return (input_data
          | 'months with tornatoes' >> beam.FlatMap(
              lambda row: [(int(row['month']), 1)] if row['tornado'] else [])
          | 'monthly count' >> beam.CombinePerKey(sum)
          | 'format' >> beam.Map(
              lambda (k, v): {'month': k, 'tornado_count': v}))

def run(argv=None):
  parser = argparse.ArgumentParser()
  parser.add_argument('--input',
                      default='clouddataflow-readonly:samples.weather_stations',
                      help=('Input BigQuery table to process specified as: '
                            'PROJECT:DATASET.TABLE or DATASET.TABLE.'))
  parser.add_argument(
      '--output',
      required=True,
      help=
      ('Output BigQuery table for results specified as: PROJECT:DATASET.TABLE '
       'or DATASET.TABLE.'))
  known_args, pipeline_args = parser.parse_known_args(argv)

  p = beam.Pipeline(argv=pipeline_args)

  # Read the table rows into a PCollection.
  rows = p | 'read' >> beam.io.Read(beam.io.BigQuerySource(known_args.input))
  counts = count_tornadoes(rows)

  # Write the output using a "Write" transform that has side effects.
  # pylint: disable=expression-not-assigned
  counts | 'write' >> beam.io.Write(
      beam.io.BigQuerySink(
          known_args.output,
          schema='month:INTEGER, tornado_count:INTEGER',
          create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
          write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE))

  # Run the pipeline (all operations are deferred until run() is called).
  p.run().wait_until_finish()

if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  run()
<commit_msg>This closes #2390<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 workflow using BigQuery sources and sinks.

The workflow will read from a table that has the 'month' and 'tornado' fields as
part of the table schema (other additional fields are ignored). The 'month'
field is a number represented as a string (e.g., '23') and the 'tornado' field
is a boolean field.

The workflow will compute the number of tornadoes in each month and output
the results to a table (created if needed) with the following schema:

- month: number
- tornado_count: number

This example uses the default behavior for BigQuery source and sinks that
represents table rows as plain Python dictionaries.
"""

from __future__ import absolute_import

import argparse
import logging

import apache_beam as beam

def count_tornadoes(input_data):
  """Workflow computing the number of tornadoes for each month that had one.

  Args:
    input_data: a PCollection of dictionaries representing table rows. Each
      dictionary will have a 'month' and a 'tornado' key as described in the
      module comment.

  Returns:
    A PCollection of dictionaries containing 'month' and 'tornado_count' keys.
    Months without tornadoes are skipped.
  """

  return (input_data
          | 'months with tornadoes' >> beam.FlatMap(
              lambda row: [(int(row['month']), 1)] if row['tornado'] else [])
          | 'monthly count' >> beam.CombinePerKey(sum)
          | 'format' >> beam.Map(
              lambda (k, v): {'month': k, 'tornado_count': v}))

def run(argv=None):
  parser = argparse.ArgumentParser()
  parser.add_argument('--input',
                      default='clouddataflow-readonly:samples.weather_stations',
                      help=('Input BigQuery table to process specified as: '
                            'PROJECT:DATASET.TABLE or DATASET.TABLE.'))
  parser.add_argument(
      '--output',
      required=True,
      help=
      ('Output BigQuery table for results specified as: PROJECT:DATASET.TABLE '
       'or DATASET.TABLE.'))
  known_args, pipeline_args = parser.parse_known_args(argv)

  p = beam.Pipeline(argv=pipeline_args)

  # Read the table rows into a PCollection.
  rows = p | 'read' >> beam.io.Read(beam.io.BigQuerySource(known_args.input))
  counts = count_tornadoes(rows)

  # Write the output using a "Write" transform that has side effects.
  # pylint: disable=expression-not-assigned
  counts | 'write' >> beam.io.Write(
      beam.io.BigQuerySink(
          known_args.output,
          schema='month:INTEGER, tornado_count:INTEGER',
          create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
          write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE))

  # Run the pipeline (all operations are deferred until run() is called).
  p.run().wait_until_finish()

if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  run()
<|endoftext|>


<commit_before>#
# Copyright 2012 Quantopian, Inc.
#
# 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.

from numbers import Number
from collections import defaultdict
from math import sqrt

from zipline.transforms.utils import EventWindow, TransformMeta

class MovingStandardDev(object):
    """
    Class that maintains a dicitonary from sids to
    MovingStandardDevWindows.  For each sid, we maintain a the
    standard deviation of all events falling within the specified
    window.
    """
    __metaclass__ = TransformMeta

    def __init__(self, market_aware=True, window_length=None, delta=None):

        self.market_aware = market_aware

        self.delta = delta
        self.window_length = window_length

        # Market-aware mode only works with full-day windows.
        if self.market_aware:
            assert self.window_length and not self.delta,\
                "Market-aware mode only works with full-day windows."

        # Non-market-aware mode requires a timedelta.
        else:
            assert self.delta and not self.window_length, \
                "Non-market-aware mode requires a timedelta."

        # No way to pass arguments to the defaultdict factory, so we
        # need to define a method to generate the correct EventWindows.
        self.sid_windows = defaultdict(self.create_window)

    def create_window(self):
        """
        Factory method for self.sid_windows.
        """
        return MovingStandardDevWindow(
            self.market_aware,
            self.window_length,
            self.delta
        )

    def update(self, event):
        """
        Update the event window for this event's sid.  Return an ndict
        from tracked fields to moving averages.
        """
        # This will create a new EventWindow if this is the first
        # message for this sid.
        window = self.sid_windows[event.sid]
        window.update(event)
        return window.get_stddev()

class MovingStandardDevWindow(EventWindow):
    """
    Iteratively calculates standard deviation for a particular sid
    over a given time window.  The expected functionality of this
    class is to be instantiated inside a MovingStandardDev.
    """

    def __init__(self, market_aware, days, delta):
        # Call the superclass constructor to set up base EventWindow
        # infrastructure.
        EventWindow.__init__(self, market_aware, days, delta)

        self.sum = 0.0
        self.sum_sqr = 0.0

    def handle_add(self, event):
        assert isinstance(event.price, Number)

        self.sum += event.price
        self.sum_sqr += event.price ** 2

    def handle_remove(self, event):
        assert isinstance(event.price, Number)

        self.sum -= event.price
        self.sum_sqr -= event.price ** 2

    def get_stddev(self):
        # Sample standard deviation is undefined for a single event or
        # no events.
        if len(self) <= 1:
            return None

        else:
            average = self.sum / len(self)
            s_squared = (self.sum_sqr - self.sum * average) \
                / (len(self) - 1)
            stddev = sqrt(s_squared)
        return stddev
<commit_msg>Update zipline/transforms/stddev.py<commit_after>#
# Copyright 2012 Quantopian, Inc.
#
# 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.

from numbers import Number
from collections import defaultdict
from math import sqrt

from zipline.transforms.utils import EventWindow, TransformMeta

class MovingStandardDev(object):
    """
    Class that maintains a dictionary from sids to
    MovingStandardDevWindows.  For each sid, we maintain a the
    standard deviation of all events falling within the specified
    window.
    """
    __metaclass__ = TransformMeta

    def __init__(self, market_aware=True, window_length=None, delta=None):

        self.market_aware = market_aware

        self.delta = delta
        self.window_length = window_length

        # Market-aware mode only works with full-day windows.
        if self.market_aware:
            assert self.window_length and not self.delta,\
                "Market-aware mode only works with full-day windows."

        # Non-market-aware mode requires a timedelta.
        else:
            assert self.delta and not self.window_length, \
                "Non-market-aware mode requires a timedelta."

        # No way to pass arguments to the defaultdict factory, so we
        # need to define a method to generate the correct EventWindows.
        self.sid_windows = defaultdict(self.create_window)

    def create_window(self):
        """
        Factory method for self.sid_windows.
        """
        return MovingStandardDevWindow(
            self.market_aware,
            self.window_length,
            self.delta
        )

    def update(self, event):
        """
        Update the event window for this event's sid.  Return an ndict
        from tracked fields to moving averages.
        """
        # This will create a new EventWindow if this is the first
        # message for this sid.
        window = self.sid_windows[event.sid]
        window.update(event)
        return window.get_stddev()

class MovingStandardDevWindow(EventWindow):
    """
    Iteratively calculates standard deviation for a particular sid
    over a given time window.  The expected functionality of this
    class is to be instantiated inside a MovingStandardDev.
    """

    def __init__(self, market_aware, days, delta):
        # Call the superclass constructor to set up base EventWindow
        # infrastructure.
        EventWindow.__init__(self, market_aware, days, delta)

        self.sum = 0.0
        self.sum_sqr = 0.0

    def handle_add(self, event):
        assert isinstance(event.price, Number)

        self.sum += event.price
        self.sum_sqr += event.price ** 2

    def handle_remove(self, event):
        assert isinstance(event.price, Number)

        self.sum -= event.price
        self.sum_sqr -= event.price ** 2

    def get_stddev(self):
        # Sample standard deviation is undefined for a single event or
        # no events.
        if len(self) <= 1:
            return None

        else:
            average = self.sum / len(self)
            s_squared = (self.sum_sqr - self.sum * average) \
                / (len(self) - 1)
            stddev = sqrt(s_squared)
        return stddev
<|endoftext|>