<commit_before>#!/usr/bin/env python
#
# 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.

import argparse
import glob
import os
import pickle
import sys

from lxml import etree
from oslo.config import cfg

from autohelp import OptionsCache  # noqa

# Swift configuration example files live in
# swift/etc/*.conf-sample
# and contain sections enclosed in [], with
# options one per line containing =
# and generally only having a single entry
# after the equals (the default value)

DBK_NS = ".//{http://docbook.org/ns/docbook}"

BASE_XML = '''<?xml version="1.0"?>
<para xmlns="http://docbook.org/ns/docbook"
  version="5.0">

Description of configuration options for
        <literal>[%s]</literal> in <filename>%s.conf</filename>

    <col width="50%%"/>
    <col width="50%%"/>

        Configuration option = Default value
        Description

</para>'''

def parse_line(line):
    """Parse a line.

    Takes a line from a swift sample configuration file and attempts
    to separate the lines with actual configuration option and default
    value from the rest. Returns None if the line doesn't appear to
    contain a valid configuration option = default value pair, and
    a pair of the config and its default if it does.
    """

    if '=' not in line:
        return None
    temp_line = line.strip('#').strip()
    config, default = temp_line.split('=', 1)
    config = config.strip()
    if ' ' in config and config[0:3] != 'set':
        if len(default.split()) > 1 or config[0].isupper():
            return None
    if len(config) < 2 or '.' in config or '<' in config or '>' in config:
        return None
    return config, default.strip()

def get_existing_options(optfiles):
    """Parse an existing XML table to compile a list of existing options."""
    options = {}
    for optfile in optfiles:
        if optfile.endswith('/swift-conf-changes.xml'):
            continue
        xml = etree.fromstring(open(optfile).read())
        tbody = xml.find(DBK_NS + "tbody")
        trlist = tbody.findall(DBK_NS + "tr")
        for tr in trlist:
            try:
                col1, col2 = tr.findall(DBK_NS + "td")
                option = col1.find(DBK_NS + "option").text
                helptext = col2.text
            except IndexError:
                continue
            if option not in options or 'No help text' in options[option]:
                # options[option.split('=',1)[0]] = helptext
                options[option] = helptext
    return options

def extract_descriptions_from_devref(swift_repo, options):
    """Extract descriptions from devref RST files.

    Loop through the devref RST files, looking for lines formatted
    such that they might contain a description of a particular
    option.
    """

    option_descs = {}
    rsts = glob.glob(swift_repo + '/doc/source/*.rst')
    for rst in rsts:
        rst_file = open(rst, 'r')
        in_option_block = False
        prev_option = None
        for line in rst_file:
            if 'Option    ' in line:
                in_option_block = True
            if in_option_block:
                if '========' in line:
                    in_option_block = False
                    continue
                if line[0] == ' ' and prev_option is not None:
                    option_descs[prev_option] = (option_descs[prev_option]
                                                 + ' ' + line.strip())
                for option in options:
                    line_parts = line.strip().split(None, 2)
                    if ('   ' in line and
                            len(line_parts) == 3 and
                            option == line_parts[0] and
                            line_parts[1] != '=' and
                            option != 'use' and
                            (option not in option_descs or
                             len(option_descs[option]) < len(line_parts[2]))):
                        option_descs[option] = line_parts[2]
                        prev_option = option
    return option_descs

def write_xml(manuals_repo, section, xml):
    """Write the XML to file."""
    sample, section_name = section.split('|')
    section_filename = (manuals_repo + '/doc/common/tables/' +
                        'swift-' + sample + '-' + section_name + '.xml')
    with open(section_filename, 'w') as fd:
        fd.write(etree.tostring(xml, pretty_print=True,
                                xml_declaration=True,
                                encoding="UTF-8"))

def new_section_xml(manuals_repo, section):
    """Create a new XML tree."""

    # The section holds 2 informations, the file in which the option was found,
    # and the section name in that file.
    sample, section_name = section.split('|')
    parser = etree.XMLParser(remove_blank_text=True)
    xml = etree.XML(BASE_XML % (section_name, sample), parser)
    return xml

def write_docbook(options, manuals_repo):
    """Create new DocBook tables.

    Writes a set of DocBook-formatted tables, one per section in swift
    configuration files.
    """
    names = options.get_option_names()
    current_section = None
    xml = None
    for full_option in sorted(names, OptionsCache._cmpopts):
        section, optname = full_option.split('/')

        if current_section != section:
            if xml is not None:
                write_xml(manuals_repo, current_section, xml)
            current_section = section
            xml = new_section_xml(manuals_repo, section)
            tbody = xml.find(DBK_NS + "tbody")

        oslo_opt = options.get_option(full_option)[1]

        tr = etree.Element('tr')
        tbody.append(tr)

        td = etree.Element('td')
        option_xml = etree.SubElement(td, 'option')
        option_xml.text = "%s" % oslo_opt.name
        option_xml.tail = " = "
        replaceable_xml = etree.SubElement(td, 'replaceable')
        replaceable_xml.text = "%s" % oslo_opt.default
        tr.append(td)

        td = etree.Element('td')
        td.text = oslo_opt.help
        tr.append(td)

    write_xml(manuals_repo, section, xml)

def read_options(swift_repo, manuals_repo, verbose):
    """Find swift configuration options.

    Uses existing tables and swift devref as a source of truth in that order to
    determine helptext for options found in sample config files.
    """

    existing_tables = glob.glob(manuals_repo + '/doc/common/tables/swift*xml')
    options = {}
    # use the existing tables to get a list of option names
    options = get_existing_options(existing_tables)
    option_descs = extract_descriptions_from_devref(swift_repo, options)
    conf_samples = glob.glob(swift_repo + '/etc/*conf-sample')
    for sample in conf_samples:
        current_section = None
        sample_file = open(sample, 'r')
        for line in sample_file:
            if '[' in line and ']\n' in line and '=' not in line:
                # It's a header line in the conf file, open a new table file
                # for this section and close any existing one
                new_line = line.strip('#').strip()
                if current_section != new_line:
                    current_section = new_line

                    base_section = os.path.basename(sample).split('.conf')[0]
                    extra_section = current_section[1:-1].replace(':', '-')
                    full_section = "%s|%s" % (base_section, extra_section)

                    continue

            # All the swift files start with a section, except the rsync
            # sample. The first items are not important for us.
            if current_section is None:
                continue

            # It's a config option line in the conf file, find out the
            # help text and write to the table file.
            parsed_line = parse_line(line)
            if parsed_line is not None:
                if (parsed_line[0] in options.keys()
                   and 'No help text' not in options[parsed_line[0]]):
                    # use the help text from existing tables
                    option_desc = options[parsed_line[0]]
                elif parsed_line[0] in option_descs:
                    # use the help text from the devref
                    option_desc = option_descs[parsed_line[0]]
                else:
                    option_desc = 'No help text available for this option.'
                    if verbose > 0:
                        print(parsed_line[0] + " has no help text")

                # \xa0 is a non-breacking space
                name = parsed_line[0]
                option_desc = option_desc.replace(u'\xa0', u' ')
                default = parsed_line[1]

                o = cfg.StrOpt(name=name, default=default, help=option_desc)
                try:
                    cfg.CONF.register_opt(o, full_section)
                except cfg.DuplicateOptError:
                    pass

def dump_options(options):
    """Dump the list of options with their attributes.

    This output is consumed by the diff_branches script.
    """
    print(pickle.dumps(options._opts_by_name))

def main():
    """Parse and write the Swift configuration options."""

    parser = argparse.ArgumentParser(
        description="Update the swift options tables.",
        usage="%(prog)s docbook|dump [-v] [-s swift_repo] [-m manuals_repo]")
    parser.add_argument('subcommand',
                        help='Action (docbook, dump).',
                        choices=['docbook', 'dump'])
    parser.add_argument('-s', '--swift-repo',
                        dest='swift_repo',
                        help="Location of the swift git repository.",
                        required=False,
                        default="./sources/swift")
    parser.add_argument('-m', '--manuals-repo',
                        dest='manuals_repo',
                        help="Location of the manuals git repository.",
                        required=False,
                        default="./sources/openstack-manuals")
    parser.add_argument('-v', '--verbose',
                        action='count',
                        default=0,
                        dest='verbose',
                        required=False)
    args = parser.parse_args()

    # Avoid cluttering the pickle output, otherwise it's not usable
    if args.subcommand == 'dump':
        args.verbose = 0

    read_options(args.swift_repo, args.manuals_repo, args.verbose)
    options = OptionsCache()

    if args.subcommand == 'docbook':
        write_docbook(options, args.manuals_repo)

    elif args.subcommand == 'dump':
        options.dump()

    return 0

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>swift flag: handle XML in help strings<commit_after>#!/usr/bin/env python
#
# 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.

import argparse
import glob
import os
import pickle
import sys

from lxml import etree
from oslo.config import cfg

from autohelp import OptionsCache  # noqa

# Swift configuration example files live in
# swift/etc/*.conf-sample
# and contain sections enclosed in [], with
# options one per line containing =
# and generally only having a single entry
# after the equals (the default value)

DBK_NS = ".//{http://docbook.org/ns/docbook}"

BASE_XML = '''<?xml version="1.0"?>
<para xmlns="http://docbook.org/ns/docbook"
  version="5.0">

Description of configuration options for
        <literal>[%s]</literal> in <filename>%s.conf</filename>

    <col width="50%%"/>
    <col width="50%%"/>

        Configuration option = Default value
        Description

</para>'''

def parse_line(line):
    """Parse a line.

    Takes a line from a swift sample configuration file and attempts
    to separate the lines with actual configuration option and default
    value from the rest. Returns None if the line doesn't appear to
    contain a valid configuration option = default value pair, and
    a pair of the config and its default if it does.
    """

    if '=' not in line:
        return None
    temp_line = line.strip('#').strip()
    config, default = temp_line.split('=', 1)
    config = config.strip()
    if ' ' in config and config[0:3] != 'set':
        if len(default.split()) > 1 or config[0].isupper():
            return None
    if len(config) < 2 or '.' in config or '<' in config or '>' in config:
        return None
    return config, default.strip()

