<commit_before>##############################################################################
# Copyright (c) 2013-2016, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://github.com/llnl/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License (as
# published by the Free Software Foundation) version 2.1, February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
import argparse
import codecs
import functools
import os
import platform
import time
import xml.dom.minidom
import xml.etree.ElementTree as ET

import llnl.util.filesystem as fs
import llnl.util.tty as tty
import spack
import spack.cmd
import spack.cmd.common.arguments as arguments
from spack.build_environment import InstallError
from spack.fetch_strategy import FetchError
from spack.package import PackageBase

description = "Build and install packages"

def setup_parser(subparser):
    subparser.add_argument(
        '--only',
        default='package,dependencies',
        dest='things_to_install',
        choices=['package', 'dependencies'],
        help="""Select the mode of installation.
The default is to install the package along with all its dependencies.
Alternatively one can decide to install only the package or only
the dependencies."""
    )
    subparser.add_argument(
        '-j', '--jobs', action='store', type=int,
        help="Explicitly set number of make jobs.  Default is #cpus.")
    subparser.add_argument(
        '--keep-prefix', action='store_true', dest='keep_prefix',
        help="Don't remove the install prefix if installation fails.")
    subparser.add_argument(
        '--keep-stage', action='store_true', dest='keep_stage',
        help="Don't remove the build stage if installation succeeds.")
    subparser.add_argument(
        '-n', '--no-checksum', action='store_true', dest='no_checksum',
        help="Do not check packages against checksum")
    subparser.add_argument(
        '-v', '--verbose', action='store_true', dest='verbose',
        help="Display verbose build output while installing.")
    subparser.add_argument(
        '--fake', action='store_true', dest='fake',
        help="Fake install. Just remove prefix and create a fake file.")

    cd_group = subparser.add_mutually_exclusive_group()
    arguments.add_common_arguments(cd_group, ['clean', 'dirty'])

    subparser.add_argument(
        'package',
        nargs=argparse.REMAINDER,
        help="spec of the package to install"
    )
    subparser.add_argument(
        '--run-tests', action='store_true', dest='run_tests',
        help="Run package level tests during installation."
    )
    subparser.add_argument(
        '--log-format',
        default=None,
        choices=['junit'],
        help="Format to be used for log files."
    )
    subparser.add_argument(
        '--log-file',
        default=None,
        help="Filename for the log file. If not passed a default will be used."
    )

# Needed for test cases
class TestResult(object):
    PASSED = 0
    FAILED = 1
    SKIPPED = 2
    ERRORED = 3

class TestSuite(object):
    def __init__(self, spec):
        self.root = ET.Element('testsuite')
        self.tests = []
        self.spec = spec

    def append(self, item):
        if not isinstance(item, TestCase):
            raise TypeError(
                'only TestCase instances may be appended to TestSuite'
            )
        self.tests.append(item)  # Append the item to the list of tests

    def dump(self, filename):
        # Prepare the header for the entire test suite
        number_of_errors = sum(
            x.result_type == TestResult.ERRORED for x in self.tests
        )
        self.root.set('errors', str(number_of_errors))
        number_of_failures = sum(
            x.result_type == TestResult.FAILED for x in self.tests
        )
        self.root.set('failures', str(number_of_failures))
        self.root.set('tests', str(len(self.tests)))
        self.root.set('name', str(self.spec))
        self.root.set('hostname', platform.node())

        for item in self.tests:
            self.root.append(item.element)

        with codecs.open(filename, 'wb', 'utf-8') as file:
            xml_string = ET.tostring(self.root)
            xml_string = xml.dom.minidom.parseString(xml_string).toprettyxml()
            file.write(xml_string)

class TestCase(object):

    results = {
        TestResult.PASSED: None,
        TestResult.SKIPPED: 'skipped',
        TestResult.FAILED: 'failure',
        TestResult.ERRORED: 'error',
    }

    def __init__(self, classname, name):
        self.element = ET.Element('testcase')
        self.element.set('classname', str(classname))
        self.element.set('name', str(name))
        self.result_type = None

    def set_duration(self, duration):
        self.element.set('time', str(duration))

    def set_result(self, result_type,
                   message=None, error_type=None, text=None):
        self.result_type = result_type
        result = TestCase.results[self.result_type]
        if result is not None and result is not TestResult.PASSED:
            subelement = ET.SubElement(self.element, result)
            if error_type is not None:
                subelement.set('type', error_type)
            if message is not None:
                subelement.set('message', str(message))
            if text is not None:
                subelement.text = text

def fetch_text(path):
    if not os.path.exists(path):
        return ''

    with codecs.open(path, 'rb', 'utf-8') as f:
        return '\n'.join(
            list(line.strip() for line in f.readlines())
        )

def junit_output(spec, test_suite):
    # Cycle once and for all on the dependencies and skip
    # the ones that are already installed. This ensures that
    # for the same spec, the same number of entries will be
    # displayed in the XML report
    for x in spec.traverse(order='post'):
        package = spack.repo.get(x)
        if package.installed:
            test_case = TestCase(package.name, x.short_spec)
            test_case.set_duration(0.0)
            test_case.set_result(
                TestResult.SKIPPED,
                message='Skipped [already installed]',
                error_type='already_installed'
            )
            test_suite.append(test_case)

    def decorator(func):
        @functools.wraps(func)
        def wrapper(self, *args, ** kwargs):

            # Check if the package has been installed already
            if self.installed:
                return

            test_case = TestCase(self.name, self.spec.short_spec)
            # Try to install the package
            try:
                # If already installed set the spec as skipped
                start_time = time.time()
                # PackageBase.do_install
                func(self, *args, **kwargs)
                duration = time.time() - start_time
                test_case.set_duration(duration)
                test_case.set_result(TestResult.PASSED)
            except InstallError:
                # Check if the package relies on dependencies that
                # did not install
                duration = time.time() - start_time
                test_case.set_duration(duration)
                if [x for x in self.spec.dependencies(('link', 'run')) if not spack.repo.get(x).installed]:  # NOQA: ignore=E501
                    test_case.set_duration(0.0)
                    test_case.set_result(
                        TestResult.SKIPPED,
                        message='Skipped [failed dependencies]',
                        error_type='dep_failed'
                    )
                else:
                    # An InstallError is considered a failure (the recipe
                    # didn't work correctly)
                    text = fetch_text(self.build_log_path)
                    test_case.set_result(
                        TestResult.FAILED,
                        message='Installation failure',
                        text=text
                    )
            except FetchError:
                # A FetchError is considered an error as
                # we didn't even start building
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.ERRORED,
                    message='Unable to fetch package',
                    text=text
                )
            except Exception:
                # Anything else is also an error
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.ERRORED,
                    message='Unexpected exception thrown during install',
                    text=text
                )
            except:
                # Anything else is also an error
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.ERRORED,
                    message='Unknown error',
                    text=text
                )

            # Try to get the log
            test_suite.append(test_case)
        return wrapper
    return decorator

def default_log_file(spec):
    """Computes the default filename for the log file and creates
    the corresponding directory if not present
    """
    fmt = 'test-{x.name}-{x.version}-{hash}.xml'
    basename = fmt.format(x=spec, hash=spec.dag_hash())
    dirname = fs.join_path(spack.var_path, 'junit-report')
    fs.mkdirp(dirname)
    return fs.join_path(dirname, basename)

def install(parser, args, **kwargs):
    if not args.package:
        tty.die("install requires at least one package argument")

    if args.jobs is not None:
        if args.jobs <= 0:
            tty.die("The -j option must be a positive integer!")

    if args.no_checksum:
        spack.do_checksum = False        # TODO: remove this global.

    # Parse cli arguments and construct a dictionary
    # that will be passed to Package.do_install API
    kwargs.update({
        'keep_prefix': args.keep_prefix,
        'keep_stage': args.keep_stage,
        'install_deps': 'dependencies' in args.things_to_install,
        'make_jobs': args.jobs,
        'run_tests': args.run_tests,
        'verbose': args.verbose,
        'fake': args.fake,
        'dirty': args.dirty
    })

    # Spec from cli
    specs = spack.cmd.parse_specs(args.package, concretize=True)
    if len(specs) != 1:
        tty.error('only one spec can be installed at a time.')
    spec = specs.pop()

    # Check if we were asked to produce some log for dashboards
    if args.log_format is not None:
        # Compute the filename for logging
        log_filename = args.log_file
        if not log_filename:
            log_filename = default_log_file(spec)
        # Create the test suite in which to log results
        test_suite = TestSuite(spec)
        # Decorate PackageBase.do_install to get installation status
        PackageBase.do_install = junit_output(
            spec, test_suite
        )(PackageBase.do_install)

    # Do the actual installation
    if args.things_to_install == 'dependencies':
        # Install dependencies as-if they were installed
        # for root (explicit=False in the DB)
        kwargs['explicit'] = False
        for s in spec.dependencies():
            p = spack.repo.get(s)
            p.do_install(**kwargs)
    else:
        package = spack.repo.get(spec)
        kwargs['explicit'] = True
        package.do_install(**kwargs)

    # Dump log file if asked to
    if args.log_format is not None:
        test_suite.dump(log_filename)