def get_existing_options(optfiles):
    """Parse an existing XML table to compile a list of existing options."""
    options = {}
    for optfile in optfiles:
        if optfile.endswith('/swift-conf-changes.xml'):
            continue
        xml = etree.fromstring(open(optfile).read())
        tbody = xml.find(DBK_NS + "tbody")
        trlist = tbody.findall(DBK_NS + "tr")
        for tr in trlist:
            try:
                col1, col2 = tr.findall(DBK_NS + "td")
                option = col1.find(DBK_NS + "option").text
                helptext = etree.tostring(col2, xml_declaration=False,
                                          method="text")
            except IndexError:
                continue
            if option not in options or 'No help text' in options[option]:
                # options[option.split('=',1)[0]] = helptext
                options[option] = helptext
    return options

def extract_descriptions_from_devref(swift_repo, options):
    """Extract descriptions from devref RST files.

    Loop through the devref RST files, looking for lines formatted
    such that they might contain a description of a particular
    option.
    """

    option_descs = {}
    rsts = glob.glob(swift_repo + '/doc/source/*.rst')
    for rst in rsts:
        rst_file = open(rst, 'r')
        in_option_block = False
        prev_option = None
        for line in rst_file:
            if 'Option    ' in line:
                in_option_block = True
            if in_option_block:
                if '========' in line:
                    in_option_block = False
                    continue
                if line[0] == ' ' and prev_option is not None:
                    option_descs[prev_option] = (option_descs[prev_option]
                                                 + ' ' + line.strip())
                for option in options:
                    line_parts = line.strip().split(None, 2)
                    if ('   ' in line and
                            len(line_parts) == 3 and
                            option == line_parts[0] and
                            line_parts[1] != '=' and
                            option != 'use' and
                            (option not in option_descs or
                             len(option_descs[option]) < len(line_parts[2]))):
                        option_descs[option] = line_parts[2]
                        prev_option = option
    return option_descs

def write_xml(manuals_repo, section, xml):
    """Write the XML to file."""
    sample, section_name = section.split('|')
    section_filename = (manuals_repo + '/doc/common/tables/' +
                        'swift-' + sample + '-' + section_name + '.xml')
    with open(section_filename, 'w') as fd:
        fd.write(etree.tostring(xml, pretty_print=True,
                                xml_declaration=True,
                                encoding="UTF-8"))

def new_section_xml(manuals_repo, section):
    """Create a new XML tree."""

    # The section holds 2 informations, the file in which the option was found,
    # and the section name in that file.
    sample, section_name = section.split('|')
    parser = etree.XMLParser(remove_blank_text=True)
    xml = etree.XML(BASE_XML % (section_name, sample), parser)
    return xml

def write_docbook(options, manuals_repo):
    """Create new DocBook tables.

    Writes a set of DocBook-formatted tables, one per section in swift
    configuration files.
    """
    names = options.get_option_names()
    current_section = None
    xml = None
    for full_option in sorted(names, OptionsCache._cmpopts):
        section, optname = full_option.split('/')

        if current_section != section:
            if xml is not None:
                write_xml(manuals_repo, current_section, xml)
            current_section = section
            xml = new_section_xml(manuals_repo, section)
            tbody = xml.find(DBK_NS + "tbody")

        oslo_opt = options.get_option(full_option)[1]

        tr = etree.Element('tr')
        tbody.append(tr)

        td = etree.Element('td')
        option_xml = etree.SubElement(td, 'option')
        option_xml.text = "%s" % oslo_opt.name
        option_xml.tail = " = "
        replaceable_xml = etree.SubElement(td, 'replaceable')
        replaceable_xml.text = "%s" % oslo_opt.default
        tr.append(td)

        td = etree.Element('td')
        td.text = oslo_opt.help
        tr.append(td)

    write_xml(manuals_repo, section, xml)

def read_options(swift_repo, manuals_repo, verbose):
    """Find swift configuration options.

    Uses existing tables and swift devref as a source of truth in that order to
    determine helptext for options found in sample config files.
    """

    existing_tables = glob.glob(manuals_repo + '/doc/common/tables/swift*xml')
    options = {}
    # use the existing tables to get a list of option names
    options = get_existing_options(existing_tables)
    option_descs = extract_descriptions_from_devref(swift_repo, options)
    conf_samples = glob.glob(swift_repo + '/etc/*conf-sample')
    for sample in conf_samples:
        current_section = None
        sample_file = open(sample, 'r')
        for line in sample_file:
            if '[' in line and ']\n' in line and '=' not in line:
                # It's a header line in the conf file, open a new table file
                # for this section and close any existing one
                new_line = line.strip('#').strip()
                if current_section != new_line:
                    current_section = new_line

                    base_section = os.path.basename(sample).split('.conf')[0]
                    extra_section = current_section[1:-1].replace(':', '-')
                    full_section = "%s|%s" % (base_section, extra_section)

                    continue

            # All the swift files start with a section, except the rsync
            # sample. The first items are not important for us.
            if current_section is None:
                continue

            # It's a config option line in the conf file, find out the
            # help text and write to the table file.
            parsed_line = parse_line(line)
            if parsed_line is not None:
                if (parsed_line[0] in options.keys()
                   and 'No help text' not in options[parsed_line[0]]):
                    # use the help text from existing tables
                    option_desc = options[parsed_line[0]]
                elif parsed_line[0] in option_descs:
                    # use the help text from the devref
                    option_desc = option_descs[parsed_line[0]]
                else:
                    option_desc = 'No help text available for this option.'
                    if verbose > 0:
                        print(parsed_line[0] + " has no help text")

                # \xa0 is a non-breacking space
                name = parsed_line[0]
                option_desc = option_desc.replace(u'\xa0', u' ')
                default = parsed_line[1]

                o = cfg.StrOpt(name=name, default=default, help=option_desc)
                try:
                    cfg.CONF.register_opt(o, full_section)
                except cfg.DuplicateOptError:
                    pass

def dump_options(options):
    """Dump the list of options with their attributes.

    This output is consumed by the diff_branches script.
    """
    print(pickle.dumps(options._opts_by_name))

def main():
    """Parse and write the Swift configuration options."""

    parser = argparse.ArgumentParser(
        description="Update the swift options tables.",
        usage="%(prog)s docbook|dump [-v] [-s swift_repo] [-m manuals_repo]")
    parser.add_argument('subcommand',
                        help='Action (docbook, dump).',
                        choices=['docbook', 'dump'])
    parser.add_argument('-s', '--swift-repo',
                        dest='swift_repo',
                        help="Location of the swift git repository.",
                        required=False,
                        default="./sources/swift")
    parser.add_argument('-m', '--manuals-repo',
                        dest='manuals_repo',
                        help="Location of the manuals git repository.",
                        required=False,
                        default="./sources/openstack-manuals")
    parser.add_argument('-v', '--verbose',
                        action='count',
                        default=0,
                        dest='verbose',
                        required=False)
    args = parser.parse_args()

    # Avoid cluttering the pickle output, otherwise it's not usable
    if args.subcommand == 'dump':
        args.verbose = 0

    read_options(args.swift_repo, args.manuals_repo, args.verbose)
    options = OptionsCache()

    if args.subcommand == 'docbook':
        write_docbook(options, args.manuals_repo)

    elif args.subcommand == 'dump':
        options.dump()

    return 0

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


<commit_before>import numpy as np
import matplotlib.pyplot as plt

class Histogram():
    def __init__(self, bins, limits=None):

        if np.isscalar(bins):
            if limits is None:
                raise ValueError('If you supply the number of bins you need to supply limits')
            self.n_bins = bins
            self.bin_edges = np.linspace(limits[0], limits[1], bins+1)
            self.limits = limits
        else:
            self.n_bins = len(bins) - 1
            self.bin_edges = bins
            self.limits = (min(bins), max(bins))

        self.entries = np.zeros(self.n_bins)
        self.n_entries = 0
        self.n_underflow = 0
        self.n_overflow = 0

        self.bin_width = np.diff(self.bin_edges)
        if np.all(self.bin_width == self.bin_width[0]):
            self.bin_width = self.bin_width[0]

        self.bin_centers = self.bin_edges[:-1] + 0.5 * self.bin_width

    def fill(self, data):
        data = np.array(data, ndmin=1, copy=False)

        entries, _ = np.histogram(data, self.bin_edges)
        self.entries += entries
        self.n_entries += entries.sum()

        self.n_underflow += np.sum(data < self.limits[0])
        self.n_overflow += np.sum(data > self.limits[1])

    def plot(self, kind='step', ax=None, **kwargs):

        if ax is None:
            ax = plt.gca()

        if kind == 'step':
            px = np.repeat(self.bin_edges, 2)
            py = np.hstack([[0], np.repeat(self.entries, 2), [0]])

            ax.plot(px, py, '-', **kwargs)

        if kind == 'bar':
            ax.bar(
                self.bin_centers,
                self.entries,
                self.bin_width,
                align='center',
                **kwargs
            )

        plt.draw_if_interactive()
<commit_msg>add stats to histogram<commit_after>from __future__ import division, absolute_import
import numpy as np
import matplotlib.pyplot as plt
from .runningstats import Statistics

class Histogram():
    """
    A 1D Histogram that can be filled and stores some statistics
    """
    def __init__(self, bins, limits=None, include_all=False):
        """
        """

        if np.isscalar(bins):
            if limits is None:
                raise ValueError(
                    'If you supply the number of bins'
                    'you need to supply limits'
                )
            self.n_bins = bins
            self.bin_edges = np.linspace(limits[0], limits[1], bins+1)
            self.limits = limits
        else:
            self.n_bins = len(bins) - 1
            self.bin_edges = bins
            self.limits = (min(bins), max(bins))

        self.entries = np.zeros(self.n_bins)
        self.n_entries = 0
        self.n_underflow = 0
        self.n_overflow = 0

        self.bin_width = np.diff(self.bin_edges)
        if np.all(self.bin_width == self.bin_width[0]):
            self.bin_width = self.bin_width[0]

        self.bin_centers = self.bin_edges[:-1] + 0.5 * self.bin_width

        self.stats = Statistics()
        self.include_all = include_all

    def fill(self, data):
        data = np.array(data, ndmin=1, copy=False)

        if self.include_all:
            self.stats.fill(data)
        else:
            a, b = self.limits
            in_limits = np.logical_and(data >= a, data <= b)
            self.stats.fill(data[in_limits])

        entries, _ = np.histogram(data, self.bin_edges)
        self.entries += entries
        self.n_entries += entries.sum()

        self.n_underflow += np.sum(data < self.limits[0])
        self.n_overflow += np.sum(data > self.limits[1])

    def plot(self, kind='step', ax=None, **kwargs):

        if ax is None:
            ax = plt.gca()

        if kind == 'step':
            px = np.repeat(self.bin_edges, 2)
            py = np.hstack([[0], np.repeat(self.entries, 2), [0]])

            ax.plot(px, py, '-', **kwargs)

        if kind == 'bar':
            ax.bar(
                self.bin_centers,
                self.entries,
                self.bin_width,
                align='center',
                **kwargs
            )

        plt.draw_if_interactive()

    def __len__(self):
        return self.entries

    def mean(self):
        return self.stats.mean()

    def std(self):
        return self.stats.std()

    def variance(self):
        return self.stats.variance()

    def kurtosis(self):
        return self.stats.kurtosis()

    def skewness(self):
        return self.stats.skewness()
<|endoftext|>


<commit_before>from pathlib import Path
from datetime import datetime
from pytz import UTC
import h5py
import numpy as np
import logging
import h5py
from struct import pack,unpack
# NOTE: need to use int64 since Numpy thru 1.11 defaults to int32 on Windows for dtype=int, and we need int64 for large files

def getRawInd(fn:Path, BytesPerImage:int, nHeadBytes:int, Nmetadata:int):
    assert isinstance(Nmetadata,int)
    if Nmetadata<1: #no header, only raw images
        fileSizeBytes = fn.stat().st_size
        if fileSizeBytes % BytesPerImage:
            logging.error(f'{fn} may not be read correctly, mismatch frame->file size')

        firstRawIndex = 1 #definition, one-based indexing
        lastRawIndex = fileSizeBytes // BytesPerImage
    else: # normal case 2013-2016
        # gets first and last raw indices from a big .DMCdata file
        with fn.open('rb') as f:
            f.seek(BytesPerImage, 0) # get first raw frame index
            firstRawIndex = meta2rawInd(f,Nmetadata)

            f.seek(-nHeadBytes, 2) #get last raw frame index
            lastRawIndex = meta2rawInd(f,Nmetadata)

    return firstRawIndex, lastRawIndex

def meta2rawInd(f,Nmetadata):
    assert isinstance(Nmetadata,int)

    if Nmetadata<1:
        rawind = None # undefined
    else:
        #FIXME works for .DMCdata version 1 only
        metad = np.fromfile(f, dtype=np.uint16, count=Nmetadata)
        metad = pack('<2H',metad[1],metad[0]) # reorder 2 uint16
        rawind = unpack('<I',metad)[0] #always a tuple

    return rawind

def req2frame(req, N:int=0):
    """
    output has to be numpy.arange for > comparison
    """
    if req is None:
        frame = np.arange(N, dtype=np.int64)
    elif isinstance(req,int): #the user is specifying a step size
        frame = np.arange(0, N, req, dtype=np.int64)
    elif len(req) == 1:
        frame = np.arange(0, N, req[0], dtype=np.int64)
    elif len(req) == 2:
        frame = np.arange(req[0], req[1], dtype=np.int64)
    elif len(req) == 3:
        # this is -1 because user is specifying one-based index
        frame = np.arange(req[0], req[1], req[2], dtype=np.int64) - 1 # keep -1 !
    else: # just return all frames
        frame = np.arange(N, dtype=np.int64)

    return frame

def dir2fn(ofn,ifn,suffix) -> Path:
    """
    ofn = filename or output directory, to create filename based on ifn
    ifn = input filename (don't overwrite!)
    suffix = desired file extension e.g. .h5
    """
    if not ofn: # no output file desired
        return

    ofn = Path(ofn).expanduser()
    ifn = Path(ifn).expanduser()
    assert ifn.is_file()

    if ofn.suffix==suffix: #must already be a filename
        pass
    else: #must be a directory
        assert ofn.is_dir(), f'create directory {ofn}'
        ofn = ofn / ifn.with_suffix(suffix).name

    try:
        assert not ofn.samefile(ifn), f'do not overwrite input file! {ifn}'
    except FileNotFoundError: # a good thing, the output file doesn't exist and hence it's not the input file
        pass

    return ofn

def splitconf(conf,key,i=None,dtype=float,fallback=None,sep=','):
    if conf is None:
        return fallback

    #if isinstance(conf, (ConfigParser,SectionProxy)):
   #     pass
    if isinstance(conf,dict):
        try:
            return conf[key][i]
        except TypeError:
            return conf[key]
        except KeyError:
            return fallback
    else:
        pass
        #raise TypeError('expecting dict or configparser')

    if i is not None:
        assert isinstance(i, (int,slice)),'single integer index only'

    if isinstance(key,(tuple,list)):
        if len(key)>1: #drilldown
            return splitconf(conf[key[0]],key[1:],i,dtype,fallback,sep)
        else:
            return splitconf(conf,key[0],i,dtype,fallback,sep)
    elif isinstance(key,str):
        val = conf.get(key,fallback=fallback)
    else:
        raise TypeError(f'invalid key type {type(key)}')

    try:
        return dtype(val.split(sep)[i])
    except (ValueError,AttributeError,IndexError):
        return fallback
    except TypeError:
        if i is None:
            try:
                return [dtype(v) for v in val.split(sep)] #return list of all values instead of just one
            except ValueError:
                return fallback
        else:
            return fallback
# %% HDF5

def setupimgh5(f, Nframetotal:int, Nrow:int, Ncol:int, dtype=np.uint16, writemode='r+'):
    """
    f: HDF5 handle (or filename)

    h: HDF5 dataset handle
    """
    if isinstance(f, (str,Path)):  # assume new HDF5 file wanted
        print('creating', f)
        with h5py.File(f, writemode, libver='latest') as F:
            return setupimgh5(F,Nframetotal,Nrow,Ncol,dtype)
    elif isinstance(f, h5py.File):
        h = f.create_dataset('/rawimg',
                 shape =  (Nframetotal,Nrow,Ncol),
                 dtype=dtype,
                 chunks= (1,Nrow,Ncol), # each image is a chunk
                 compression='gzip',
                 compression_opts=1, #no difference in filesize from 1 to 5, except much faster to use lower numbers!
                 shuffle= True,
                 fletcher32= True,
                 track_times= True)
        h.attrs["CLASS"] = np.string_("IMAGE")
        h.attrs["IMAGE_VERSION"] = np.string_("1.2")
        h.attrs["IMAGE_SUBCLASS"] = np.string_("IMAGE_GRAYSCALE")
        h.attrs["DISPLAY_ORIGIN"] = np.string_("LL")
        h.attrs['IMAGE_WHITE_IS_ZERO'] = np.uint8(0)
    else:
        raise TypeError(f'{type(f)} is not correct, must be filename or h5py.File HDF5 file handle')

    return h

def vid2h5(data, ut1, rawind, ticks, outfn, P, cmdlog='', i:int=0, Nfile:int=1):
    assert outfn,'must provide a filename to write'

    outfn = Path(outfn).expanduser()
    print('writing', outfn)
    #%% saving
    if outfn.suffix == '.h5':
        """
        Reference: https://www.hdfgroup.org/HDF5/doc/ADGuide/ImageSpec.html
        Thanks to Eric Piel of Delmic for pointing out this spec
        * the HDF5 attributess set are necessary to put HDFView into image mode and enables
        other conforming readers to easily play images stacks as video.
        * the string_() calls are necessary to make fixed length strings per HDF5 spec
        """
        #NOTE write mode r+ to not overwrite incremental images
        if outfn.is_file():
            writemode = 'r+' # append
        else:
            writemode = 'w'

        N = data.shape[0] * Nfile
        ind = slice(i,i+data.shape[0])

        with h5py.File(outfn, writemode, libver='latest') as f:
            if data is not None:
                if 'rawimg' not in f:  # first run
                    setupimgh5(f, N, data.shape[1], data.shape[2])

                f['/rawimg'][ind, ...] = data

            if ut1 is not None:
                print(f'writing from {datetime.utcfromtimestamp(ut1[0]).replace(tzinfo=UTC)} to {datetime.utcfromtimestamp(ut1[-1]).replace(tzinfo=UTC)}')
                if 'ut1_unix' not in f:
                    fut1 = f.create_dataset('/ut1_unix', shape=(N,), dtype=float, fletcher32=True)
                    fut1.attrs['units'] = 'seconds since Unix epoch Jan 1 1970 midnight'

                f['/ut1_unix'][ind] = ut1

            if rawind is not None:
                if 'rawind' not in f:
                    fri = f.create_dataset('/rawind', shape=(N,), dtype=np.int64, fletcher32=True)
                    fri.attrs['units'] = 'one-based index since camera program started this session'

                f['/rawind'][ind] = rawind

            if ticks is not None:
                if 'ticks' not in f:
                    ftk = f.create_dataset('/ticks', shape=(N,), dtype=np.uint64, fletcher32=True)  # Uint64
                    ftk.attrs['units'] = 'FPGA tick counter for each image frame'

                f['/ticks'][ind] = ticks

            if 'params' not in f:
                cparam = np.array((
                           P['kineticsec'],
                           P['rotccw'],
                           P['transpose']==True,
                           P['flipud']==True,
                           P['fliplr']==True,
                            1),
                           dtype=[('kineticsec','f8'),
                                  ('rotccw',    'i1'),
                                  ('transpose', 'i1'),
                                  ('flipud',    'i1'),
                                  ('fliplr',    'i1'),
                                  ('questionable_ut1','i1')]
                           )

                f.create_dataset('/params',data=cparam) #cannot use fletcher32 here, typeerror

            if 'sensorloc' not in f and 'sensorloc' in P:
                l = P['sensorloc']
                lparam = np.array((l[0],l[1],l[2]),
                                  dtype=[('lat','f8'),('lon','f8'),('alt_m','f8')])

                Ld = f.create_dataset('/sensorloc', data=lparam) #cannot use fletcher32 here, typeerror
                Ld.attrs['units'] = 'WGS-84 lat (deg),lon (deg), altitude (meters)'

            if 'cmdlog' not in f:
                if isinstance(cmdlog,(tuple,list)):
                    cmdlog = ' '.join(cmdlog)
                f['/cmdlog'] = str(cmdlog) #cannot use fletcher32 here, typeerror

            if 'header' not in f and 'header' in P:
                f['/header'] = str(P['header'])

    elif outfn.suffix == '.fits':
        from astropy.io import fits
        #NOTE the with... syntax does NOT yet work with astropy.io.fits
        hdu = fits.PrimaryHDU(data)
        hdu.writeto(outfn,clobber=False,checksum=True)
        #no close
        """
        Note: the orientation of this FITS in NASA FV program and the preview
        image shown in Python should/must have the same orientation and pixel indexing')
        """

    elif outfn.suffix == '.mat':
        from scipy.io import savemat
        matdata = {'imgdata':data.transpose(1,2,0)} #matlab is fortran order
        savemat(outfn,matdata,oned_as='column')
    else:
        raise ValueError(f'what kind of file is {outfn}')