<commit_msg>CDash fixes (#2438)<commit_after>##############################################################################
# Copyright (c) 2013-2016, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://github.com/llnl/spack
# Please also see the LICENSE file for our notice and the LGPL.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License (as
# published by the Free Software Foundation) version 2.1, February 1999.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
# conditions of the GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
##############################################################################
import argparse
import codecs
import functools
import os
import platform
import time
import xml.dom.minidom
import xml.etree.ElementTree as ET

import llnl.util.filesystem as fs
import llnl.util.tty as tty
import spack
import spack.cmd
import spack.cmd.common.arguments as arguments
from spack.build_environment import InstallError
from spack.fetch_strategy import FetchError
from spack.package import PackageBase

description = "Build and install packages"

def setup_parser(subparser):
    subparser.add_argument(
        '--only',
        default='package,dependencies',
        dest='things_to_install',
        choices=['package', 'dependencies'],
        help="""Select the mode of installation.
The default is to install the package along with all its dependencies.
Alternatively one can decide to install only the package or only
the dependencies."""
    )
    subparser.add_argument(
        '-j', '--jobs', action='store', type=int,
        help="Explicitly set number of make jobs.  Default is #cpus.")
    subparser.add_argument(
        '--keep-prefix', action='store_true', dest='keep_prefix',
        help="Don't remove the install prefix if installation fails.")
    subparser.add_argument(
        '--keep-stage', action='store_true', dest='keep_stage',
        help="Don't remove the build stage if installation succeeds.")
    subparser.add_argument(
        '-n', '--no-checksum', action='store_true', dest='no_checksum',
        help="Do not check packages against checksum")
    subparser.add_argument(
        '-v', '--verbose', action='store_true', dest='verbose',
        help="Display verbose build output while installing.")
    subparser.add_argument(
        '--fake', action='store_true', dest='fake',
        help="Fake install. Just remove prefix and create a fake file.")

    cd_group = subparser.add_mutually_exclusive_group()
    arguments.add_common_arguments(cd_group, ['clean', 'dirty'])

    subparser.add_argument(
        'package',
        nargs=argparse.REMAINDER,
        help="spec of the package to install"
    )
    subparser.add_argument(
        '--run-tests', action='store_true', dest='run_tests',
        help="Run package level tests during installation."
    )
    subparser.add_argument(
        '--log-format',
        default=None,
        choices=['junit'],
        help="Format to be used for log files."
    )
    subparser.add_argument(
        '--log-file',
        default=None,
        help="Filename for the log file. If not passed a default will be used."
    )

# Needed for test cases
class TestResult(object):
    PASSED = 0
    FAILED = 1
    SKIPPED = 2
    ERRORED = 3

class TestSuite(object):
    def __init__(self, spec):
        self.root = ET.Element('testsuite')
        self.tests = []
        self.spec = spec

    def append(self, item):
        if not isinstance(item, TestCase):
            raise TypeError(
                'only TestCase instances may be appended to TestSuite'
            )
        self.tests.append(item)  # Append the item to the list of tests

    def dump(self, filename):
        # Prepare the header for the entire test suite
        number_of_errors = sum(
            x.result_type == TestResult.ERRORED for x in self.tests
        )
        self.root.set('errors', str(number_of_errors))
        number_of_failures = sum(
            x.result_type == TestResult.FAILED for x in self.tests
        )
        self.root.set('failures', str(number_of_failures))
        self.root.set('tests', str(len(self.tests)))
        self.root.set('name', self.spec.short_spec)
        self.root.set('hostname', platform.node())

        for item in self.tests:
            self.root.append(item.element)

        with codecs.open(filename, 'wb', 'utf-8') as file:
            xml_string = ET.tostring(self.root)
            xml_string = xml.dom.minidom.parseString(xml_string).toprettyxml()
            file.write(xml_string)

class TestCase(object):

    results = {
        TestResult.PASSED: None,
        TestResult.SKIPPED: 'skipped',
        TestResult.FAILED: 'failure',
        TestResult.ERRORED: 'error',
    }

    def __init__(self, classname, name):
        self.element = ET.Element('testcase')
        self.element.set('classname', str(classname))
        self.element.set('name', str(name))
        self.result_type = None

    def set_duration(self, duration):
        self.element.set('time', str(duration))

    def set_result(self, result_type,
                   message=None, error_type=None, text=None):
        self.result_type = result_type
        result = TestCase.results[self.result_type]
        if result is not None and result is not TestResult.PASSED:
            subelement = ET.SubElement(self.element, result)
            if error_type is not None:
                subelement.set('type', error_type)
            if message is not None:
                subelement.set('message', str(message))
            if text is not None:
                subelement.text = text

def fetch_text(path):
    if not os.path.exists(path):
        return ''

    with codecs.open(path, 'rb', 'utf-8') as f:
        return '\n'.join(
            list(line.strip() for line in f.readlines())
        )

def junit_output(spec, test_suite):
    # Cycle once and for all on the dependencies and skip
    # the ones that are already installed. This ensures that
    # for the same spec, the same number of entries will be
    # displayed in the XML report
    for x in spec.traverse(order='post'):
        package = spack.repo.get(x)
        if package.installed:
            test_case = TestCase(package.name, x.short_spec)
            test_case.set_duration(0.0)
            test_case.set_result(
                TestResult.SKIPPED,
                message='Skipped [already installed]',
                error_type='already_installed'
            )
            test_suite.append(test_case)

    def decorator(func):
        @functools.wraps(func)
        def wrapper(self, *args, ** kwargs):

            # Check if the package has been installed already
            if self.installed:
                return

            test_case = TestCase(self.name, self.spec.short_spec)
            # Try to install the package
            try:
                # If already installed set the spec as skipped
                start_time = time.time()
                # PackageBase.do_install
                func(self, *args, **kwargs)
                duration = time.time() - start_time
                test_case.set_duration(duration)
                test_case.set_result(TestResult.PASSED)
            except InstallError:
                # Check if the package relies on dependencies that
                # did not install
                duration = time.time() - start_time
                test_case.set_duration(duration)
                if [x for x in self.spec.dependencies(('link', 'run')) if not spack.repo.get(x).installed]:  # NOQA: ignore=E501
                    test_case.set_duration(0.0)
                    test_case.set_result(
                        TestResult.SKIPPED,
                        message='Skipped [failed dependencies]',
                        error_type='dep_failed'
                    )
                else:
                    # An InstallError is considered a failure (the recipe
                    # didn't work correctly)
                    text = fetch_text(self.build_log_path)
                    test_case.set_result(
                        TestResult.FAILED,
                        message='Installation failure',
                        text=text
                    )
            except FetchError:
                # A FetchError is considered an error as
                # we didn't even start building
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.FAILED,
                    message='Unable to fetch package',
                    text=text
                )
            except Exception:
                # Anything else is also an error
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.FAILED,
                    message='Unexpected exception thrown during install',
                    text=text
                )
            except:
                # Anything else is also an error
                duration = time.time() - start_time
                test_case.set_duration(duration)
                text = fetch_text(self.build_log_path)
                test_case.set_result(
                    TestResult.FAILED,
                    message='Unknown error',
                    text=text
                )

            # Try to get the log
            test_suite.append(test_case)
        return wrapper
    return decorator

def default_log_file(spec):
    """Computes the default filename for the log file and creates
    the corresponding directory if not present
    """
    fmt = 'test-{x.name}-{x.version}-{hash}.xml'
    basename = fmt.format(x=spec, hash=spec.dag_hash())
    dirname = fs.join_path(spack.var_path, 'junit-report')
    fs.mkdirp(dirname)
    return fs.join_path(dirname, basename)

def install(parser, args, **kwargs):
    if not args.package:
        tty.die("install requires at least one package argument")

    if args.jobs is not None:
        if args.jobs <= 0:
            tty.die("The -j option must be a positive integer!")

    if args.no_checksum:
        spack.do_checksum = False        # TODO: remove this global.

    # Parse cli arguments and construct a dictionary
    # that will be passed to Package.do_install API
    kwargs.update({
        'keep_prefix': args.keep_prefix,
        'keep_stage': args.keep_stage,
        'install_deps': 'dependencies' in args.things_to_install,
        'make_jobs': args.jobs,
        'run_tests': args.run_tests,
        'verbose': args.verbose,
        'fake': args.fake,
        'dirty': args.dirty
    })

    # Spec from cli
    specs = spack.cmd.parse_specs(args.package, concretize=True)
    if len(specs) != 1:
        tty.error('only one spec can be installed at a time.')
    spec = specs.pop()

    # Check if we were asked to produce some log for dashboards
    if args.log_format is not None:
        # Compute the filename for logging
        log_filename = args.log_file
        if not log_filename:
            log_filename = default_log_file(spec)
        # Create the test suite in which to log results
        test_suite = TestSuite(spec)
        # Decorate PackageBase.do_install to get installation status
        PackageBase.do_install = junit_output(
            spec, test_suite
        )(PackageBase.do_install)

    # Do the actual installation
    if args.things_to_install == 'dependencies':
        # Install dependencies as-if they were installed
        # for root (explicit=False in the DB)
        kwargs['explicit'] = False
        for s in spec.dependencies():
            p = spack.repo.get(s)
            p.do_install(**kwargs)
    else:
        package = spack.repo.get(spec)
        kwargs['explicit'] = True
        package.do_install(**kwargs)

    # Dump log file if asked to
    if args.log_format is not None:
        test_suite.dump(log_filename)
<|endoftext|>


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

#TODO: Support references
#TODO: Support inline schema

import re, copy