def imgwriteincr(fn, imgs, imgslice):
    """
    writes HDF5 huge image files in increments
    """
    if isinstance(imgslice,int):
        if imgslice and not (imgslice % 2000):
            print(f'appending images {imgslice} to {fn}')

    if isinstance(fn, Path):
        assert fn.suffix == '.h5','Expecting to write .h5 file' # avoid accidental overwriting of source file due to misspecified command line

        with h5py.File(fn, 'r+', libver='latest') as f:
            f['/rawimg'][imgslice,:,:] = imgs
    elif isinstance(fn, h5py.File):
        f['/rawimg'][imgslice,:,:] = imgs
    else:
        raise TypeError(f'fn must be Path or h5py.File instead of {type(fn)}')<commit_msg>fix indexing, more metadata<commit_after>from pathlib import Path
from datetime import datetime
from pytz import UTC
import h5py
import numpy as np
import logging
from struct import pack,unpack
# NOTE: need to use int64 since Numpy thru 1.11 defaults to int32 on Windows for dtype=int, and we need int64 for large files

def getRawInd(fn:Path, BytesPerImage:int, nHeadBytes:int, Nmetadata:int):
    assert isinstance(Nmetadata,int)
    if Nmetadata<1: #no header, only raw images
        fileSizeBytes = fn.stat().st_size
        if fileSizeBytes % BytesPerImage:
            logging.error(f'{fn} may not be read correctly, mismatch frame->file size')

        firstRawIndex = 1 #definition, one-based indexing
        lastRawIndex = fileSizeBytes // BytesPerImage
    else: # normal case 2013-2016
        # gets first and last raw indices from a big .DMCdata file
        with fn.open('rb') as f:
            f.seek(BytesPerImage, 0) # get first raw frame index
            firstRawIndex = meta2rawInd(f,Nmetadata)

            f.seek(-nHeadBytes, 2) #get last raw frame index
            lastRawIndex = meta2rawInd(f,Nmetadata)

    return firstRawIndex, lastRawIndex

def meta2rawInd(f,Nmetadata):
    assert isinstance(Nmetadata,int)

    if Nmetadata<1:
        rawind = None # undefined
    else:
        #FIXME works for .DMCdata version 1 only
        metad = np.fromfile(f, dtype=np.uint16, count=Nmetadata)
        metad = pack('<2H',metad[1],metad[0]) # reorder 2 uint16
        rawind = unpack('<I',metad)[0] #always a tuple

    return rawind

def req2frame(req, N:int=0):
    """
    output has to be numpy.arange for > comparison
    """
    if req is None:
        frame = np.arange(N, dtype=np.int64)
    elif isinstance(req,int): #the user is specifying a step size
        frame = np.arange(0, N, req, dtype=np.int64)
    elif len(req) == 1:
        frame = np.arange(0, N, req[0], dtype=np.int64)
    elif len(req) == 2:
        frame = np.arange(req[0], req[1], dtype=np.int64)
    elif len(req) == 3:
        # this is -1 because user is specifying one-based index
        frame = np.arange(req[0], req[1], req[2], dtype=np.int64) - 1 # keep -1 !
    else: # just return all frames
        frame = np.arange(N, dtype=np.int64)

    return frame

def dir2fn(ofn,ifn,suffix) -> Path:
    """
    ofn = filename or output directory, to create filename based on ifn
    ifn = input filename (don't overwrite!)
    suffix = desired file extension e.g. .h5
    """
    if not ofn: # no output file desired
        return

    ofn = Path(ofn).expanduser()
    ifn = Path(ifn).expanduser()
    assert ifn.is_file()

    if ofn.suffix==suffix: #must already be a filename
        pass
    else: #must be a directory
        assert ofn.is_dir(), f'create directory {ofn}'
        ofn = ofn / ifn.with_suffix(suffix).name

    try:
        assert not ofn.samefile(ifn), f'do not overwrite input file! {ifn}'
    except FileNotFoundError: # a good thing, the output file doesn't exist and hence it's not the input file
        pass

    return ofn

def splitconf(conf,key,i=None,dtype=float,fallback=None,sep=','):
    if conf is None:
        return fallback

    #if isinstance(conf, (ConfigParser,SectionProxy)):
   #     pass
    if isinstance(conf,dict):
        try:
            return conf[key][i]
        except TypeError:
            return conf[key]
        except KeyError:
            return fallback
    else:
        pass
        #raise TypeError('expecting dict or configparser')

    if i is not None:
        assert isinstance(i, (int,slice)),'single integer index only'

    if isinstance(key,(tuple,list)):
        if len(key)>1: #drilldown
            return splitconf(conf[key[0]],key[1:],i,dtype,fallback,sep)
        else:
            return splitconf(conf,key[0],i,dtype,fallback,sep)
    elif isinstance(key,str):
        val = conf.get(key,fallback=fallback)
    else:
        raise TypeError(f'invalid key type {type(key)}')

    try:
        return dtype(val.split(sep)[i])
    except (ValueError,AttributeError,IndexError):
        return fallback
    except TypeError:
        if i is None:
            try:
                return [dtype(v) for v in val.split(sep)] #return list of all values instead of just one
            except ValueError:
                return fallback
        else:
            return fallback
# %% HDF5

def setupimgh5(f, Nframetotal:int, Nrow:int, Ncol:int, dtype=np.uint16, writemode='r+'):
    """
    f: HDF5 handle (or filename)

    h: HDF5 dataset handle
    """
    if isinstance(f, (str,Path)):  # assume new HDF5 file wanted
        print('creating', f)
        with h5py.File(f, writemode, libver='latest') as F:
            return setupimgh5(F,Nframetotal,Nrow,Ncol,dtype)
    elif isinstance(f, h5py.File):
        h = f.create_dataset('/rawimg',
                 shape =  (Nframetotal,Nrow,Ncol),
                 dtype=dtype,
                 chunks= (1,Nrow,Ncol), # each image is a chunk
                 compression='gzip',
                 compression_opts=1, #no difference in filesize from 1 to 5, except much faster to use lower numbers!
                 shuffle= True,
                 fletcher32= True,
                 track_times= True)
        h.attrs["CLASS"] = np.string_("IMAGE")
        h.attrs["IMAGE_VERSION"] = np.string_("1.2")
        h.attrs["IMAGE_SUBCLASS"] = np.string_("IMAGE_GRAYSCALE")
        h.attrs["DISPLAY_ORIGIN"] = np.string_("LL")
        h.attrs['IMAGE_WHITE_IS_ZERO'] = np.uint8(0)
    else:
        raise TypeError(f'{type(f)} is not correct, must be filename or h5py.File HDF5 file handle')

    return h