class JSONSchemaValidator:
  '''
  Implementation of the json-schema validator that adheres to the
  JSON Schema Proposal 2nd Draft.
  '''

  # Map of schema types to their equivalent in the python types module
  _typesmap = {
    "string": [str, unicode],
    "integer": int,
    "number": [int, float],
    "boolean": bool,
    "object": dict,
    "array": list,
    "null": type(None),
    "any": None
  }

  _ignored = ('identity', 'options', 'readonly', 'transient', 'hidden')
  # extends not in here b/c it should raise error

  # Default schema property values.
  _schemadefault = {
    "id": None,
    "type": None,
    "properties": None,
    "items": None,
    "optional": False,
    "additionalProperties": None,
    "requires": None,
    "minimum": None,
    "maximum": None,
    "minItems": None,
    "maxItems": None,
    "pattern": None,
    "maxLength": None,
    "minLength": None,
    "enum": None,
    "title": None,
    "description": None,
    "format": None,
    "default": None,
    "maxDecimal": None,
    "disallow": None,
  }

  _refmap = {}

  _interactive_mode = True

  def __init__(self, interactive_mode=True):
    self._interactive_mode = interactive_mode

  def validate_id(self, x, fieldname, schema, ID=None):
    '''
    Validates a schema id and adds it to the schema reference map
    '''
    if ID is not None:
      if ID == "$":
        raise ValueError("Reference id for field '%s' cannot equal '$'" % fieldname)
      self._refmap[ID] = schema
    return x

  def validate_type(self, x, fieldname, schema, fieldtype=None):
    '''
    Validates that the fieldtype specified is correct for the given
    data
    '''

    converted_fieldtype = self._convert_type(fieldtype)

    # We need to know if the field exists or if it's just Null
    fieldexists = True
    try:
      value = x[fieldname]
    except KeyError:
      fieldexists = False
    finally:
      value = x.get(fieldname)

    if converted_fieldtype is not None and fieldexists:
      if isinstance(converted_fieldtype, list):
        # Match if type matches any one of the types in the list
        datavalid = False
        for eachtype in converted_fieldtype:
          try:
            self.validate_type(x, fieldname, eachtype, eachtype)
            datavalid = True
            break
          except ValueError:
            pass
        if not datavalid:
          raise ValueError("Value %r for field '%s' is not of type %s" % (value, fieldname, fieldtype))
      elif isinstance(converted_fieldtype, dict):
        try:
          self.__validate(fieldname, x, converted_fieldtype)
        except ValueError,e:
          raise e
      else:
        if type(value) != converted_fieldtype:
          raise ValueError("Value %r for field '%s' is not of type %s" % (value, fieldname, fieldtype))
    return x

  def validate_properties(self, x, fieldname, schema, properties=None):
    '''
    Validates properties of a JSON object by processing the object's
    schema recursively
    '''
    if properties is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if isinstance(value, dict):
        if isinstance(properties, dict):
          for eachProp in properties.keys():
            self.__validate(eachProp, value, properties.get(eachProp))
        else:
          raise ValueError("Properties definition of field '%s' is not an object" % fieldname)
    return x

  def validate_items(self, x, fieldname, schema, items=None):
    '''
    Validates that all items in the list for the given field match the
    given schema
    '''
    if items is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if isinstance(value, list):
        if isinstance(items, list):
          if len(items) == len(value):
            for itemIndex in range(len(items)):
              try:
                self.validate(value[itemIndex], items[itemIndex])
              except ValueError, e:
                raise ValueError("Failed to validate field '%s' list schema: %s" % (fieldname, e))
          else:
            raise ValueError("Length of list %r for field '%s' is not equal to length of schema list" % (value, fieldname))
        elif isinstance(items, dict):
          for eachItem in value:
              try:
                self._validate(eachItem, items)
              except ValueError, e:
                raise ValueError("Failed to validate field '%s' list schema: %s" % (fieldname, e))
        else:
          raise ValueError("Properties definition of field '%s' is not a list or an object" % fieldname)
    return x

  def validate_optional(self, x, fieldname, schema, optional=False):
    '''
    Validates that the given field is present if optional is false
    '''
    # Make sure the field is present
    if fieldname not in x.keys() and not optional:
      raise ValueError("Required field '%s' is missing" % fieldname)
    return x

  def validate_additionalProperties(self, x, fieldname, schema, additionalProperties=None):
    '''
    Validates additional properties of a JSON object that were not
    specifically defined by the properties property
    '''
    if additionalProperties is not None:
      # If additionalProperties is the boolean value True then we accept any
      # additional properties.
      if isinstance(additionalProperties, bool) and additionalProperties:
        return x

      value = x.get(fieldname)
      if isinstance(additionalProperties, (dict, bool)):
        properties = schema.get("properties")
        if properties is None:
          properties = {}
        for eachProperty in value.keys():
          if eachProperty not in properties:
            # If additionalProperties is the boolean value False then we
            # don't accept any additional properties.
            if isinstance(additionalProperties, bool) and not additionalProperties:
              raise ValueError("Additional properties not defined by 'properties' are not allowed in field '%s'" % fieldname)
            self.__validate(eachProperty, value, additionalProperties)
      else:
        raise ValueError("additionalProperties schema definition for field '%s' is not an object" % fieldname)
    return x

  def validate_requires(self, x, fieldname, schema, requires=None):
    if x.get(fieldname) is not None and requires is not None:
      if x.get(requires) is None:
        raise ValueError("Field '%s' is required by field '%s'" % (requires, fieldname))
    return x

  def validate_minimum(self, x, fieldname, schema, minimum=None):
    '''
    Validates that the field is longer than or equal to the minimum
    length if specified
    '''
    if minimum is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if value is not None:
        if type(value) in (int, float) and value < minimum:
          raise ValueError("Value %r for field '%s' is less than minimum value: %f" % (value, fieldname, minimum))
        elif isinstance(value, list) and len(value) < minimum:
          raise ValueError("Value %r for field '%s' has fewer values than the minimum: %f" % (value, fieldname, minimum))
    return x

  def validate_maximum(self, x, fieldname, schema, maximum=None):
    '''
    Validates that the field is shorter than or equal to the maximum
    length if specified.
    '''
    if maximum is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if value is not None:
        if type(value) in (int, float) and value > maximum:
          raise ValueError("Value %r for field '%s' is greater than maximum value: %f" % (value, fieldname, maximum))
        elif isinstance(value, list) and len(value) > maximum:
          raise ValueError("Value %r for field '%s' has more values than the maximum: %f" % (value, fieldname, maximum))
    return x

  def validate_minItems(self, x, fieldname, schema, minitems=None):
    '''
    Validates that the number of items in the given field is equal to or
    more than the minimum amount.
    '''
    if minitems is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if value is not None:
        if isinstance(value, list) and len(value) < minitems:
          raise ValueError("Value %r for field '%s' must have a minimum of %d items" % (fieldname, fieldname, minitems))
    return x

  def validate_maxItems(self, x, fieldname, schema, maxitems=None):
    '''
    Validates that the number of items in the given field is equal to or
    less than the maximum amount.
    '''
    if maxitems is not None and x.get(fieldname) is not None:
      value = x.get(fieldname)
      if value is not None:
        if isinstance(value, list) and len(value) > maxitems:
          raise ValueError("Value %r for field '%s' must have a maximum of %d items" % (value, fieldname, maxitems))
    return x

  def validate_pattern(self, x, fieldname, schema, pattern=None):
    '''
    Validates that the given field, if a string, matches the given
    regular expression.
    '''
    value = x.get(fieldname)
    if pattern is not None and \
       self._is_string_type(value):
      p = re.compile(pattern)
      if not p.match(value):
        raise ValueError("Value %r for field '%s' does not match regular expression '%s'" % (value, fieldname, pattern))
    return x

  def validate_maxLength(self, x, fieldname, schema, length=None):
    '''
    Validates that the value of the given field is shorter than or equal
    to the specified length if a string
    '''
    value = x.get(fieldname)
    if length is not None and \
       self._is_string_type(value) and \
       len(value) > length:
      raise ValueError("Length of value %r for field '%s' must be less than or equal to %f" % (value, fieldname, length))
    return x

  def validate_minLength(self, x, fieldname, schema, length=None):
    '''
    Validates that the value of the given field is longer than or equal
    to the specified length if a string
    '''
    value = x.get(fieldname)
    if length is not None and \
       self._is_string_type(value) and \
       len(value) < length:
      raise ValueError("Length of value %r for field '%s' must be more than or equal to %f" % (value, fieldname, length))
    return x

  def validate_enum(self, x, fieldname, schema, options=None):
    '''
    Validates that the value of the field is equal to one of the
    specified option values
    '''
    value = x.get(fieldname)
    if options is not None and value is not None:
      if not isinstance(options, list):
        raise ValueError("Enumeration %r for field '%s' is not a list type", (options, fieldname))
      if value not in options:
        raise ValueError("Value %r for field '%s' is not in the enumeration: %r" % (value, fieldname, options))
    return x

  def validate_title(self, x, fieldname, schema, title=None):
    if not self._is_string_type(title):
      raise ValueError("The title for field '%s' must be a string" % fieldname);
    return x

  def validate_description(self, x, fieldname, schema, description=None):
    if not self._is_string_type(description):
      raise ValueError("The description for field '%s' must be a string." % fieldname);
    return x

  def validate_format(self, x, fieldname, schema, format=None):
    '''
    Validates that the value of the field matches the predifined format
    specified.
    '''
    # No definitions are currently defined for formats
    return x

  def validate_default(self, x, fieldname, schema, default=None):
    '''
    Adds default data to the original json document if the document is
    not readonly
    '''
    if self._interactive_mode and fieldname not in x.keys() and default is not None:
      if not schema.get("readonly"):
        x[fieldname] = default
    return x

  def validate_maxDecimal(self, x, fieldname, schema, maxdecimal=None):
    '''
    Validates that the value of the given field has less than or equal
    to the maximum number of decimal places given
    '''
    value = x.get(fieldname)
    if maxdecimal is not None and value is not None:
      maxdecstring = str(value)
      if len(maxdecstring[maxdecstring.find(".")+1:]) > maxdecimal:
        raise ValueError("Value %r for field '%s' must not have more than %d decimal places" % (value, fieldname, maxdecimal))
    return x

  def validate_disallow(self, x, fieldname, schema, disallow=None):
    '''
    Validates that the value of the given field does not match the
    disallowed type.
    '''
    if disallow is not None:
      try:
        self.validate_type(x, fieldname, schema, disallow)
      except ValueError:
        return x
      raise ValueError("Value %r of type %s is disallowed for field '%s'" % (x.get(fieldname), disallow, fieldname))
    return x

  def _convert_type(self, fieldtype):
    if isinstance(fieldtype, (type, dict)):
      return fieldtype
    elif isinstance(fieldtype, list):
      converted_fields = []
      for subfieldtype in fieldtype:
        converted_fields.append(self._convert_type(subfieldtype))
      return converted_fields
    elif fieldtype is None:
      return None
    else:
      fieldtype = str(fieldtype)
      if fieldtype in self._typesmap.keys():
        return self._typesmap[fieldtype]
      else:
        raise ValueError("Field type '%s' is not supported." % fieldtype)

  def validate(self, data, schema):
    '''
    Validates a piece of json data against the provided json-schema.
    '''

    #TODO: Validate the schema object here.

    self._refmap = {
      '$': schema
    }
    # Wrap the data in a dictionary
    self._validate(data, schema)

  def _validate(self, data, schema):
    self.__validate("_data", {"_data": data}, schema)

  def __validate(self, fieldname, data, schema):

    if schema is not None:
      if not isinstance(schema, dict):
        raise ValueError("Schema structure is invalid.");

      # Produce a copy of the schema object since we will make changes to
      # it to process default values. Deep copy is not necessary since we will
      # produce a copy of sub items on the next recursive call.

      new_schema = copy.copy(schema)
      #Initialize defaults
      for schemaprop in self._schemadefault.keys():
        if schemaprop not in new_schema:
          new_schema[schemaprop] = self._schemadefault[schemaprop]

      for schemaprop in new_schema:

        validatorname = "validate_"+schemaprop

        try:
          validator = getattr(self, validatorname)
          # Pass the original schema object but the value of the property from
          # copy in order to validate default values.
          validator(data, fieldname, schema, new_schema.get(schemaprop))
        except AttributeError, e:
          if schemaprop not in self._ignored:
            raise ValueError("Schema property '%s' is not supported" %
                             schemaprop)

    return data

  def _is_string_type(self, value):
    return value is None or isinstance(value, basestring)