def vid2h5(data, ut1, rawind, ticks, outfn, P, cmdlog='', i:int=0, Nfile:int=1):
    assert outfn,'must provide a filename to write'

    outfn = Path(outfn).expanduser()
    txtupd = f'converting file # {i} / {Nfile}'
    if 'spoolfn' in P:
        txtupd += f' from {P["spoolfn"].name}'
    txtupd += f' to {outfn}'
    print(txtupd)
    #%% saving
    if outfn.suffix == '.h5':
        """
        Reference: https://www.hdfgroup.org/HDF5/doc/ADGuide/ImageSpec.html
        Thanks to Eric Piel of Delmic for pointing out this spec
        * the HDF5 attributess set are necessary to put HDFView into image mode and enables
        other conforming readers to easily play images stacks as video.
        * the string_() calls are necessary to make fixed length strings per HDF5 spec
        """
        #NOTE write mode r+ to not overwrite incremental images
        if outfn.is_file():
            writemode = 'r+' # append
        else:
            writemode = 'w'

        if data is not None:

            N = data.shape[0] * Nfile
            ind = slice(i*data.shape[0], (i+1)*data.shape[0])
        else: # FIXME haven't reevaluated for update only case
            for q in (ut1,rawind,ticks):
                if q is not None:
                    N = len(q)
                    break
            ind = slice(None)

        with h5py.File(outfn, writemode, libver='latest') as f:
            if data is not None:
                if 'rawimg' not in f:  # first run
                    setupimgh5(f, N, data.shape[1], data.shape[2])

                f['/rawimg'][ind, ...] = data

            if ut1 is not None:
                print(f'writing from {datetime.utcfromtimestamp(ut1[0]).replace(tzinfo=UTC)} to {datetime.utcfromtimestamp(ut1[-1]).replace(tzinfo=UTC)}')
                if 'ut1_unix' not in f:
                    fut1 = f.create_dataset('/ut1_unix', shape=(N,), dtype=float, fletcher32=True)
                    fut1.attrs['units'] = 'seconds since Unix epoch Jan 1 1970 midnight'

                f['/ut1_unix'][ind] = ut1

            if rawind is not None:
                if 'rawind' not in f:
                    fri = f.create_dataset('/rawind', shape=(N,), dtype=np.int64, fletcher32=True)
                    fri.attrs['units'] = 'one-based index since camera program started this session'

                f['/rawind'][ind] = rawind

            if ticks is not None:
                if 'ticks' not in f:
                    ftk = f.create_dataset('/ticks', shape=(N,), dtype=np.uint64, fletcher32=True)  # Uint64
                    ftk.attrs['units'] = 'FPGA tick counter for each image frame'

                f['/ticks'][ind] = ticks

            if 'spoolfn' in P:
                # http://docs.h5py.org/en/latest/strings.html
                if 'spoolfn' not in f:
                    fsp = f.create_dataset('/spoolfn', shape=(N,), dtype=h5py.special_dtype(vlen=bytes))
                    fsp.attrs['description'] = 'input filename data was extracted from'

                f['/spoolfn'][ind] = P['spoolfn'].name

            if 'params' not in f:
                cparam = np.array((
                           P['kineticsec'],
                           P['rotccw'],
                           P['transpose']==True,
                           P['flipud']==True,
                           P['fliplr']==True,
                            1),
                           dtype=[('kineticsec','f8'),
                                  ('rotccw',    'i1'),
                                  ('transpose', 'i1'),
                                  ('flipud',    'i1'),
                                  ('fliplr',    'i1'),
                                  ('questionable_ut1','i1')]
                           )

                f.create_dataset('/params',data=cparam) #cannot use fletcher32 here, typeerror

            if 'sensorloc' not in f and 'sensorloc' in P:
                l = P['sensorloc']
                lparam = np.array((l[0],l[1],l[2]),
                                  dtype=[('lat','f8'),('lon','f8'),('alt_m','f8')])

                Ld = f.create_dataset('/sensorloc', data=lparam) #cannot use fletcher32 here, typeerror
                Ld.attrs['units'] = 'WGS-84 lat (deg),lon (deg), altitude (meters)'

            if 'cmdlog' not in f:
                if isinstance(cmdlog,(tuple,list)):
                    cmdlog = ' '.join(cmdlog)
                f['/cmdlog'] = str(cmdlog) #cannot use fletcher32 here, typeerror

            if 'header' not in f and 'header' in P:
                f['/header'] = str(P['header'])

            if 'hdf5version' not in f:
                f['/hdf5version'] = h5py.version.hdf5_version_tuple

    elif outfn.suffix == '.fits':
        from astropy.io import fits
        #NOTE the with... syntax does NOT yet work with astropy.io.fits
        hdu = fits.PrimaryHDU(data)
        hdu.writeto(outfn,clobber=False,checksum=True)
        #no close
        """
        Note: the orientation of this FITS in NASA FV program and the preview
        image shown in Python should/must have the same orientation and pixel indexing')
        """

    elif outfn.suffix == '.mat':
        from scipy.io import savemat
        matdata = {'imgdata':data.transpose(1,2,0)} #matlab is fortran order
        savemat(outfn,matdata,oned_as='column')
    else:
        raise ValueError(f'what kind of file is {outfn}')

def imgwriteincr(fn, imgs, imgslice):
    """
    writes HDF5 huge image files in increments
    """
    if isinstance(imgslice,int):
        if imgslice and not (imgslice % 2000):
            print(f'appending images {imgslice} to {fn}')

    if isinstance(fn, Path):
        assert fn.suffix == '.h5','Expecting to write .h5 file' # avoid accidental overwriting of source file due to misspecified command line

        with h5py.File(fn, 'r+', libver='latest') as f:
            f['/rawimg'][imgslice,:,:] = imgs
    elif isinstance(fn, h5py.File):
        f['/rawimg'][imgslice,:,:] = imgs
    else:
        raise TypeError(f'fn must be Path or h5py.File instead of {type(fn)}')<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