__all__ = [ 'JSONSchemaValidator' ]
<commit_msg>PEP8 (but lines still too long)<commit_after>#!/usr/bin/env python
#:coding=utf-8:

#TODO: Support references
#TODO: Support inline schema

import re
import copy

class JSONSchemaValidator(object):
    '''
    Implementation of the json-schema validator that adheres to the
    JSON Schema Proposal 2nd Draft.
    '''

    # Map of schema types to their equivalent in the python types module
    _typesmap = {
        "string": [str, unicode],
        "integer": int,
        "number": [int, float],
        "boolean": bool,
        "object": dict,
        "array": list,
        "null": type(None),
        "any": None,
    }

    _ignored = ('identity', 'options', 'readonly', 'transient', 'hidden')
    # extends not in here b/c it should raise error

    # Default schema property values.
    _schemadefault = {
        "id": None,
        "type": None,
        "properties": None,
        "items": None,
        "optional": False,
        "additionalProperties": None,
        "requires": None,
        "minimum": None,
        "maximum": None,
        "minItems": None,
        "maxItems": None,
        "pattern": None,
        "maxLength": None,
        "minLength": None,
        "enum": None,
        "title": None,
        "description": None,
        "format": None,
        "default": None,
        "maxDecimal": None,
        "disallow": None,
    }

    _refmap = {}

    _interactive_mode = True

    def __init__(self, interactive_mode=True):
        self._interactive_mode = interactive_mode

    def validate_id(self, x, fieldname, schema, ID=None):
        '''
        Validates a schema id and adds it to the schema reference map
        '''
        if ID is not None:
            if ID == "$":
                raise ValueError("Reference id for field '%s' cannot equal '$'" % fieldname)
            self._refmap[ID] = schema
        return x

    def validate_type(self, x, fieldname, schema, fieldtype=None):
        '''
        Validates that the fieldtype specified is correct for the given
        data
        '''

        converted_fieldtype = self._convert_type(fieldtype)

        # We need to know if the field exists or if it's just Null
        fieldexists = True
        try:
            value = x[fieldname]
        except KeyError:
            fieldexists = False
        finally:
            value = x.get(fieldname)

        if converted_fieldtype is not None and fieldexists:
            if isinstance(converted_fieldtype, list):
                # Match if type matches any one of the types in the list
                datavalid = False
                for eachtype in converted_fieldtype:
                    try:
                        self.validate_type(x, fieldname, eachtype, eachtype)
                        datavalid = True
                        break
                    except ValueError:
                        pass
                if not datavalid:
                    raise ValueError("Value %r for field '%s' is not of type %s" % (value, fieldname, fieldtype))
            elif isinstance(converted_fieldtype, dict):
                try:
                    self.__validate(fieldname, x, converted_fieldtype)
                except ValueError, e:
                    raise e
            else:
                if type(value) != converted_fieldtype:
                    raise ValueError("Value %r for field '%s' is not of type %s" % (value, fieldname, fieldtype))
        return x

    def validate_properties(self, x, fieldname, schema, properties=None):
        '''
        Validates properties of a JSON object by processing the object's
        schema recursively
        '''
        if properties is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if isinstance(value, dict):
                if isinstance(properties, dict):
                    for eachProp in properties.keys():
                        self.__validate(eachProp, value, properties.get(eachProp))
                else:
                    raise ValueError("Properties definition of field '%s' is not an object" % fieldname)
        return x

    def validate_items(self, x, fieldname, schema, items=None):
        '''
        Validates that all items in the list for the given field match the
        given schema
        '''
        if items is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if isinstance(value, list):
                if isinstance(items, list):
                    if len(items) == len(value):
                        for itemIndex in range(len(items)):
                            try:
                                self.validate(value[itemIndex], items[itemIndex])
                            except ValueError, e:
                                raise ValueError("Failed to validate field '%s' list schema: %s" % (fieldname, e))
                    else:
                        raise ValueError("Length of list %r for field '%s' is not equal to length of schema list" % (value, fieldname))
                elif isinstance(items, dict):
                    for eachItem in value:
                            try:
                                self._validate(eachItem, items)
                            except ValueError, e:
                                raise ValueError("Failed to validate field '%s' list schema: %s" % (fieldname, e))
                else:
                    raise ValueError("Properties definition of field '%s' is not a list or an object" % fieldname)
        return x

    def validate_optional(self, x, fieldname, schema, optional=False):
        '''
        Validates that the given field is present if optional is false
        '''
        # Make sure the field is present
        if fieldname not in x.keys() and not optional:
            raise ValueError("Required field '%s' is missing" % fieldname)
        return x

    def validate_additionalProperties(self, x, fieldname, schema, additionalProperties=None):
        '''
        Validates additional properties of a JSON object that were not
        specifically defined by the properties property
        '''
        if additionalProperties is not None:
            # If additionalProperties is the boolean value True then we accept any
            # additional properties.
            if isinstance(additionalProperties, bool) and additionalProperties:
                return x

            value = x.get(fieldname)
            if isinstance(additionalProperties, (dict, bool)):
                properties = schema.get("properties")
                if properties is None:
                    properties = {}
                for eachProperty in value.keys():
                    if eachProperty not in properties:
                        # If additionalProperties is the boolean value False then we
                        # don't accept any additional properties.
                        if isinstance(additionalProperties, bool) and not additionalProperties:
                            raise ValueError("Additional properties not defined by 'properties' are not allowed in field '%s'" % fieldname)
                        self.__validate(eachProperty, value, additionalProperties)
            else:
                raise ValueError("additionalProperties schema definition for field '%s' is not an object" % fieldname)
        return x

    def validate_requires(self, x, fieldname, schema, requires=None):
        if x.get(fieldname) is not None and requires is not None:
            if x.get(requires) is None:
                raise ValueError("Field '%s' is required by field '%s'" % (requires, fieldname))
        return x

    def validate_minimum(self, x, fieldname, schema, minimum=None):
        '''
        Validates that the field is longer than or equal to the minimum
        length if specified
        '''
        if minimum is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if value is not None:
                if type(value) in (int, float) and value < minimum:
                    raise ValueError("Value %r for field '%s' is less than minimum value: %f" % (value, fieldname, minimum))
                elif isinstance(value, list) and len(value) < minimum:
                    raise ValueError("Value %r for field '%s' has fewer values than the minimum: %f" % (value, fieldname, minimum))
        return x

    def validate_maximum(self, x, fieldname, schema, maximum=None):
        '''
        Validates that the field is shorter than or equal to the maximum
        length if specified.
        '''
        if maximum is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if value is not None:
                if type(value) in (int, float) and value > maximum:
                    raise ValueError("Value %r for field '%s' is greater than maximum value: %f" % (value, fieldname, maximum))
                elif isinstance(value, list) and len(value) > maximum:
                    raise ValueError("Value %r for field '%s' has more values than the maximum: %f" % (value, fieldname, maximum))
        return x

    def validate_minItems(self, x, fieldname, schema, minitems=None):
        '''
        Validates that the number of items in the given field is equal to or
        more than the minimum amount.
        '''
        if minitems is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if value is not None:
                if isinstance(value, list) and len(value) < minitems:
                    raise ValueError("Value %r for field '%s' must have a minimum of %d items" % (fieldname, fieldname, minitems))
        return x

    def validate_maxItems(self, x, fieldname, schema, maxitems=None):
        '''
        Validates that the number of items in the given field is equal to or
        less than the maximum amount.
        '''
        if maxitems is not None and x.get(fieldname) is not None:
            value = x.get(fieldname)
            if value is not None:
                if isinstance(value, list) and len(value) > maxitems:
                    raise ValueError("Value %r for field '%s' must have a maximum of %d items" % (value, fieldname, maxitems))
        return x

    def validate_pattern(self, x, fieldname, schema, pattern=None):
        '''
        Validates that the given field, if a string, matches the given
        regular expression.
        '''
        value = x.get(fieldname)
        if pattern is not None and \
             self._is_string_type(value):
            p = re.compile(pattern)
            if not p.match(value):
                raise ValueError("Value %r for field '%s' does not match regular expression '%s'" % (value, fieldname, pattern))
        return x

    def validate_maxLength(self, x, fieldname, schema, length=None):
        '''
        Validates that the value of the given field is shorter than or equal
        to the specified length if a string
        '''
        value = x.get(fieldname)
        if length is not None and \
             self._is_string_type(value) and \
             len(value) > length:
            raise ValueError("Length of value %r for field '%s' must be less than or equal to %f" % (value, fieldname, length))
        return x

    def validate_minLength(self, x, fieldname, schema, length=None):
        '''
        Validates that the value of the given field is longer than or equal
        to the specified length if a string
        '''
        value = x.get(fieldname)
        if length is not None and \
             self._is_string_type(value) and \
             len(value) < length:
            raise ValueError("Length of value %r for field '%s' must be more than or equal to %f" % (value, fieldname, length))
        return x

    def validate_enum(self, x, fieldname, schema, options=None):
        '''
        Validates that the value of the field is equal to one of the
        specified option values
        '''
        value = x.get(fieldname)
        if options is not None and value is not None:
            if not isinstance(options, list):
                raise ValueError("Enumeration %r for field '%s' is not a list type", (options, fieldname))
            if value not in options:
                raise ValueError("Value %r for field '%s' is not in the enumeration: %r" % (value, fieldname, options))
        return x

    def validate_title(self, x, fieldname, schema, title=None):
        if not self._is_string_type(title):
            raise ValueError("The title for field '%s' must be a string" % fieldname)
        return x

    def validate_description(self, x, fieldname, schema, description=None):
        if not self._is_string_type(description):
            raise ValueError("The description for field '%s' must be a string." % fieldname)
        return x

    def validate_format(self, x, fieldname, schema, format=None):
        '''
        Validates that the value of the field matches the predifined format
        specified.
        '''
        # No definitions are currently defined for formats
        return x

    def validate_default(self, x, fieldname, schema, default=None):
        '''
        Adds default data to the original json document if the document is
        not readonly
        '''
        if self._interactive_mode and fieldname not in x.keys() and default is not None:
            if not schema.get("readonly"):
                x[fieldname] = default
        return x

    def validate_maxDecimal(self, x, fieldname, schema, maxdecimal=None):
        '''
        Validates that the value of the given field has less than or equal
        to the maximum number of decimal places given
        '''
        value = x.get(fieldname)
        if maxdecimal is not None and value is not None:
            maxdecstring = str(value)
            if len(maxdecstring[maxdecstring.find(".") + 1:]) > maxdecimal:
                raise ValueError("Value %r for field '%s' must not have more than %d decimal places" % (value, fieldname, maxdecimal))
        return x

    def validate_disallow(self, x, fieldname, schema, disallow=None):
        '''
        Validates that the value of the given field does not match the
        disallowed type.
        '''
        if disallow is not None:
            try:
                self.validate_type(x, fieldname, schema, disallow)
            except ValueError:
                return x
            raise ValueError("Value %r of type %s is disallowed for field '%s'" % (x.get(fieldname), disallow, fieldname))
        return x

    def _convert_type(self, fieldtype):
        if isinstance(fieldtype, (type, dict)):
            return fieldtype
        elif isinstance(fieldtype, list):
            converted_fields = []
            for subfieldtype in fieldtype:
                converted_fields.append(self._convert_type(subfieldtype))
            return converted_fields
        elif fieldtype is None:
            return None
        else:
            fieldtype = str(fieldtype)
            if fieldtype in self._typesmap.keys():
                return self._typesmap[fieldtype]
            else:
                raise ValueError("Field type '%s' is not supported." % fieldtype)

    def validate(self, data, schema):
        '''
        Validates a piece of json data against the provided json-schema.
        '''

        #TODO: Validate the schema object here.

        self._refmap = {
            '$': schema,
        }
        # Wrap the data in a dictionary
        self._validate(data, schema)

    def _validate(self, data, schema):
        self.__validate("_data", {"_data": data}, schema)

    def __validate(self, fieldname, data, schema):

        if schema is not None:
            if not isinstance(schema, dict):
                raise ValueError("Schema structure is invalid.")

            # Produce a copy of the schema object since we will make changes to
            # it to process default values. Deep copy is not necessary since we will
            # produce a copy of sub items on the next recursive call.

            new_schema = copy.copy(schema)
            #Initialize defaults
            for schemaprop in self._schemadefault.keys():
                if schemaprop not in new_schema:
                    new_schema[schemaprop] = self._schemadefault[schemaprop]

            for schemaprop in new_schema:

                validatorname = "validate_" + schemaprop

                try:
                    validator = getattr(self, validatorname)
                    # Pass the original schema object but the value of the property from
                    # copy in order to validate default values.
                    validator(data, fieldname, schema, new_schema.get(schemaprop))
                except AttributeError, e:
                    if schemaprop not in self._ignored:
                        raise ValueError("Schema property '%s' is not supported" %
                                                         schemaprop)

        return data

    def _is_string_type(self, value):
        return value is None or isinstance(value, basestring)

__all__ = ['JSONSchemaValidator']
<|endoftext|>


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

import readline
import shlex
import sys

class JSHError(Exception):
	pass

class JSH(object):
	section = None
	def __init__(self, layout, prompt='> ', section_delims=('(', ')'), ignore_case=False):
		self.layout = layout
		self.prompt = prompt
		self.section_delims = section_delims
		self.ignore_case = ignore_case
		readline.parse_and_bind('tab: complete')
		readline.parse_and_bind('"?": "\\C-v?\\t\\d"')
		readline.parse_and_bind('" ": "\\t"')
		readline.parse_and_bind('"\\r": "\\C-v\\n\\t\\d\\n"')
		readline.set_completer_delims(' ')
		readline.set_completer(self.completer)
	def get_prompt(self):
		return self.prompt.format(section='{0}{1}{2}'.format(self.section_delims[0], self.section, self.section_delims[1]) if self.section else '')
	def read_and_execute(self):
		command = self.get_input()
		if command:
			return self.run_command(command)
	def get_input(self):
		try:
			command = raw_input(self.get_prompt()).strip()
		except EOFError:
			print
			raise
		except KeyboardInterrupt:
			print
			return None
		return command
	def commands(self):
		def walk(level, path=[], paths=[]):
			new_paths = []
			if type(level) == dict:
				for key in level.keys():
					if key in ['\t', '?']:
						continue
					if key is None or len(path) == 1 and key == self.section:
						new_path = path
					elif key == str and '?' in level[str]:
						compl = level[str]['?']
						new_path = path + ['<{0}>'.format(compl if type(compl) == str else compl[0])]
					elif key == str:
						new_path = path + ['<...>']
					else:
						new_path = path + [key]
					new_paths.extend(walk(level[key], new_path, paths))
			else:
				new_paths.append(path)
			return paths + new_paths
		return sorted(walk(self.layout))
	@property
	def completer(self):
		def complete(text, state):
			stext = text.rstrip('?\n')

			try:
				parts = shlex.split(readline.get_line_buffer().rstrip('?')[:readline.get_endidx()].lower())
			except:
				if state == 0:
					return text + ' '
				else:
					return None

			if not stext:
				parts.append(None)
			else:
				parts.pop()

			depth = 0
			completions = {}
			args = []
			level = self.layout
			while True:
				if hasattr(level, '__call__'):
					break
				elif type(level) == dict and '\t' in level and str in level and (not parts or parts[0] is None):
					try:
						possibilities = level['\t'](self, *args)
						if isinstance(possibilities, dict):
							dynamic = possibilities
						else:
							dynamic = dict((comp, '') for comp in possibilities)
					except:
						dynamic = {}
					completions.update(dynamic)
					completions.update(dict((key, level[key].get('?', '') if type(level[key]) == dict else '') for key in level.keys()))
					break
				elif type(level) == dict and (parts and parts[0] is not None and (parts[0] in level or str in level) or depth == 1 and self.section is not None and self.section in level):
					if parts and parts[0] in level:
						level = level[parts[0]]
						parts.pop(0)
					elif depth == 1 and self.section is not None and self.section in level:
						level = level[self.section]
					elif parts and parts[0] is not None and str in level:
						level = level[str]
						args.append(parts.pop(0))
					depth += 1
				elif type(level) == dict:
					completions = dict((key, level[key].get('?', '') if type(level[key]) == dict else '') for key in level.keys())
					break
				else:
					break

			if stext in completions:
				completions = {stext: completions[stext]}
			else:
				completions = dict((key, value) for key, value in completions.iteritems() if key not in [None, str, '\t', '?'] and key.startswith(stext))

			if text == '\n' or text.endswith('\n') and (len(completions) > 1 or str in level):
				return None
			elif text.endswith('?'):
				print
				if str in level and type(level[str]) == dict and '?' in level[str] and '\t' not in level[str]:
					compl = level[str]['?']
					if type(compl) == str:
						completions['<{0}>'.format(compl)] = ''
					else:
						completions['<{0}>'.format(compl[0])] = compl[1]
				if None in level:
					completions['<[Enter]>'] = 'Execute this command'
				if completions:
					just = max(map(len, completions.keys()))
					print 'Possible completions:'
					for key in sorted(completions.keys(), key=lambda comp: '!!!' + comp if comp.startswith('<') else comp):
						print '  {0}   {1}'.format(key.ljust(just), completions[key])
				else:
					print 'No valid completions'
				print '{0}{1}'.format(self.get_prompt(), readline.get_line_buffer()),
				return None
			else:
				if str in level and '\t' not in level:
					completions[text] = ''
				comp_strings = [completion + ' ' for completion in set(completions.keys())]
				if len(comp_strings) > state:
					return sorted(comp_strings)[state]
				else:
					return None
		return complete
	def run_command(self, command):
		try:
			parts = shlex.split(command.lower() if self.ignore_case else command)
		except ValueError as err:
			raise JSHError('Parse error: {0}'.format(err.message.lower()))
		parts.append(None)

		args = []
		consumed = []
		depth = 0
		level = self.layout
		while True:
			if hasattr(level, '__call__'):
				if parts and parts[-1] is None:
					parts.pop()
				if parts:
					raise JSHError("Unexpected arguments{0}: '{1}'".format(
						" after '{0}'".format(' '.join(consumed)) if consumed else '',
						' '.join(parts)
					))
				return level(self, *args)
			elif type(level) == dict and (parts and parts[0] in level or parts and parts[0] is not None and str in level or depth == 1 and self.section is not None and self.section in level):
				if parts and parts[0] in level:
					level = level[parts[0]]
					consumed.append(parts.pop(0))
				elif depth == 1 and self.section is not None and self.section in level:
					level = level[self.section]
				elif parts and parts[0] is not None and str in level:
					args.append(parts[0])
					level = level[str]
					consumed.append(parts.pop(0))
				depth += 1
			else:
				if parts and parts[-1] is None:
					parts.pop()
				if parts:
					raise JSHError("Parse error at '{0}'{1}".format(
						' '.join(parts),
						" after '{0}'".format(' '.join(consumed)) if consumed else '',
					))
				else:
					raise JSHError("Incomplete command '{0}'".format(' '.join(consumed)))