recursive list of hdf5 file contents
can also simply in Terminal type:   h5ls -r myFile.h5
this is just an example of traversing the h5 contents, it may not be useful
in and of itself.
Michael Hirsch 2014
"""

import h5py

def h5lister(h5fn):

    with h5py.File(h5fn,'r') as f:

        for dkey,dval in dict(f).iteritems():

            if isinstance(dval,h5py.Dataset):
                print(dval)
            elif isinstance(dval,h5py.Group):
                print(dval)
                for gkey,gval in dict(dval).items():
                    print(gval)

if __name__ == '__main__':
	import sys
	h5lister(sys.argv[1])
<commit_msg>move to pybashutils<commit_after><|endoftext|>


<commit_before>from Crypto.Cipher import AES
from pkcs7 import PKCS7Encoder
from m3u8.model import Key

from collections import namedtuple
import errno
import os
import urllib2
import urlparse

MPEG_TS_HEADER = '474000'.decode('hex')

import m3u8

def consume(m3u8_uri, destination_path, new_key=False):
    '''
    Given a ``m3u8_uri``, downloads all files to disk
    The remote path structure is maintained under ``destination_path``

    - new_key:
        If False, keeps existing encryption
        If None, decrypts file
    '''
    playlist = m3u8.load(m3u8_uri)

    if playlist.is_variant:
        return consume_variant_playlist(playlist, m3u8_uri, destination_path)
    else:
        return consume_single_playlist(playlist, m3u8_uri, destination_path, new_key)

def consume_variant_playlist(playlist, m3u8_uri, destination_path):
    full_path = build_full_path(destination_path, m3u8_uri)
    save_m3u8(playlist, m3u8_uri, full_path)

    for p in playlist.playlists:
        consume(p.absolute_uri, destination_path)
    return True

def consume_single_playlist(playlist, m3u8_uri, destination_path, new_key=False):
    full_path = build_full_path(destination_path, m3u8_uri)
    downloaded_key = download_key(playlist, destination_path)
    downloaded_segments = download_segments(playlist, full_path, new_key)

    m3u8_has_changed = downloaded_key or any(downloaded_segments)
    if m3u8_has_changed:
        save_m3u8(playlist, m3u8_uri, full_path, new_key)

    return m3u8_has_changed

def build_intermediate_path(m3u8_uri):
    url_path = urlparse.urlparse(m3u8_uri).path
    return os.path.dirname(url_path)

def build_full_path(destination_path, m3u8_uri):
    intermediate_path = build_intermediate_path(m3u8_uri)[1:] # ignore first "/"
    full_path = os.path.join(destination_path, intermediate_path)
    ensure_directory_exists(full_path)
    return full_path

def ensure_directory_exists(directory):
    try:
        os.makedirs(directory)
    except OSError as error:
        if error.errno != errno.EEXIST:
            raise

def download_key(playlist, destination_path):
    if playlist.key:
        filename = download_to_file(playlist.key.absolute_uri, destination_path)
        with open(filename, 'rb') as f:
            playlist.key.key_value = f.read()
    return False

def download_segments(playlist, destination_path, new_key):
    segments = [segment.absolute_uri for segment in playlist.segments]
    return [download_to_file(uri, destination_path, playlist.key, new_key) for uri in segments]

def save_m3u8(playlist, m3u8_uri, full_path, new_key=False):
    playlist.basepath = build_intermediate_path(m3u8_uri)
    if new_key:
        save_new_key(new_key, full_path)
        playlist.version = "2"
        playlist.key = new_key
    elif new_key is None:
        playlist.key = None
    filename = os.path.join(full_path, os.path.basename(m3u8_uri))
    playlist.dump(filename)

def save_new_key(new_key, destination_path):
    filename = os.path.join(destination_path, os.path.basename(new_key.uri))
    if not os.path.exists(filename):
        with open(filename, 'wb') as f:
            f.write(new_key.key_value)

def download_to_file(uri, destination_path, old_key=None, new_key=False):
    '''
    Retrives the file if it does not exist locally and changes the encryption if needed.

    '''
    filename = os.path.join(destination_path, os.path.basename(uri))
    if not os.path.exists(filename):
        request = urllib2.urlopen(url=uri)
        raw = request.read()
        if new_key is not False:
            plain = decrypt(raw, old_key) if old_key else raw
            raw = encrypt(plain, new_key) if new_key else plain
        with open(filename, 'wb') as f:
            f.write(raw)
        return filename
    return False

def random_key(key_name):
    class IV:
        def __init__(self, iv):
            self.iv = iv

        def __str__(self):
            return '0X' + self.iv.encode('hex')

    key = Key(method='AES-128', uri=key_name, baseuri=None,  iv=IV(os.urandom(16)))
    key.key_value = os.urandom(16)
    return key

def get_key_iv(key):
    iv = str(key.iv)[2:] # Removes 0X prefix
    return iv.decode('hex')

def encrypt(data, key):
    encoder = PKCS7Encoder()
    padded_text = encoder.encode(data)
    encryptor = AES.new(key.key_value, AES.MODE_CBC, get_key_iv(key))
    encrypted_data = encryptor.encrypt(padded_text)

    return encrypted_data

def decrypt(data, key):
    encoder = PKCS7Encoder()
    decryptor = AES.new(key.key_value, AES.MODE_CBC, get_key_iv(key))
    plain = decryptor.decrypt(data)

    return encoder.decode(plain)
<commit_msg>consumer: Removed unused imports<commit_after>from Crypto.Cipher import AES
from pkcs7 import PKCS7Encoder

import errno
import os
import urllib2
import urlparse

import m3u8

def consume(m3u8_uri, destination_path, new_key=False):
    '''
    Given a ``m3u8_uri``, downloads all files to disk
    The remote path structure is maintained under ``destination_path``

    - new_key:
        If False, keeps existing encryption
        If None, decrypts file
    '''
    playlist = m3u8.load(m3u8_uri)

    if playlist.is_variant:
        return consume_variant_playlist(playlist, m3u8_uri, destination_path)
    else:
        return consume_single_playlist(playlist, m3u8_uri, destination_path, new_key)

def consume_variant_playlist(playlist, m3u8_uri, destination_path):
    full_path = build_full_path(destination_path, m3u8_uri)
    save_m3u8(playlist, m3u8_uri, full_path)

    for p in playlist.playlists:
        consume(p.absolute_uri, destination_path)
    return True

def consume_single_playlist(playlist, m3u8_uri, destination_path, new_key=False):
    full_path = build_full_path(destination_path, m3u8_uri)
    downloaded_key = download_key(playlist, destination_path)
    downloaded_segments = download_segments(playlist, full_path, new_key)

    m3u8_has_changed = downloaded_key or any(downloaded_segments)
    if m3u8_has_changed:
        save_m3u8(playlist, m3u8_uri, full_path, new_key)

    return m3u8_has_changed

def build_intermediate_path(m3u8_uri):
    url_path = urlparse.urlparse(m3u8_uri).path
    return os.path.dirname(url_path)

def build_full_path(destination_path, m3u8_uri):
    intermediate_path = build_intermediate_path(m3u8_uri)[1:] # ignore first "/"
    full_path = os.path.join(destination_path, intermediate_path)
    ensure_directory_exists(full_path)
    return full_path

def ensure_directory_exists(directory):
    try:
        os.makedirs(directory)
    except OSError as error:
        if error.errno != errno.EEXIST:
            raise

def download_key(playlist, destination_path):
    if playlist.key:
        filename = download_to_file(playlist.key.absolute_uri, destination_path)
        with open(filename, 'rb') as f:
            playlist.key.key_value = f.read()
    return False

def download_segments(playlist, destination_path, new_key):
    segments = [segment.absolute_uri for segment in playlist.segments]
    return [download_to_file(uri, destination_path, playlist.key, new_key) for uri in segments]

def save_m3u8(playlist, m3u8_uri, full_path, new_key=False):
    playlist.basepath = build_intermediate_path(m3u8_uri)
    if new_key:
        save_new_key(new_key, full_path)
        playlist.version = "2"
        playlist.key = new_key
    elif new_key is None:
        playlist.key = None
    filename = os.path.join(full_path, os.path.basename(m3u8_uri))
    playlist.dump(filename)

def save_new_key(new_key, destination_path):
    filename = os.path.join(destination_path, os.path.basename(new_key.uri))
    if not os.path.exists(filename):
        with open(filename, 'wb') as f:
            f.write(new_key.key_value)

def download_to_file(uri, destination_path, old_key=None, new_key=False):
    '''
    Retrives the file if it does not exist locally and changes the encryption if needed.

    '''
    filename = os.path.join(destination_path, os.path.basename(uri))
    if not os.path.exists(filename):
        request = urllib2.urlopen(url=uri)
        raw = request.read()
        if new_key is not False:
            plain = decrypt(raw, old_key) if old_key else raw
            raw = encrypt(plain, new_key) if new_key else plain
        with open(filename, 'wb') as f:
            f.write(raw)
        return filename
    return False

def random_key(key_name):
    class IV:
        def __init__(self, iv):
            self.iv = iv

        def __str__(self):
            return '0X' + self.iv.encode('hex')

    key = m3u8.model.Key(method='AES-128', uri=key_name, baseuri=None,  iv=IV(os.urandom(16)))
    key.key_value = os.urandom(16)
    return key

def get_key_iv(key):
    iv = str(key.iv)[2:] # Removes 0X prefix
    return iv.decode('hex')

def encrypt(data, key):
    encoder = PKCS7Encoder()
    padded_text = encoder.encode(data)
    encryptor = AES.new(key.key_value, AES.MODE_CBC, get_key_iv(key))
    encrypted_data = encryptor.encrypt(padded_text)

    return encrypted_data

def decrypt(data, key):
    encoder = PKCS7Encoder()
    decryptor = AES.new(key.key_value, AES.MODE_CBC, get_key_iv(key))
    plain = decryptor.decrypt(data)

    return encoder.decode(plain)
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
from uuid import uuid4
from datetime import datetime, timedelta
from random import choice
import hashlib
import logging

import sqlalchemy as sa
from sqlalchemy.orm import relationship
from sqlalchemy import or_
from ujson import dumps
from tornado.concurrent import return_future

from holmes.models import Base
from holmes.utils import get_domain_from_url

class Page(Base):
    __tablename__ = "pages"

    id = sa.Column(sa.Integer, primary_key=True)
    url = sa.Column('url', sa.String(2000), nullable=False)
    url_hash = sa.Column('url_hash', sa.String(128), nullable=False)
    uuid = sa.Column('uuid', sa.String(36), default=uuid4, nullable=False)
    created_date = sa.Column('created_date', sa.DateTime, default=datetime.utcnow, nullable=False)

    domain_id = sa.Column('domain_id', sa.Integer, sa.ForeignKey('domains.id'))

    reviews = relationship("Review", backref="page", foreign_keys='[Review.page_id]')

    last_review_id = sa.Column('last_review_id', sa.Integer, sa.ForeignKey('reviews.id'))
    last_review = relationship("Review", foreign_keys=[last_review_id])
    last_review_date = sa.Column('last_review_date', sa.DateTime, nullable=True)
    last_review_uuid = sa.Column('last_review_uuid', sa.String(36), nullable=True)

    last_modified = sa.Column('last_modified', sa.DateTime, nullable=True)
    expires = sa.Column('expires', sa.DateTime, nullable=True)

    violations_count = sa.Column('violations_count', sa.Integer, server_default='0', nullable=False)

    score = sa.Column('score', sa.Float, default=0.0, nullable=False)

    def to_dict(self):
        return {
            'uuid': str(self.uuid),
            'url': self.url,
            'lastModified': self.last_modified,
            'expires': self.expires,
            'score': self.score
        }

    def __str__(self):
        return str(self.uuid)

    def __repr__(self):
        return str(self)

    def get_violations_per_day(self, db):
        from holmes.models import Review, Violation  # Prevent circular dependency

        violations = db \
            .query(
                sa.func.year(Review.completed_date).label('year'),
                sa.func.month(Review.completed_date).label('month'),
                sa.func.day(Review.completed_date).label('day'),
                sa.func.count(Violation.id).label('violation_count'),
                sa.func.sum(Violation.points).label('violation_points')
            ).join(
                Page, Page.id == Review.page_id
            ).join(
                Violation, Violation.review_id == Review.id
            ).filter(Review.is_complete == True).filter(Review.page_id == self.id) \
            .group_by(
                sa.func.year(Review.completed_date),
                sa.func.month(Review.completed_date),
                sa.func.day(Review.completed_date),
            ) \
            .order_by(Review.completed_date) \
            .all()

        result = []

        for day in violations:
            dt = "%d-%d-%d" % (day.year, day.month, day.day)
            result.append({
                "completedAt": dt,
                "violation_count": int(day.violation_count),
                "violation_points": int(day.violation_points)
            })

        return result

    @classmethod
    def by_uuid(cls, uuid, db):
        return db.query(Page).filter(Page.uuid == uuid).first()

    @classmethod
    def get_page_count(cls, db):
        return int(db.query(sa.func.count(Page.id)).scalar())

    @classmethod
    def update_scores(cls, individual_score, db):
        for i in range(3):
            db.begin(subtransactions=True)
            try:
                db.query(Page).update({'score': Page.score + individual_score})
                db.flush()
                db.commit()
                break
            except Exception:
                err = sys.exc_info()[1]
                if 'Deadlock found' in str(err):
                    logging.error('Deadlock happened! Trying again (try number %d)! (Details: %s)' % (i, str(err)))
                else:
                    db.rollback()
                    raise

    @classmethod
    def get_next_job(cls, db, expiration, cache, lock_expiration):
        from holmes.models import Settings, Domain

        expired_time = datetime.now() - timedelta(seconds=expiration)

        settings = Settings.instance(db)

        active_domains = [item.id for item in db.query(Domain.id).filter(Domain.is_active).all()]

        if not active_domains:
            return None

        pages_query = db.query(Page.uuid, Page.url, Page.score, Page.last_review_date) \
                        .filter(Page.last_review_date == None) \
                        .filter(Page.domain_id.in_(active_domains))

        pages_in_need_of_review = pages_query.order_by(Page.score.desc())[:200]

        if len(pages_in_need_of_review) == 0:
            pages_query = db.query(Page.uuid, Page.url, Page.score, Page.last_review_date) \
                            .filter(Page.last_review_date <= expired_time) \
                            .filter(Page.domain_id.in_(active_domains))

            pages_in_need_of_review = pages_query.order_by(Page.score.desc())[:200]

            if len(pages_in_need_of_review) == 0:
                if settings.lambda_score > 0:
                    cls.update_pages_score_by(settings, settings.lambda_score, db)
                return None

        if settings.lambda_score > pages_in_need_of_review[0].score:
            cls.update_pages_score_by(settings, settings.lambda_score, db)
            return None

        page = choice(pages_in_need_of_review)

        for i in range(10):
            lock = cache.has_next_job_lock(page.url, lock_expiration)
            if lock is None:
                page = choice(pages_in_need_of_review)
            else:
                break

        if lock is None:
            return None

        return {
            'page': str(page.uuid),
            'url': page.url,
            'score': page.score,
            'lock': lock
        }

    @classmethod
    def update_pages_score_by(cls, settings, score, db):
        settings.lambda_score = 0
        page_count = cls.get_page_count(db)
        individual_score = float(score) / float(page_count)
        cls.update_scores(individual_score, db)

    @classmethod
    @return_future
    def add_page(cls, db, cache, url, score, fetch_method, publish_method, config, callback):
        domain_name, domain_url = get_domain_from_url(url)
        if not url or not domain_name:
            callback((False, url, {
                'reason': 'invalid_url',
                'url': url,
                'status': None,
                'details': 'Domain name could not be determined.'
            }))
            return

        logging.info('Obtaining "%s"...' % url)

        fetch_method(
            url,
            cls.handle_request(cls.handle_add_page(db, cache, url, score, publish_method, callback)),
            proxy_host=config.HTTP_PROXY_HOST,
            proxy_port=config.HTTP_PROXY_PORT
        )

    @classmethod
    def handle_request(cls, callback):
        def handle(*args, **kw):
            response = args[-1]  # supports (url, response) and just response

            if hasattr(response, 'status_code'):
                status_code = response.status_code
            elif hasattr(response, 'code'):
                status_code = response.code
            else:
                status_code = 400

            if hasattr(response, 'body'):
                text = response.body
            elif hasattr(response, 'text'):
                text = response.text
            else:
                text = 'Empty response.text'

            callback(status_code, text, response.effective_url)

        return handle

    @classmethod
    def handle_add_page(cls, db, cache, url, score, publish_method, callback):
        def handle(code, body, effective_url):
            if code > 399:
                callback((False, url, {
                    'reason': 'invalid_url',
                    'url': url,
                    'status': code,
                    'details': body
                }))
                return

            if effective_url != url:
                callback((False, url, {
                    'reason': 'redirect',
                    'url': url,
                    'effectiveUrl': effective_url
                }))
                return

            domain = cls.add_domain(url, db, publish_method)
            page_uuid = cls.insert_or_update_page(url, score, domain, db, publish_method)

            callback((True, url, page_uuid))

            cache.increment_page_count(domain)
            cache.increment_page_count()

        return handle

    @classmethod
    def insert_or_update_page(cls, url, score, domain, db, publish_method):
        url = url.encode('utf-8')
        url_hash = hashlib.sha512(url).hexdigest()
        page = db.query(Page).filter(Page.url_hash==url_hash).first()

        if page:
            for i in range(3):
                db.begin(subtransactions=True)
                try:
                    db.query(Page).filter(Page.id == page.id).update({'score': Page.score + score})
                    db.flush()
                    db.commit()
                    break
                except Exception:
                    err = sys.exc_info()[1]
                    if 'Deadlock found' in str(err):
                        logging.error('Deadlock happened! Trying again (try number %d)! (Details: %s)' % (i, str(err)))
                    else:
                        db.rollback()
                        raise

            return page.uuid

        db.begin(subtransactions=True)
        try:
            page = Page(url=url, url_hash=url_hash, domain=domain, score=score)
            db.add(page)
            db.flush()
            db.commit()
        except Exception:
            db.rollback()
            err = sys.exc_info()[1]
            if 'Duplicate entry' in str(err):
                logging.error('Duplicate entry! (Details: %s)' % str(err))
            else:
                raise

        publish_method(dumps({
            'type': 'new-page',
            'pageUrl': str(url)
        }))

        return page.uuid

    @classmethod
    def add_domain(cls, url, db, publish_method):
        from holmes.models import Domain

        domain_name, domain_url = get_domain_from_url(url)

        domains = db.query(Domain).filter(or_(
            Domain.name == domain_name,
            Domain.name == domain_name.rstrip('/'),
            Domain.name == "%s/" % domain_name
        )).all()

        if not domains:
            domain = None
        else:
            domain = domains[0]

        if not domain:
            url_hash = hashlib.sha512(domain_url).hexdigest()
            domain = Domain(url=domain_url, url_hash=url_hash, name=domain_name)
            db.add(domain)
            db.flush()
            db.commit()

            publish_method(dumps({
                'type': 'new-domain',
                'domainUrl': str(domain_url)
            }))

        return domain
<commit_msg>Fixing page count<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
from uuid import uuid4
from datetime import datetime, timedelta
from random import choice
import hashlib
import logging

import sqlalchemy as sa
from sqlalchemy.orm import relationship
from sqlalchemy import or_
from ujson import dumps
from tornado.concurrent import return_future

from holmes.models import Base
from holmes.utils import get_domain_from_url

class Page(Base):
    __tablename__ = "pages"

    id = sa.Column(sa.Integer, primary_key=True)
    url = sa.Column('url', sa.String(2000), nullable=False)
    url_hash = sa.Column('url_hash', sa.String(128), nullable=False)
    uuid = sa.Column('uuid', sa.String(36), default=uuid4, nullable=False)
    created_date = sa.Column('created_date', sa.DateTime, default=datetime.utcnow, nullable=False)

    domain_id = sa.Column('domain_id', sa.Integer, sa.ForeignKey('domains.id'))

    reviews = relationship("Review", backref="page", foreign_keys='[Review.page_id]')

    last_review_id = sa.Column('last_review_id', sa.Integer, sa.ForeignKey('reviews.id'))
    last_review = relationship("Review", foreign_keys=[last_review_id])
    last_review_date = sa.Column('last_review_date', sa.DateTime, nullable=True)
    last_review_uuid = sa.Column('last_review_uuid', sa.String(36), nullable=True)

    last_modified = sa.Column('last_modified', sa.DateTime, nullable=True)
    expires = sa.Column('expires', sa.DateTime, nullable=True)

    violations_count = sa.Column('violations_count', sa.Integer, server_default='0', nullable=False)

    score = sa.Column('score', sa.Float, default=0.0, nullable=False)

    def to_dict(self):
        return {
            'uuid': str(self.uuid),
            'url': self.url,
            'lastModified': self.last_modified,
            'expires': self.expires,
            'score': self.score
        }

    def __str__(self):
        return str(self.uuid)

    def __repr__(self):
        return str(self)

    def get_violations_per_day(self, db):
        from holmes.models import Review, Violation  # Prevent circular dependency

        violations = db \
            .query(
                sa.func.year(Review.completed_date).label('year'),
                sa.func.month(Review.completed_date).label('month'),
                sa.func.day(Review.completed_date).label('day'),
                sa.func.count(Violation.id).label('violation_count'),
                sa.func.sum(Violation.points).label('violation_points')
            ).join(
                Page, Page.id == Review.page_id
            ).join(
                Violation, Violation.review_id == Review.id
            ).filter(Review.is_complete == True).filter(Review.page_id == self.id) \
            .group_by(
                sa.func.year(Review.completed_date),
                sa.func.month(Review.completed_date),
                sa.func.day(Review.completed_date),
            ) \
            .order_by(Review.completed_date) \
            .all()

        result = []

        for day in violations:
            dt = "%d-%d-%d" % (day.year, day.month, day.day)
            result.append({
                "completedAt": dt,
                "violation_count": int(day.violation_count),
                "violation_points": int(day.violation_points)
            })

        return result

    @classmethod
    def by_uuid(cls, uuid, db):
        return db.query(Page).filter(Page.uuid == uuid).first()

    @classmethod
    def get_page_count(cls, db):
        return int(db.query(sa.func.count(Page.id)).scalar())

    @classmethod
    def update_scores(cls, individual_score, db):
        for i in range(3):
            db.begin(subtransactions=True)
            try:
                db.query(Page).update({'score': Page.score + individual_score})
                db.flush()
                db.commit()
                break
            except Exception:
                err = sys.exc_info()[1]
                if 'Deadlock found' in str(err):
                    logging.error('Deadlock happened! Trying again (try number %d)! (Details: %s)' % (i, str(err)))
                else:
                    db.rollback()
                    raise

    @classmethod
    def get_next_job(cls, db, expiration, cache, lock_expiration):
        from holmes.models import Settings, Domain

        expired_time = datetime.now() - timedelta(seconds=expiration)

        settings = Settings.instance(db)

        active_domains = [item.id for item in db.query(Domain.id).filter(Domain.is_active).all()]

        if not active_domains:
            return None

        pages_query = db.query(Page.uuid, Page.url, Page.score, Page.last_review_date) \
                        .filter(Page.last_review_date == None) \
                        .filter(Page.domain_id.in_(active_domains))

        pages_in_need_of_review = pages_query.order_by(Page.score.desc())[:200]

        if len(pages_in_need_of_review) == 0:
            pages_query = db.query(Page.uuid, Page.url, Page.score, Page.last_review_date) \
                            .filter(Page.last_review_date <= expired_time) \
                            .filter(Page.domain_id.in_(active_domains))

            pages_in_need_of_review = pages_query.order_by(Page.score.desc())[:200]

            if len(pages_in_need_of_review) == 0:
                if settings.lambda_score > 0:
                    cls.update_pages_score_by(settings, settings.lambda_score, db)
                return None

        if settings.lambda_score > pages_in_need_of_review[0].score:
            cls.update_pages_score_by(settings, settings.lambda_score, db)
            return None

        page = choice(pages_in_need_of_review)

        for i in range(10):
            lock = cache.has_next_job_lock(page.url, lock_expiration)
            if lock is None:
                page = choice(pages_in_need_of_review)
            else:
                break

        if lock is None:
            return None

        return {
            'page': str(page.uuid),
            'url': page.url,
            'score': page.score,
            'lock': lock
        }

    @classmethod
    def update_pages_score_by(cls, settings, score, db):
        settings.lambda_score = 0
        page_count = cls.get_page_count(db)
        individual_score = float(score) / float(page_count)
        cls.update_scores(individual_score, db)

    @classmethod
    @return_future
    def add_page(cls, db, cache, url, score, fetch_method, publish_method, config, callback):
        domain_name, domain_url = get_domain_from_url(url)
        if not url or not domain_name:
            callback((False, url, {
                'reason': 'invalid_url',
                'url': url,
                'status': None,
                'details': 'Domain name could not be determined.'
            }))
            return

        logging.info('Obtaining "%s"...' % url)

        fetch_method(
            url,
            cls.handle_request(cls.handle_add_page(db, cache, url, score, publish_method, callback)),
            proxy_host=config.HTTP_PROXY_HOST,
            proxy_port=config.HTTP_PROXY_PORT
        )

    @classmethod
    def handle_request(cls, callback):
        def handle(*args, **kw):
            response = args[-1]  # supports (url, response) and just response

            if hasattr(response, 'status_code'):
                status_code = response.status_code
            elif hasattr(response, 'code'):
                status_code = response.code
            else:
                status_code = 400

            if hasattr(response, 'body'):
                text = response.body
            elif hasattr(response, 'text'):
                text = response.text
            else:
                text = 'Empty response.text'

            callback(status_code, text, response.effective_url)

        return handle

    @classmethod
    def handle_add_page(cls, db, cache, url, score, publish_method, callback):
        def handle(code, body, effective_url):
            if code > 399:
                callback((False, url, {
                    'reason': 'invalid_url',
                    'url': url,
                    'status': code,
                    'details': body
                }))
                return

            if effective_url != url:
                callback((False, url, {
                    'reason': 'redirect',
                    'url': url,
                    'effectiveUrl': effective_url
                }))
                return

            domain = cls.add_domain(url, db, publish_method)
            page_uuid = cls.insert_or_update_page(url, score, domain, db, publish_method, cache)

            callback((True, url, page_uuid))

        return handle

    @classmethod
    def insert_or_update_page(cls, url, score, domain, db, publish_method, cache):
        url = url.encode('utf-8')
        url_hash = hashlib.sha512(url).hexdigest()
        page = db.query(Page).filter(Page.url_hash==url_hash).first()

        if page:
            for i in range(3):
                db.begin(subtransactions=True)
                try:
                    db.query(Page).filter(Page.id == page.id).update({'score': Page.score + score})
                    db.flush()
                    db.commit()
                    break
                except Exception:
                    err = sys.exc_info()[1]
                    if 'Deadlock found' in str(err):
                        logging.error('Deadlock happened! Trying again (try number %d)! (Details: %s)' % (i, str(err)))
                    else:
                        db.rollback()
                        raise

            return page.uuid

        db.begin(subtransactions=True)
        try:
            page = Page(url=url, url_hash=url_hash, domain=domain, score=score)
            db.add(page)
            db.flush()
            db.commit()
            cache.increment_page_count(domain)
            cache.increment_page_count()
        except Exception:
            db.rollback()
            err = sys.exc_info()[1]
            if 'Duplicate entry' in str(err):
                logging.error('Duplicate entry! (Details: %s)' % str(err))
            else:
                raise

        publish_method(dumps({
            'type': 'new-page',
            'pageUrl': str(url)
        }))

        return page.uuid

    @classmethod
    def add_domain(cls, url, db, publish_method):
        from holmes.models import Domain

        domain_name, domain_url = get_domain_from_url(url)

        domains = db.query(Domain).filter(or_(
            Domain.name == domain_name,
            Domain.name == domain_name.rstrip('/'),
            Domain.name == "%s/" % domain_name
        )).all()

        if not domains:
            domain = None
        else:
            domain = domains[0]

        if not domain:
            url_hash = hashlib.sha512(domain_url).hexdigest()
            domain = Domain(url=domain_url, url_hash=url_hash, name=domain_name)
            db.add(domain)
            db.flush()
            db.commit()

            publish_method(dumps({
                'type': 'new-domain',
                'domainUrl': str(domain_url)
            }))

        return domain
<|endoftext|>


<commit_before>import os
import sys
import tempfile
import traceback
from os.path import dirname, isdir, join, realpath

from aspen import resources, sockets
try:
    from nose.tools import with_setup
except ImportError:
    with_setup = None
    nose_tb = traceback.format_exc()

CWD = os.getcwd()
FSFIX = os.path.realpath(os.path.join(tempfile.gettempdir(), 'fsfix'))

def convert_path(path):
    """Given a Unix path, convert it for the current platform.
    """
    return os.sep.join(path.split('/'))

def convert_paths(paths):
    """Given a tuple of Unix paths, convert them for the current platform.
    """
    return tuple([convert_path(p) for p in paths])

def mk(*treedef):
    """Given a treedef, build a filesystem fixture in FSFIX.

    treedef is a sequence of strings and tuples. If a string, it is interpreted
    as a path to a directory that should be created. If a tuple, the first
    element is a path to a file, the second is the contents of the file. We do
    it this way to ease cross-platform testing.

    """
    root = FSFIX
    os.mkdir(root)
    for item in treedef:
        if isinstance(item, basestring):
            path = convert_path(item.lstrip('/'))
            path = os.sep.join([root, path])
            os.makedirs(path)
        elif isinstance(item, tuple):
            filepath, contents = item
            path = convert_path(filepath.lstrip('/'))
            path = os.sep.join([root, path])
            parent = dirname(path)
            if not isdir(parent):
                os.makedirs(parent)
            file(path, 'w').write(contents)

def fix(path=''):
    """Given a relative path, return an absolute path under FSFIX.

    The incoming path is in UNIX form (/foo/bar.html). The outgoing path is in
    native form, with symlinks removed.

    """
    path = os.sep.join([FSFIX] + path.split('/'))
    return realpath(path)

def rm():
    """Remove the filesystem fixture at FSFIX.
    """
    root = FSFIX
    if isdir(root):
        for root, dirs, files in os.walk(root, topdown=False):
            for name in dirs:
                os.rmdir(join(root, name))
            for name in files:
                os.remove(join(root, name))
        os.rmdir(root)

def teardown():
    """Standard teardown function.

    - reset the current working directory
    - remove FSFIX = %{tempdir}/fsfix
    - reset Aspen's global state
    - remove '.aspen' from sys.path
    - remove 'foo' from sys.modules
    - clear out sys.path_importer_cache

    """
    os.chdir(CWD)
    rm()
    # Reset some process-global caches. Hrm ...
    resources.__cache__ = {}
    sockets.__sockets__ = {}
    sockets.__channels__ = {}
    sys.path_importer_cache = {} # see test_weird.py
    if 'fsfix' in sys.path[0]:
        sys.path = sys.path[1:]
    if 'foo' in sys.modules:
        del sys.modules['foo']

teardown() # start clean

def attach_teardown(context):
    """Given a namespace dictionary, attach the teardown function.

    This function is designed for nose and will raise NotImplementedError at
    runtime with an additional ImportError traceback if nose is not available.
    To use it put this line at the end of your test modules:

        attach_teardown(globals())

    """
    if with_setup is None:
        raise NotImplementedError("This function is designed for nose, which "
                                  "couldn't be imported:" + os.sep + nose_tb)
    for name, func in context.items():
        if name.startswith('test_'):
            func = with_setup(teardown=teardown)(func) # non-destructive

def torndown(func):
    func.teardown = teardown
    return func

def path(*parts):
    """Given relative path parts, convert to absolute path on the filesystem.
    """
    return realpath(join(*parts))
<commit_msg>replace custom rmtree code with shutils.rmtree<commit_after>import os
import sys
import shutil
import tempfile
import traceback
from os.path import dirname, isdir, join, realpath
from aspen import resources, sockets
try:
    from nose.tools import with_setup
except ImportError:
    with_setup = None
    nose_tb = traceback.format_exc()

CWD = os.getcwd()
FSFIX = os.path.realpath(os.path.join(tempfile.gettempdir(), 'fsfix'))

def convert_path(path):
    """Given a Unix path, convert it for the current platform.
    """
    return os.sep.join(path.split('/'))

def convert_paths(paths):
    """Given a tuple of Unix paths, convert them for the current platform.
    """
    return tuple([convert_path(p) for p in paths])

def mk(*treedef):
    """Given a treedef, build a filesystem fixture in FSFIX.

    treedef is a sequence of strings and tuples. If a string, it is interpreted
    as a path to a directory that should be created. If a tuple, the first
    element is a path to a file, the second is the contents of the file. We do
    it this way to ease cross-platform testing.

    """
    root = FSFIX
    os.mkdir(root)
    for item in treedef:
        if isinstance(item, basestring):
            path = convert_path(item.lstrip('/'))
            path = os.sep.join([root, path])
            os.makedirs(path)
        elif isinstance(item, tuple):
            filepath, contents = item
            path = convert_path(filepath.lstrip('/'))
            path = os.sep.join([root, path])
            parent = dirname(path)
            if not isdir(parent):
                os.makedirs(parent)
            file(path, 'w').write(contents)

def fix(path=''):
    """Given a relative path, return an absolute path under FSFIX.

    The incoming path is in UNIX form (/foo/bar.html). The outgoing path is in
    native form, with symlinks removed.

    """
    path = os.sep.join([FSFIX] + path.split('/'))
    return realpath(path)

def rm():
    """Remove the filesystem fixture at FSFIX.
    """
    if isdir(FSFIX):
        shutil.rmtree(FSFIX)

def teardown():
    """Standard teardown function.

    - reset the current working directory
    - remove FSFIX = %{tempdir}/fsfix
    - reset Aspen's global state
    - remove '.aspen' from sys.path
    - remove 'foo' from sys.modules
    - clear out sys.path_importer_cache

    """
    os.chdir(CWD)
    rm()
    # Reset some process-global caches. Hrm ...
    resources.__cache__ = {}
    sockets.__sockets__ = {}
    sockets.__channels__ = {}
    sys.path_importer_cache = {} # see test_weird.py
    if 'fsfix' in sys.path[0]:
        sys.path = sys.path[1:]
    if 'foo' in sys.modules:
        del sys.modules['foo']

teardown() # start clean

def attach_teardown(context):
    """Given a namespace dictionary, attach the teardown function.

    This function is designed for nose and will raise NotImplementedError at
    runtime with an additional ImportError traceback if nose is not available.
    To use it put this line at the end of your test modules:

        attach_teardown(globals())

    """
    if with_setup is None:
        raise NotImplementedError("This function is designed for nose, which "
                                  "couldn't be imported:" + os.sep + nose_tb)
    for name, func in context.items():
        if name.startswith('test_'):
            func = with_setup(teardown=teardown)(func) # non-destructive

def torndown(func):
    func.teardown = teardown
    return func

def path(*parts):
    """Given relative path parts, convert to absolute path on the filesystem.
    """
    return realpath(join(*parts))
<|endoftext|>