def exit(whatever):
	if type(whatever) == JSH:
		sys.exit(0)
	else:
		def quit(cli):
			sys.exit(whatever)
		return quit

def show_commands(cli):
	print 'Available commands:'
	for command in cli.commands():
		print '  {0}'.format(' '.join(command))

def set_section(section):
	def set_section(cli):
		cli.section = section
	return set_section

<commit_msg>Even better help sorting<commit_after>#!/usr/bin/env python

import readline
import shlex
import sys

class JSHError(Exception):
	pass

class JSH(object):
	section = None
	def __init__(self, layout, prompt='> ', section_delims=('(', ')'), ignore_case=False):
		self.layout = layout
		self.prompt = prompt
		self.section_delims = section_delims
		self.ignore_case = ignore_case
		readline.parse_and_bind('tab: complete')
		readline.parse_and_bind('"?": "\\C-v?\\t\\d"')
		readline.parse_and_bind('" ": "\\t"')
		readline.parse_and_bind('"\\r": "\\C-v\\n\\t\\d\\n"')
		readline.set_completer_delims(' ')
		readline.set_completer(self.completer)
	def get_prompt(self):
		return self.prompt.format(section='{0}{1}{2}'.format(self.section_delims[0], self.section, self.section_delims[1]) if self.section else '')
	def read_and_execute(self):
		command = self.get_input()
		if command:
			return self.run_command(command)
	def get_input(self):
		try:
			command = raw_input(self.get_prompt()).strip()
		except EOFError:
			print
			raise
		except KeyboardInterrupt:
			print
			return None
		return command
	def commands(self):
		def walk(level, path=[], paths=[]):
			new_paths = []
			if type(level) == dict:
				for key in level.keys():
					if key in ['\t', '?']:
						continue
					if key is None or len(path) == 1 and key == self.section:
						new_path = path
					elif key == str and '?' in level[str]:
						compl = level[str]['?']
						new_path = path + ['<{0}>'.format(compl if type(compl) == str else compl[0])]
					elif key == str:
						new_path = path + ['<...>']
					else:
						new_path = path + [key]
					new_paths.extend(walk(level[key], new_path, paths))
			else:
				new_paths.append(path)
			return paths + new_paths
		return sorted(walk(self.layout))
	@property
	def completer(self):
		def complete(text, state):
			stext = text.rstrip('?\n')

			try:
				parts = shlex.split(readline.get_line_buffer().rstrip('?')[:readline.get_endidx()].lower())
			except:
				if state == 0:
					return text + ' '
				else:
					return None

			if not stext:
				parts.append(None)
			else:
				parts.pop()

			depth = 0
			completions = {}
			args = []
			level = self.layout
			while True:
				if hasattr(level, '__call__'):
					break
				elif type(level) == dict and '\t' in level and str in level and (not parts or parts[0] is None):
					try:
						possibilities = level['\t'](self, *args)
						if isinstance(possibilities, dict):
							dynamic = possibilities
						else:
							dynamic = dict((comp, '') for comp in possibilities)
					except:
						dynamic = {}
					completions.update(dynamic)
					completions.update(dict((key, level[key].get('?', '') if type(level[key]) == dict else '') for key in level.keys()))
					break
				elif type(level) == dict and (parts and parts[0] is not None and (parts[0] in level or str in level) or depth == 1 and self.section is not None and self.section in level):
					if parts and parts[0] in level:
						level = level[parts[0]]
						parts.pop(0)
					elif depth == 1 and self.section is not None and self.section in level:
						level = level[self.section]
					elif parts and parts[0] is not None and str in level:
						level = level[str]
						args.append(parts.pop(0))
					depth += 1
				elif type(level) == dict:
					completions = dict((key, level[key].get('?', '') if type(level[key]) == dict else '') for key in level.keys())
					break
				else:
					break

			if stext in completions:
				completions = {stext: completions[stext]}
			else:
				completions = dict((key, value) for key, value in completions.iteritems() if key not in [None, str, '\t', '?'] and key.startswith(stext))

			if text == '\n' or text.endswith('\n') and (len(completions) > 1 or str in level):
				return None
			elif text.endswith('?'):
				print
				if str in level and type(level[str]) == dict and '?' in level[str] and '\t' not in level[str]:
					compl = level[str]['?']
					if type(compl) == str:
						completions['<{0}>'.format(compl)] = ''
					else:
						completions['<{0}>'.format(compl[0])] = compl[1]
				if None in level:
					completions['<[Enter]>'] = 'Execute this command'
				if completions:
					just = max(map(len, completions.keys()))
					print 'Possible completions:'
					for key in sorted(completions.keys(), key=lambda comp: '!!' if comp.startswith('<[') else ('!' + comp if comp.startswith('<') else comp)):
						print '  {0}   {1}'.format(key.ljust(just), completions[key])
				else:
					print 'No valid completions'
				print '{0}{1}'.format(self.get_prompt(), readline.get_line_buffer()),
				return None
			else:
				if str in level and '\t' not in level:
					completions[text] = ''
				comp_strings = [completion + ' ' for completion in set(completions.keys())]
				if len(comp_strings) > state:
					return sorted(comp_strings)[state]
				else:
					return None
		return complete
	def run_command(self, command):
		try:
			parts = shlex.split(command.lower() if self.ignore_case else command)
		except ValueError as err:
			raise JSHError('Parse error: {0}'.format(err.message.lower()))
		parts.append(None)

		args = []
		consumed = []
		depth = 0
		level = self.layout
		while True:
			if hasattr(level, '__call__'):
				if parts and parts[-1] is None:
					parts.pop()
				if parts:
					raise JSHError("Unexpected arguments{0}: '{1}'".format(
						" after '{0}'".format(' '.join(consumed)) if consumed else '',
						' '.join(parts)
					))
				return level(self, *args)
			elif type(level) == dict and (parts and parts[0] in level or parts and parts[0] is not None and str in level or depth == 1 and self.section is not None and self.section in level):
				if parts and parts[0] in level:
					level = level[parts[0]]
					consumed.append(parts.pop(0))
				elif depth == 1 and self.section is not None and self.section in level:
					level = level[self.section]
				elif parts and parts[0] is not None and str in level:
					args.append(parts[0])
					level = level[str]
					consumed.append(parts.pop(0))
				depth += 1
			else:
				if parts and parts[-1] is None:
					parts.pop()
				if parts:
					raise JSHError("Parse error at '{0}'{1}".format(
						' '.join(parts),
						" after '{0}'".format(' '.join(consumed)) if consumed else '',
					))
				else:
					raise JSHError("Incomplete command '{0}'".format(' '.join(consumed)))

def exit(whatever):
	if type(whatever) == JSH:
		sys.exit(0)
	else:
		def quit(cli):
			sys.exit(whatever)
		return quit

def show_commands(cli):
	print 'Available commands:'
	for command in cli.commands():
		print '  {0}'.format(' '.join(command))

def set_section(section):
	def set_section(cli):
		cli.section = section
	return set_section

<|endoftext|>


<commit_before>import pandas as pd
import numpy as np
import re
import json
import os
import warnings

"""
Helper functions to load the article data. The main method to use
is load_data().
"""

# Caution! Modifying this in code will have no effect since the
# default arguments are populated with this reference at creation
# time, so post-hoc modifications will do nothing.
__data_folder = os.path.join(os.path.split(__file__)[0], '..', 'data')

def clean_string(s):
    """
    Clean all the HTML/Unicode nastiness out of a string.
    Replaces newlines with spaces.
    """

    return s.replace('\r', '').replace('\n', ' ').replace('\xa0', ' ').strip()

def load_articles(data_folder=__data_folder, nrows=None):
    """
    Loads the articles CSV. Can optionally only load the first
    `nrows` number of rows.
    """
    column_names = ['id',
                    'feedname',
                    'url',
                    'orig_html',
                    'title',
                    'bodytext',
                    'relevant',
                    'created',
                    'last_modified',
                    'news_source_id',
                    'author']

    return pd.read_csv(os.path.join(data_folder,
                                    'newsarticles_article.csv'),
                       header=None,
                       names=column_names,
                       nrows=nrows,
                       dtype={'orig_html': str, 'author': str})

def load_taggings(data_folder=__data_folder):
    """Loads the type-of-crime human tagging of the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc.set_index('id', drop=True, inplace=True)

    uc_tags_column_names = ['id', 'usercoding_id', 'category_id']

    uc_tags = pd.read_csv(os.path.join(data_folder,
                                       'newsarticles_usercoding_categories.csv'),
                          header=None,
                          names=uc_tags_column_names)
    uc_tags.set_index('usercoding_id', drop=True, inplace=True)

    uc_tags['article_id'] = uc.loc[uc_tags.index, 'article_id']
    return uc_tags

def load_locations(data_folder=__data_folder):
    """Load the human-extracted locations from the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc['locations'] = uc['locations'].apply(lambda x: json.loads(x))

    return uc

def load_categories(data_folder=__data_folder):
    """Loads the mapping of id to names/abbrevations of categories"""
    column_names = ['id', 'category_name', 'abbreviation', 'created',
                    'active', 'kind']

    return pd.read_csv(os.path.join(data_folder, 'newsarticles_category.csv'),
                       header=None,
                       names=column_names)

def load_data(data_folder=__data_folder, nrows=None):
    """
    Creates a dataframe of the article information and k-hot encodes the tags
    into columns called cat_NUMBER. The k-hot encoding is done assuming that the
    categories are 1-indexed and there are as many categories as the maximum
    value of the numerical cateogry_id column.

    Inputs:
        data_folder:
            A folder containing the data files in CSV format.
        nrows:
            Number of articles to load. Defaults to all, which uses about 4
            GB of memory.
    """

    df = load_articles(data_folder=data_folder, nrows=nrows)
    df['relevant'] = df['relevant'] == 't'
    df.rename(columns={'id': 'article_id'}, inplace=True)
    df.set_index('article_id', drop=True, inplace=True)
    # hopefully this will save some memory/space, can add back if needed
    del(df['orig_html'])

    tags_df = load_taggings(data_folder)
    # will help cacheing
    tags_df.sort_values(by='article_id', inplace=True)
    tags_df = tags_df.loc[tags_df['article_id'].isin(df.index.intersection(tags_df['article_id']))]

    locs_df = load_locations(data_folder)
    locs_df.sort_values(by='article_id', inplace=True)
    locs_df = locs_df.loc[locs_df['article_id'].isin(df.index.intersection(locs_df['article_id']))]

    # init with empty lists
    df['locations'] = np.empty([df.shape[0], 0]).tolist()
    df.loc[locs_df['article_id'].values, 'locations'] = locs_df['locations'].values

    def find_loc_in_string(locs, string):
        """
        The locations are generated from JavaScript, which means there's
        going to be some problems getting things to line up exactly and
        neatly. This function will hopefully performa all necessary
        transformations to find the given location text within the
        larger string.

        Inputs:
            locs: list of locations as loaded by load_locations
            string: bodytext of article in which to find locs
        Returns:
            updated_locs: list of locations as loaded by
                load_locations, but with a couple
                extra fields ('cleaned text' and 'cleaned span').
        """

        for i, loc in enumerate(locs):
            loc_text = loc['text']

            loc_text = clean_string(loc_text)
            string = clean_string(string)

            loc['cleaned text'] = loc_text

            spans = [x.span() for x in re.finditer(re.escape(loc_text), string)]
            if spans:
                # The string may have occurred multiple times, and since the spans
                # don't line up perfectly we can't know which one is the "correct" one.
                # Best we can do is find the python span closest to the expected
                # javascript span.
                closest = np.abs(np.argmin(np.array([x[0] for x in spans]) - loc['start']))
                loc['cleaned span'] = spans[closest]

            locs[i] = loc

        return locs

    df['locations'] = df.apply(
        lambda r: find_loc_in_string(r['locations'], r['bodytext']),
        axis=1
    )

    num_no_match = df['locations'].apply(
        lambda locs: any([('cleaned span' not in loc) for loc in locs])
    ).sum()
    if num_no_match:
        warnings.warn(str(num_no_match) + ' location strings were not found in the bodytext.',
                      RuntimeWarning)

    categories_df = load_categories(data_folder)
    categories_df.set_index('id', drop=True, inplace=True)

    # tags_df['category_id'] = tags_df['category_id'].astype(str)
    tags_df['category_abbreviation'] = (categories_df
                                        ['abbreviation']
                                        [tags_df['category_id']]
                                        .values)

    if np.setdiff1d(tags_df['article_id'].values, df.index.values).size:
        warnings.warn('Tags were found for article IDs that do not exist.',
                      RuntimeWarning)

    article_ids = tags_df['article_id'].values
    cat_abbreviations = tags_df['category_abbreviation'].values

    # for some reason, some articles that are tagged don't show up
    # in the articles CSV. filter those out.
    existing_ids_filter = np.isin(article_ids, df.index.values)

    article_ids = article_ids[existing_ids_filter]
    cat_abbreviations = cat_abbreviations[existing_ids_filter]

    for i in range(categories_df.shape[0]):
        cat_name = categories_df.loc[i+1, 'abbreviation']
        df[cat_name] = 0
        df[cat_name] = df[cat_name].astype('int8') # save on that memory!
        df.loc[article_ids[cat_abbreviations == cat_name], cat_name] = 1

    df.loc[df['bodytext'].isnull(), 'bodytext'] = ''

    return df

def load_crime_data(data_folder=__data_folder):
    crimes = pd.read_csv(os.path.join(data_folder, 'Crimes.csv'))
    crimes = crimes[crimes['Year'] > 2010]

    crime_string = pd.Series('', crimes.index)

    # ['ID', 'Case Number', 'Date', 'Block', 'IUCR', 'Primary Type',
    #  'Description', 'Location Description', 'Arrest', 'Domestic', 'Beat',
    #  'District', 'Ward', 'Community Area', 'FBI Code', 'X Coordinate',
    #  'Y Coordinate', 'Year', 'Updated On', 'Latitude', 'Longitude',
    #  'Location']

    # TODO: synonyms on this for month name, weekday name,
    # time of day (e.g. afternoon), etc.
    crime_string += crimes['Date'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Primary Type'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Description'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Location Description'] + ' '

    # TODO: synonyms?
    iucr = pd.read_csv(os.path.join(data_folder, 'IUCR.csv'))
    iucr.set_index('IUCR', drop=True, inplace=True)
    idx = iucr.index
    idx_values = idx.values
    idx_values[idx.str.len() == 3] = '0' + idx_values[idx.str.len() == 3]
    crime_string += (iucr.loc[crimes['IUCR'], 'PRIMARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    crime_string += (iucr.loc[crimes['IUCR'], 'SECONDARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    community_areas = pd.read_csv(os.path.join(data_folder, 'CommAreas.csv'))
    community_areas.set_index('AREA_NUM_1', inplace=True, drop=True)
    crime_string += (community_areas.loc[crimes['Community Area'], 'COMMUNITY']
                     .fillna('')
                     .values
                     + ' ')

    return crimes, crime_string

def load_ner_data(data_folder=__data_folder):
    """
    Loads ner.csv from the specified data folder.

    The column 'stag' is a binary value indicating whether or not
    the row corresponds to the entity "geo". Typically, you will
    want to use column 'word' to predict the column 'stag'.
    """
    df = pd.read_csv(os.path.join(data_folder, 'ner.csv'),
                     encoding="ISO-8859-1",
                     error_bad_lines=False)

    df.dropna(subset=['word', 'tag'], inplace=True)
    df.reset_index(inplace=True)
    df['stag'] = (df['tag'] == 'B-geo') | (df['tag'] == 'I-geo')

    return df
<commit_msg>Drop unneeded index on NER loading.<commit_after>import pandas as pd
import numpy as np
import re
import json
import os
import warnings

"""
Helper functions to load the article data. The main method to use
is load_data().
"""

# Caution! Modifying this in code will have no effect since the
# default arguments are populated with this reference at creation
# time, so post-hoc modifications will do nothing.
__data_folder = os.path.join(os.path.split(__file__)[0], '..', 'data')

def clean_string(s):
    """
    Clean all the HTML/Unicode nastiness out of a string.
    Replaces newlines with spaces.
    """

    return s.replace('\r', '').replace('\n', ' ').replace('\xa0', ' ').strip()

def load_articles(data_folder=__data_folder, nrows=None):
    """
    Loads the articles CSV. Can optionally only load the first
    `nrows` number of rows.
    """
    column_names = ['id',
                    'feedname',
                    'url',
                    'orig_html',
                    'title',
                    'bodytext',
                    'relevant',
                    'created',
                    'last_modified',
                    'news_source_id',
                    'author']

    return pd.read_csv(os.path.join(data_folder,
                                    'newsarticles_article.csv'),
                       header=None,
                       names=column_names,
                       nrows=nrows,
                       dtype={'orig_html': str, 'author': str})

def load_taggings(data_folder=__data_folder):
    """Loads the type-of-crime human tagging of the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc.set_index('id', drop=True, inplace=True)

    uc_tags_column_names = ['id', 'usercoding_id', 'category_id']

    uc_tags = pd.read_csv(os.path.join(data_folder,
                                       'newsarticles_usercoding_categories.csv'),
                          header=None,
                          names=uc_tags_column_names)
    uc_tags.set_index('usercoding_id', drop=True, inplace=True)

    uc_tags['article_id'] = uc.loc[uc_tags.index, 'article_id']
    return uc_tags

def load_locations(data_folder=__data_folder):
    """Load the human-extracted locations from the articles."""
    uc_column_names = ['id', 'date', 'relevant',
                       'article_id', 'user_id', 'locations']

    uc = pd.read_csv(os.path.join(data_folder,
                                  'newsarticles_usercoding.csv'),
                     header=None,
                     names=uc_column_names)

    uc['locations'] = uc['locations'].apply(lambda x: json.loads(x))

    return uc

def load_categories(data_folder=__data_folder):
    """Loads the mapping of id to names/abbrevations of categories"""
    column_names = ['id', 'category_name', 'abbreviation', 'created',
                    'active', 'kind']

    return pd.read_csv(os.path.join(data_folder, 'newsarticles_category.csv'),
                       header=None,
                       names=column_names)

def load_data(data_folder=__data_folder, nrows=None):
    """
    Creates a dataframe of the article information and k-hot encodes the tags
    into columns called cat_NUMBER. The k-hot encoding is done assuming that the
    categories are 1-indexed and there are as many categories as the maximum
    value of the numerical cateogry_id column.

    Inputs:
        data_folder:
            A folder containing the data files in CSV format.
        nrows:
            Number of articles to load. Defaults to all, which uses about 4
            GB of memory.
    """

    df = load_articles(data_folder=data_folder, nrows=nrows)
    df['relevant'] = df['relevant'] == 't'
    df.rename(columns={'id': 'article_id'}, inplace=True)
    df.set_index('article_id', drop=True, inplace=True)
    # hopefully this will save some memory/space, can add back if needed
    del(df['orig_html'])

    tags_df = load_taggings(data_folder)
    # will help cacheing
    tags_df.sort_values(by='article_id', inplace=True)
    tags_df = tags_df.loc[tags_df['article_id'].isin(df.index.intersection(tags_df['article_id']))]

    locs_df = load_locations(data_folder)
    locs_df.sort_values(by='article_id', inplace=True)
    locs_df = locs_df.loc[locs_df['article_id'].isin(df.index.intersection(locs_df['article_id']))]

    # init with empty lists
    df['locations'] = np.empty([df.shape[0], 0]).tolist()
    df.loc[locs_df['article_id'].values, 'locations'] = locs_df['locations'].values

    def find_loc_in_string(locs, string):
        """
        The locations are generated from JavaScript, which means there's
        going to be some problems getting things to line up exactly and
        neatly. This function will hopefully performa all necessary
        transformations to find the given location text within the
        larger string.

        Inputs:
            locs: list of locations as loaded by load_locations
            string: bodytext of article in which to find locs
        Returns:
            updated_locs: list of locations as loaded by
                load_locations, but with a couple
                extra fields ('cleaned text' and 'cleaned span').
        """

        for i, loc in enumerate(locs):
            loc_text = loc['text']

            loc_text = clean_string(loc_text)
            string = clean_string(string)

            loc['cleaned text'] = loc_text

            spans = [x.span() for x in re.finditer(re.escape(loc_text), string)]
            if spans:
                # The string may have occurred multiple times, and since the spans
                # don't line up perfectly we can't know which one is the "correct" one.
                # Best we can do is find the python span closest to the expected
                # javascript span.
                closest = np.abs(np.argmin(np.array([x[0] for x in spans]) - loc['start']))
                loc['cleaned span'] = spans[closest]

            locs[i] = loc

        return locs

    df['locations'] = df.apply(
        lambda r: find_loc_in_string(r['locations'], r['bodytext']),
        axis=1
    )

    num_no_match = df['locations'].apply(
        lambda locs: any([('cleaned span' not in loc) for loc in locs])
    ).sum()
    if num_no_match:
        warnings.warn(str(num_no_match) + ' location strings were not found in the bodytext.',
                      RuntimeWarning)

    categories_df = load_categories(data_folder)
    categories_df.set_index('id', drop=True, inplace=True)

    # tags_df['category_id'] = tags_df['category_id'].astype(str)
    tags_df['category_abbreviation'] = (categories_df
                                        ['abbreviation']
                                        [tags_df['category_id']]
                                        .values)

    if np.setdiff1d(tags_df['article_id'].values, df.index.values).size:
        warnings.warn('Tags were found for article IDs that do not exist.',
                      RuntimeWarning)

    article_ids = tags_df['article_id'].values
    cat_abbreviations = tags_df['category_abbreviation'].values

    # for some reason, some articles that are tagged don't show up
    # in the articles CSV. filter those out.
    existing_ids_filter = np.isin(article_ids, df.index.values)

    article_ids = article_ids[existing_ids_filter]
    cat_abbreviations = cat_abbreviations[existing_ids_filter]

    for i in range(categories_df.shape[0]):
        cat_name = categories_df.loc[i+1, 'abbreviation']
        df[cat_name] = 0
        df[cat_name] = df[cat_name].astype('int8') # save on that memory!
        df.loc[article_ids[cat_abbreviations == cat_name], cat_name] = 1

    df.loc[df['bodytext'].isnull(), 'bodytext'] = ''

    return df

def load_crime_data(data_folder=__data_folder):
    crimes = pd.read_csv(os.path.join(data_folder, 'Crimes.csv'))
    crimes = crimes[crimes['Year'] > 2010]

    crime_string = pd.Series('', crimes.index)

    # ['ID', 'Case Number', 'Date', 'Block', 'IUCR', 'Primary Type',
    #  'Description', 'Location Description', 'Arrest', 'Domestic', 'Beat',
    #  'District', 'Ward', 'Community Area', 'FBI Code', 'X Coordinate',
    #  'Y Coordinate', 'Year', 'Updated On', 'Latitude', 'Longitude',
    #  'Location']

    # TODO: synonyms on this for month name, weekday name,
    # time of day (e.g. afternoon), etc.
    crime_string += crimes['Date'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Primary Type'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Description'] + ' '

    # TODO: synonyms?
    crime_string += crimes['Location Description'] + ' '

    # TODO: synonyms?
    iucr = pd.read_csv(os.path.join(data_folder, 'IUCR.csv'))
    iucr.set_index('IUCR', drop=True, inplace=True)
    idx = iucr.index
    idx_values = idx.values
    idx_values[idx.str.len() == 3] = '0' + idx_values[idx.str.len() == 3]
    crime_string += (iucr.loc[crimes['IUCR'], 'PRIMARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    crime_string += (iucr.loc[crimes['IUCR'], 'SECONDARY DESCRIPTION']
                     .fillna('')
                     .values
                     + ' ')
    community_areas = pd.read_csv(os.path.join(data_folder, 'CommAreas.csv'))
    community_areas.set_index('AREA_NUM_1', inplace=True, drop=True)
    crime_string += (community_areas.loc[crimes['Community Area'], 'COMMUNITY']
                     .fillna('')
                     .values
                     + ' ')

    return crimes, crime_string

def load_ner_data(data_folder=__data_folder):
    """
    Loads ner.csv from the specified data folder.

    The column 'stag' is a binary value indicating whether or not
    the row corresponds to the entity "geo". Typically, you will
    want to use column 'word' to predict the column 'stag'.
    """
    df = pd.read_csv(os.path.join(data_folder, 'ner.csv'),
                     encoding="ISO-8859-1",
                     error_bad_lines=False)

    df.dropna(subset=['word', 'tag'], inplace=True)
    df.reset_index(inplace=True, drop=True)
    df['stag'] = (df['tag'] == 'B-geo') | (df['tag'] == 'I-geo')

    return df
<|endoftext|>