<commit_before># Copyright 2009 Google Inc. All Rights Reserved.
#
# 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.

"""Module for all nameserver health checks."""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import random
import sys
import time

WILDCARD_DOMAINS = ('live.com.', 'blogspot.com.', 'wordpress.com.')
LIKELY_HIJACKS = ['www.google.com.', 'windowsupdate.microsoft.com.', 'www.paypal.com.']

# How many checks to consider when calculating ns check_duration
SHARED_CACHE_TIMEOUT_MULTIPLIER = 1.25
ROOT_SERVER_TIMEOUT_MULTIPLIER = 0.5
CENSORSHIP_TIMEOUT_MULTIPLIER = 2
MAX_STORE_ATTEMPTS = 4
TOTAL_WILDCARDS_TO_STORE = 2
MAX_PORT_BEHAVIOR_TRIES = 2

class NameServerHealthChecks(object):
  """Health checks for a nameserver."""

  def TestAnswers(self, record_type, record, expected, critical=True, timeout=None):
    """Test to see that an answer returns correct IP's.

    Args:
      record_type: text record type for NS query (A, CNAME, etc)
      record: string to query for
      expected: tuple of strings expected in all answers
      critical: If this query fails, should it count against the server.
      timeout: timeout for query in seconds (int)

    Returns:
      (is_broken, error_msg, duration)
    """
    is_broken = False
    unmatched_answers = []
    if not timeout:
      timeout = self.health_timeout
    (response, duration, error_msg) = self.TimedRequest(record_type, record, timeout)
    if not response:
      if not error_msg:
        error_msg = 'No response'
      is_broken = True
    elif not response.answer:
      if critical:
        is_broken = True
      # Avoid preferring broken DNS servers that respond quickly
      duration = self.health_timeout
      error_msg = 'No answer for %s' % record
    else:
      unmatched_answers = []
      found_usable_record = False

      for answer in response.answer:
        if found_usable_record:
          break

        # Process the first sane rdata object available in the answers
        for rdata in answer:
          # CNAME
          if rdata.rdtype == 5:
            reply = str(rdata.target)
          # A Record
          elif rdata.rdtype == 1:
            reply = str(rdata.address)
          else:
            continue

          found_usable_record = True
          found_match = False
          for string in expected:
            if reply.startswith(string) or reply.endswith(string):
              found_match = True
              break
          if not found_match:
            unmatched_answers.append(reply)

      if unmatched_answers:
        hijack_text = ', '.join(unmatched_answers).rstrip('.')
        if record in LIKELY_HIJACKS:
          error_msg = '%s is hijacked: %s' % (record.rstrip('.'), hijack_text)
        else:
          error_msg = '%s appears incorrect: %s' % (record.rstrip('.'), hijack_text)

    return (is_broken, error_msg, duration)

  def TestBindVersion(self):
    """Test for BIND version. This acts as a pretty decent ping."""
    (unused_response, duration, error_msg) = self.RequestVersion()
    # Sometimes nameservers aren't able to respond to this request in a way that
    # dnspython likes.
    is_broken = False
    return (is_broken, error_msg, duration)

  def TestNodeId(self):
    """Get the current node id."""
    self.RequestNodeId()
    return (False, False, 0.0)

  def TestNegativeResponse(self, prefix=None):
    """Test for NXDOMAIN hijaaking."""
    is_broken = False
    if prefix:
      hostname = prefix
      warning_suffix = ' (%s)' % prefix
    else:
      hostname = 'test'
      warning_suffix = ''
    poison_test = '%s.nb%s.google.com.' % (hostname, random.random())
    (response, duration, error_msg) = self.TimedRequest('A', poison_test,
                                                        timeout=self.health_timeout*2)
    if not response:
      if not error_msg:
        error_msg = 'No response'
      is_broken = True
    elif response.answer:
      error_msg = 'NXDOMAIN Hijacking' + warning_suffix

    return (is_broken, error_msg, duration)

  def TestRootNsResponse(self):
    is_broken = False
    error_msg = None
    (response, duration, error_msg) = self.TimedRequest('NS', '.')
    if not response:
      is_broken = True
      if not error_msg:
        error_msg = 'No response'
    return (is_broken, error_msg, duration)

  def TestWwwNegativeResponse(self):
    return self.TestNegativeResponse(prefix='www')

  def TestARootServerResponse(self):
    return self.TestAnswers('A', 'a.root-servers.net.', '198.41.0.4')

  def TestPortBehavior(self, tries=0):
    """This is designed to be called multiple times to retry bad results."""

    if self.port_behavior:
      if 'UNKNOWN' not in self.port_behavior:
        return (False, None, 0)

    tries += 1
    response = self.TimedRequest('TXT', 'porttest.dns-oarc.net.', timeout=5)[0]
    if response and response.answer:
      if len(response.answer) > 1:
        text = response.answer[1].to_rdataset().to_text()
        self.port_behavior = text.split('"')[1]

    if (not self.port_behavior or 'UNKNOWN' in self.port_behavior) and tries < MAX_PORT_BEHAVIOR_TRIES:
      time.sleep(1)
      return self.TestPortBehavior(tries=tries)

#    print "%s behavior: %s (tries=%s)" % (self, self.port_behavior, tries)
    return (False, None, 0)

  def StoreWildcardCache(self):
    """Store a set of wildcard records."""
    timeout = self.health_timeout * SHARED_CACHE_TIMEOUT_MULTIPLIER
    attempted = []

    while len(self.cache_checks) != TOTAL_WILDCARDS_TO_STORE:
      if len(attempted) == MAX_STORE_ATTEMPTS:
        self.disabled = True
        self.AddFailure('Could not recursively query: %s' % ', '.join(attempted))
        return False
      domain = random.choice(WILDCARD_DOMAINS)
      hostname = 'namebench%s.%s' % (random.randint(1, 2**32), domain)
      attempted.append(hostname)
      response = self.TimedRequest('A', hostname, timeout=timeout)[0]
      if response and response.answer:
        self.cache_checks.append((hostname, response, self.timer()))
      else:
        sys.stdout.write('x')

  def TestSharedCache(self, other_ns):
    """Is this nameserver sharing a cache with another nameserver?

    Args:
      other_ns: A nameserver to compare it to.

    Returns:
      A tuple containing:
        - Boolean of whether or not this host has a shared cache
        - The faster NameServer object
        - The slower NameServer object
    """
    timeout = self.health_timeout * SHARED_CACHE_TIMEOUT_MULTIPLIER
    checked = []
    shared = False

    if self.disabled or other_ns.disabled:
      return False

    if not other_ns.cache_checks:
      print '%s has no cache checks (disabling - how did this happen?)' % other_ns
      other_ns.disabled = True
      return False

    for (ref_hostname, ref_response, ref_timestamp) in other_ns.cache_checks:
      response = self.TimedRequest('A', ref_hostname, timeout=timeout)[0]
      # Retry once - this *may* cause false positives however, as the TTL may be updated.
      if not response or not response.answer:
        sys.stdout.write('x')
        response = self.TimedRequest('A', ref_hostname, timeout=timeout)[0]

      if response and response.answer:
        ref_ttl = ref_response.answer[0].ttl
        ttl = response.answer[0].ttl
        delta = abs(ref_ttl - ttl)
        query_age = self.timer() - ref_timestamp
        delta_age_delta = abs(query_age - delta)

        if delta > 0 and delta_age_delta < 2:
          return other_ns
      else:
        sys.stdout.write('!')
      checked.append(ref_hostname)

    if not checked:
      self.AddFailure('Failed to test %s wildcard caches'  % len(other_ns.cache_checks))
    return shared

  def CheckCensorship(self, tests):
    """Check to see if results from a nameserver are being censored."""
    for (check, expected) in tests:
      (req_type, req_name) = check.split(' ')
      expected_values = expected.split(',')
      result = self.TestAnswers(req_type.upper(), req_name, expected_values,
                                timeout=self.health_timeout * CENSORSHIP_TIMEOUT_MULTIPLIER)
      warning = result[1]
      if warning:
        self.AddWarning(warning)

  def CheckHealth(self, sanity_checks=None, fast_check=False, final_check=False, port_check=False):
    """Qualify a nameserver to see if it is any good."""

    is_fatal = False
    if fast_check:
      tests = [(self.TestRootNsResponse, [])]
      is_fatal = True
      sanity_checks = []
    elif final_check:
      tests = [(self.TestWwwNegativeResponse, []), (self.TestPortBehavior, []), (self.TestNodeId, []), (self.TestBindVersion, [])]
    elif port_check:
      tests = [(self.TestPortBehavior, []), (self.TestNodeId, [])]
    else:
      tests = [(self.TestNegativeResponse, [])]

    if sanity_checks:
      for (check, expected_value) in sanity_checks:
        (req_type, req_name) = check.split(' ')
        expected_values = expected_value.split(',')
        tests.append((self.TestAnswers, [req_type.upper(), req_name, expected_values]))

    for test in tests:
      (function, args) = test
      (is_broken, warning, duration) = function(*args)
      if args:
        test_name = args[1]
      else:
        test_name = function.__name__

      self.checks.append((test_name, is_broken, warning, duration))
      if is_broken:
        self.AddFailure('%s: %s' % (test_name, warning), fatal=is_fatal)
      if warning:
        # Special case for NXDOMAIN de-duplication
        if not ('NXDOMAIN' in warning and 'NXDOMAIN Hijacking' in self.warnings):
          self.AddWarning(warning)
      if self.disabled:
        break

    return self.disabled

<commit_msg>Add RCODE checking. Move back to ARootServer check by default, due to differences in response sizes from RootNs<commit_after># Copyright 2009 Google Inc. All Rights Reserved.
#
# 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.

"""Module for all nameserver health checks."""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import random
import sys
import time
import util

from dns import rcode

WILDCARD_DOMAINS = ('live.com.', 'blogspot.com.', 'wordpress.com.')
LIKELY_HIJACKS = ['www.google.com.', 'windowsupdate.microsoft.com.', 'www.paypal.com.']

# How many checks to consider when calculating ns check_duration
SHARED_CACHE_TIMEOUT_MULTIPLIER = 1.25
ROOT_SERVER_TIMEOUT_MULTIPLIER = 0.5
CENSORSHIP_TIMEOUT_MULTIPLIER = 2
MAX_STORE_ATTEMPTS = 4
TOTAL_WILDCARDS_TO_STORE = 2
MAX_PORT_BEHAVIOR_TRIES = 2

FATAL_RCODES = ['REFUSED', 'NOTAUTH']

class NameServerHealthChecks(object):
  """Health checks for a nameserver."""

  def TestAnswers(self, record_type, record, expected, critical=False, timeout=None):
    """Test to see that an answer returns correct IP's.

    Args:
      record_type: text record type for NS query (A, CNAME, etc)
      record: string to query for
      expected: tuple of strings expected in all answers
      critical: If this query fails, should it count against the server.
      timeout: timeout for query in seconds (int)

    Returns:
      (is_broken, error_msg, duration)
    """
    is_broken = False
    unmatched_answers = []
    if not timeout:
      timeout = self.health_timeout
    (response, duration, error_msg) = self.TimedRequest(record_type, record, timeout)
    if response:
      response_code = rcode.to_text(response.rcode())
      if response_code in FATAL_RCODES:
        error_msg = 'Responded with: %s' % response_code
        if critical:
          is_broken = True
      elif not response.answer:
        # Avoid preferring broken DNS servers that respond quickly
        duration = util.SecondsToMilliseconds(self.health_timeout)
        error_msg = 'No answer (%s): %s' % (response_code, record)
        is_broken = True
      else:
        found_usable_record = False
        for answer in response.answer:
          if found_usable_record:
            break

          # Process the first sane rdata object available in the answers
          for rdata in answer:
            # CNAME
            if rdata.rdtype == 5:
              reply = str(rdata.target)
            # A Record
            elif rdata.rdtype == 1:
              reply = str(rdata.address)
            else:
              continue

            found_usable_record = True
            found_match = False
            for string in expected:
              if reply.startswith(string) or reply.endswith(string):
                found_match = True
                break
            if not found_match:
              unmatched_answers.append(reply)

        if unmatched_answers:
          hijack_text = ', '.join(unmatched_answers).rstrip('.')
          if record in LIKELY_HIJACKS:
            error_msg = '%s is hijacked: %s' % (record.rstrip('.'), hijack_text)
          else:
            error_msg = '%s appears incorrect: %s' % (record.rstrip('.'), hijack_text)
    else:
      if not error_msg:
        error_msg = 'No response'
      is_broken = True

    if is_broken:
      print "%s is broken on %s (%s)" % (self, record, error_msg)
    return (is_broken, error_msg, duration)

  def TestBindVersion(self):
    """Test for BIND version. This acts as a pretty decent ping."""
    (unused_response, duration, error_msg) = self.RequestVersion()
    # Sometimes nameservers aren't able to respond to this request in a way that
    # dnspython likes.
    is_broken = False
    return (is_broken, error_msg, duration)

  def TestNodeId(self):
    """Get the current node id."""
    self.RequestNodeId()
    return (False, False, 0.0)

  def TestNegativeResponse(self, prefix=None):
    """Test for NXDOMAIN hijaaking."""
    is_broken = False
    if prefix:
      hostname = prefix
      warning_suffix = ' (%s)' % prefix
    else:
      hostname = 'test'
      warning_suffix = ''
    poison_test = '%s.nb%s.google.com.' % (hostname, random.random())
    (response, duration, error_msg) = self.TimedRequest('A', poison_test,
                                                        timeout=self.health_timeout*2)
    if not response:
      if not error_msg:
        error_msg = 'No response'
      is_broken = True
    elif response.answer:
      error_msg = 'NXDOMAIN Hijacking' + warning_suffix

    return (is_broken, error_msg, duration)

  def TestRootNsResponse(self):
    """Test a . NS response.

    NOTE: This is a bad way to gauge performance of a nameserver, as the
    response length varies between nameserver configurations.
    """
    is_broken = False
    error_msg = None
    (response, duration, error_msg) = self.TimedRequest('NS', '.')
    if not response:
      response_code = None
      is_broken = True
      if not error_msg:
        error_msg = 'No response'
    else:
      response_code = rcode.to_text(response.rcode())
      if response_code in FATAL_RCODES:
        error_msg = response_code
        is_broken = True

    return (is_broken, error_msg, duration)

  def TestWwwNegativeResponse(self):
    return self.TestNegativeResponse(prefix='www')

  def TestARootServerResponse(self):
    return self.TestAnswers('A', 'a.root-servers.net.', '198.41.0.4', critical=True)

  def TestPortBehavior(self, tries=0):
    """This is designed to be called multiple times to retry bad results."""

    if self.port_behavior:
      if 'UNKNOWN' not in self.port_behavior:
        return (False, None, 0)

    tries += 1
    response = self.TimedRequest('TXT', 'porttest.dns-oarc.net.', timeout=5)[0]
    if response and response.answer:
      if len(response.answer) > 1:
        text = response.answer[1].to_rdataset().to_text()
        self.port_behavior = text.split('"')[1]

    if (not self.port_behavior or 'UNKNOWN' in self.port_behavior) and tries < MAX_PORT_BEHAVIOR_TRIES:
      time.sleep(1)
      return self.TestPortBehavior(tries=tries)

#    print "%s behavior: %s (tries=%s)" % (self, self.port_behavior, tries)
    return (False, None, 0)

  def StoreWildcardCache(self):
    """Store a set of wildcard records."""
    timeout = self.health_timeout * SHARED_CACHE_TIMEOUT_MULTIPLIER
    attempted = []

    while len(self.cache_checks) != TOTAL_WILDCARDS_TO_STORE:
      if len(attempted) == MAX_STORE_ATTEMPTS:
        self.disabled = True
        self.AddFailure('Could not recursively query: %s' % ', '.join(attempted))
        return False
      domain = random.choice(WILDCARD_DOMAINS)
      hostname = 'namebench%s.%s' % (random.randint(1, 2**32), domain)
      attempted.append(hostname)
      response = self.TimedRequest('A', hostname, timeout=timeout)[0]
      if response and response.answer:
        self.cache_checks.append((hostname, response, self.timer()))
      else:
        sys.stdout.write('x')

  def TestSharedCache(self, other_ns):
    """Is this nameserver sharing a cache with another nameserver?

    Args:
      other_ns: A nameserver to compare it to.

    Returns:
      A tuple containing:
        - Boolean of whether or not this host has a shared cache
        - The faster NameServer object
        - The slower NameServer object
    """
    timeout = self.health_timeout * SHARED_CACHE_TIMEOUT_MULTIPLIER
    checked = []
    shared = False

    if self.disabled or other_ns.disabled:
      return False

    if not other_ns.cache_checks:
      print '%s has no cache checks (disabling - how did this happen?)' % other_ns
      other_ns.disabled = True
      return False

    for (ref_hostname, ref_response, ref_timestamp) in other_ns.cache_checks:
      response = self.TimedRequest('A', ref_hostname, timeout=timeout)[0]
      # Retry once - this *may* cause false positives however, as the TTL may be updated.
      if not response or not response.answer:
        sys.stdout.write('x')
        response = self.TimedRequest('A', ref_hostname, timeout=timeout)[0]

      if response and response.answer:
        ref_ttl = ref_response.answer[0].ttl
        ttl = response.answer[0].ttl
        delta = abs(ref_ttl - ttl)
        query_age = self.timer() - ref_timestamp
        delta_age_delta = abs(query_age - delta)

        if delta > 0 and delta_age_delta < 2:
          return other_ns
      else:
        sys.stdout.write('!')
      checked.append(ref_hostname)

    if not checked:
      self.AddFailure('Failed to test %s wildcard caches'  % len(other_ns.cache_checks))
    return shared

  def CheckCensorship(self, tests):
    """Check to see if results from a nameserver are being censored."""
    for (check, expected) in tests:
      (req_type, req_name) = check.split(' ')
      expected_values = expected.split(',')
      result = self.TestAnswers(req_type.upper(), req_name, expected_values,
                                timeout=self.health_timeout * CENSORSHIP_TIMEOUT_MULTIPLIER)
      warning = result[1]
      if warning:
        self.AddWarning(warning)

  def CheckHealth(self, sanity_checks=None, fast_check=False, final_check=False, port_check=False):
    """Qualify a nameserver to see if it is any good."""

    is_fatal = False
    if fast_check:
      tests = [(self.TestARootServerResponse, [])]
      is_fatal = True
      sanity_checks = []
    elif final_check:
      tests = [(self.TestWwwNegativeResponse, []), (self.TestPortBehavior, []), (self.TestNodeId, [])]
    elif port_check:
      tests = [(self.TestPortBehavior, []), (self.TestNodeId, [])]
    else:
      # Put the bind version here so that we have a great minimum latency measurement.
      tests = [(self.TestNegativeResponse, []), (self.TestBindVersion, [])]

    if sanity_checks:
      for (check, expected_value) in sanity_checks:
        (req_type, req_name) = check.split(' ')
        expected_values = expected_value.split(',')
        tests.append((self.TestAnswers, [req_type.upper(), req_name, expected_values]))

    for test in tests:
      (function, args) = test
      (is_broken, warning, duration) = function(*args)
      if args:
        test_name = args[1]
      else:
        test_name = function.__name__

      self.checks.append((test_name, is_broken, warning, duration))
      if is_broken:
        self.AddFailure('%s: %s' % (test_name, warning), fatal=is_fatal)
      if warning:
        # Special case for NXDOMAIN de-duplication
        if not ('NXDOMAIN' in warning and 'NXDOMAIN Hijacking' in self.warnings):
          self.AddWarning(warning)
      if self.disabled:
        break

    return self.disabled

<|endoftext|>


<commit_before>"""distutils.msvccompiler

Contains MSVCCompiler, an implementation of the abstract CCompiler class
for the Microsoft Visual Studio."""

# created 1999/08/19, Perry Stoll
# hacked by Robin Becker and Thomas Heller to do a better job of
#   finding DevStudio (through the registry)

__revision__ = "$Id$"

import sys, os, string
from types import *
from distutils.errors import \
     DistutilsExecError, DistutilsPlatformError, \
     CompileError, LibError, LinkError
from distutils.ccompiler import \
     CCompiler, gen_preprocess_options, gen_lib_options

_can_read_reg = 0
try:
    import _winreg
    _can_read_reg = 1
    hkey_mod = _winreg

    RegOpenKeyEx = _winreg.OpenKeyEx
    RegEnumKey = _winreg.EnumKey
    RegEnumValue = _winreg.EnumValue
    RegError = _winreg.error

except ImportError:
    try:
        import win32api
        import win32con
        _can_read_reg = 1
        hkey_mod = win32con

        RegOpenKeyEx = win32api.RegOpenKeyEx
        RegEnumKey = win32api.RegEnumKey
        RegEnumValue = win32api.RegEnumValue
        RegError = win32api.error

    except ImportError:
        pass

if _can_read_reg:
    HKEY_CLASSES_ROOT = hkey_mod.HKEY_CLASSES_ROOT
    HKEY_LOCAL_MACHINE = hkey_mod.HKEY_LOCAL_MACHINE
    HKEY_CURRENT_USER = hkey_mod.HKEY_CURRENT_USER
    HKEY_USERS = hkey_mod.HKEY_USERS

def get_devstudio_versions ():
    """Get list of devstudio versions from the Windows registry.  Return a
       list of strings containing version numbers; the list will be
       empty if we were unable to access the registry (eg. couldn't import
       a registry-access module) or the appropriate registry keys weren't
       found."""

    if not _can_read_reg:
        return []

    K = 'Software\\Microsoft\\Devstudio'
    L = []
    for base in (HKEY_CLASSES_ROOT,
                 HKEY_LOCAL_MACHINE,
                 HKEY_CURRENT_USER,
                 HKEY_USERS):
        try:
            k = RegOpenKeyEx(base,K)
            i = 0
            while 1:
                try:
                    p = RegEnumKey(k,i)
                    if p[0] in '123456789' and p not in L:
                        L.append(p)
                except RegError:
                    break
                i = i + 1
        except RegError:
            pass
    L.sort()
    L.reverse()
    return L

# get_devstudio_versions ()

def get_msvc_paths (path, version='6.0', platform='x86'):
    """Get a list of devstudio directories (include, lib or path).  Return
       a list of strings; will be empty list if unable to access the
       registry or appropriate registry keys not found."""

    if not _can_read_reg:
        return []

    L = []
    if path=='lib':
        path= 'Library'
    path = string.upper(path + ' Dirs')
    K = ('Software\\Microsoft\\Devstudio\\%s\\' +
         'Build System\\Components\\Platforms\\Win32 (%s)\\Directories') % \
        (version,platform)
    for base in (HKEY_CLASSES_ROOT,
                 HKEY_LOCAL_MACHINE,
                 HKEY_CURRENT_USER,
                 HKEY_USERS):
        try:
            k = RegOpenKeyEx(base,K)
            i = 0
            while 1:
                try:
                    (p,v,t) = RegEnumValue(k,i)
                    if string.upper(p) == path:
                        V = string.split(v,';')
                        for v in V:
                            if v == '' or v in L: continue
                            L.append(v)
                        break
                    i = i + 1
                except RegError:
                    break
        except RegError:
            pass
    return L

# get_msvc_paths()

def find_exe (exe, version_number):
    """Try to find an MSVC executable program 'exe' (from version
       'version_number' of MSVC) in several places: first, one of the MSVC
       program search paths from the registry; next, the directories in the
       PATH environment variable.  If any of those work, return an absolute
       path that is known to exist.  If none of them work, just return the
       original program name, 'exe'."""

    for p in get_msvc_paths ('path', version_number):
        fn = os.path.join (os.path.abspath(p), exe)
        if os.path.isfile(fn):
            return fn

    # didn't find it; try existing path
    for p in string.split (os.environ['Path'],';'):
        fn = os.path.join(os.path.abspath(p),exe)
        if os.path.isfile(fn):
            return fn

    return exe                          # last desperate hope

def set_path_env_var (name, version_number):
    """Set environment variable 'name' to an MSVC path type value obtained
       from 'get_msvc_paths()'.  This is equivalent to a SET command prior
       to execution of spawned commands."""

    p = get_msvc_paths (name, version_number)
    if p:
        os.environ[name] = string.join (p,';')

class MSVCCompiler (CCompiler) :
    """Concrete class that implements an interface to Microsoft Visual C++,
       as defined by the CCompiler abstract class."""

    compiler_type = 'msvc'

    # Just set this so CCompiler's constructor doesn't barf.  We currently
    # don't use the 'set_executables()' bureaucracy provided by CCompiler,
    # as it really isn't necessary for this sort of single-compiler class.
    # Would be nice to have a consistent interface with UnixCCompiler,
    # though, so it's worth thinking about.
    executables = {}

    # Private class data (need to distinguish C from C++ source for compiler)
    _c_extensions = ['.c']
    _cpp_extensions = ['.cc','.cpp']

    # Needed for the filename generation methods provided by the
    # base class, CCompiler.
    src_extensions = _c_extensions + _cpp_extensions
    obj_extension = '.obj'
    static_lib_extension = '.lib'
    shared_lib_extension = '.dll'
    static_lib_format = shared_lib_format = '%s%s'
    exe_extension = '.exe'

    def __init__ (self,
                  verbose=0,
                  dry_run=0,
                  force=0):

        CCompiler.__init__ (self, verbose, dry_run, force)
        versions = get_devstudio_versions ()

        if versions:
            version = versions[0]  # highest version

            self.cc   = find_exe("cl.exe", version)
            self.link = find_exe("link.exe", version)
            self.lib  = find_exe("lib.exe", version)
            set_path_env_var ('lib', version)
            set_path_env_var ('include', version)
            path=get_msvc_paths('path', version)
            try:
                for p in string.split(os.environ['path'],';'):
                    path.append(p)
            except KeyError:
                pass
            os.environ['path'] = string.join(path,';')
        else:
            # devstudio not found in the registry
            self.cc = "cl.exe"
            self.link = "link.exe"
            self.lib = "lib.exe"

        self.preprocess_options = None
        self.compile_options = [ '/nologo', '/Ox', '/MD', '/W3' ]
        self.compile_options_debug = ['/nologo', '/Od', '/MDd', '/W3', '/Z7', '/D_DEBUG']

        self.ldflags_shared = ['/DLL', '/nologo', '/INCREMENTAL:NO']
        self.ldflags_shared_debug = [
            '/DLL', '/nologo', '/INCREMENTAL:no', '/pdb:None', '/DEBUG'
            ]
        self.ldflags_static = [ '/nologo']

    # -- Worker methods ------------------------------------------------

    def compile (self,
                 sources,
                 output_dir=None,
                 macros=None,
                 include_dirs=None,
                 debug=0,
                 extra_preargs=None,
                 extra_postargs=None):

        (output_dir, macros, include_dirs) = \
            self._fix_compile_args (output_dir, macros, include_dirs)
        (objects, skip_sources) = self._prep_compile (sources, output_dir)

        if extra_postargs is None:
            extra_postargs = []

        pp_opts = gen_preprocess_options (macros, include_dirs)
        compile_opts = extra_preargs or []
        compile_opts.append ('/c')
        if debug:
            compile_opts.extend (self.compile_options_debug)
        else:
            compile_opts.extend (self.compile_options)

        for i in range (len (sources)):
            src = sources[i] ; obj = objects[i]
            ext = (os.path.splitext (src))[1]

            if skip_sources[src]:
                self.announce ("skipping %s (%s up-to-date)" % (src, obj))
            else:
                if ext in self._c_extensions:
                    input_opt = "/Tc" + src
                elif ext in self._cpp_extensions:
                    input_opt = "/Tp" + src

                output_opt = "/Fo" + obj

                self.mkpath (os.path.dirname (obj))
                try:
                    self.spawn ([self.cc] + compile_opts + pp_opts +
                                [input_opt, output_opt] +
                                extra_postargs)
                except DistutilsExecError, msg:
                    raise CompileError, msg

        return objects

    # compile ()

    def create_static_lib (self,
                           objects,
                           output_libname,
                           output_dir=None,
                           debug=0,
                           extra_preargs=None,
                           extra_postargs=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        output_filename = \
            self.library_filename (output_libname, output_dir=output_dir)

        if self._need_link (objects, output_filename):
            lib_args = objects + ['/OUT:' + output_filename]
            if debug:
                pass                    # XXX what goes here?
            if extra_preargs:
                lib_args[:0] = extra_preargs
            if extra_postargs:
                lib_args.extend (extra_postargs)
            try:
                self.spawn ([self.lib] + lib_args)
            except DistutilsExecError, msg:
                raise LibError, msg

        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # create_static_lib ()

    def link_shared_lib (self,
                         objects,
                         output_libname,
                         output_dir=None,
                         libraries=None,
                         library_dirs=None,
                         runtime_library_dirs=None,
                         export_symbols=None,
                         debug=0,
                         extra_preargs=None,
                         extra_postargs=None,
                         build_temp=None):

        self.link_shared_object (objects,
                                 self.shared_library_name(output_libname),
                                 output_dir=output_dir,
                                 libraries=libraries,
                                 library_dirs=library_dirs,
                                 runtime_library_dirs=runtime_library_dirs,
                                 export_symbols=export_symbols,
                                 debug=debug,
                                 extra_preargs=extra_preargs,
                                 extra_postargs=extra_postargs,
                                 build_temp=build_temp)

    def link_shared_object (self,
                            objects,
                            output_filename,
                            output_dir=None,
                            libraries=None,
                            library_dirs=None,
                            runtime_library_dirs=None,
                            export_symbols=None,
                            debug=0,
                            extra_preargs=None,
                            extra_postargs=None,
                            build_temp=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        (libraries, library_dirs, runtime_library_dirs) = \
            self._fix_lib_args (libraries, library_dirs, runtime_library_dirs)

        if runtime_library_dirs:
            self.warn ("I don't know what to do with 'runtime_library_dirs': "
                       + str (runtime_library_dirs))

        lib_opts = gen_lib_options (self,
                                    library_dirs, runtime_library_dirs,
                                    libraries)
        if output_dir is not None:
            output_filename = os.path.join (output_dir, output_filename)

        if self._need_link (objects, output_filename):

            if debug:
                ldflags = self.ldflags_shared_debug
            else:
                ldflags = self.ldflags_shared

            export_opts = []
            for sym in (export_symbols or []):
                export_opts.append("/EXPORT:" + sym)

            ld_args = (ldflags + lib_opts + export_opts +
                       objects + ['/OUT:' + output_filename])

            if extra_preargs:
                ld_args[:0] = extra_preargs
            if extra_postargs:
                ld_args.extend (extra_postargs)

            print "link_shared_object():"
            print "  output_filename =", output_filename
            print "  mkpath'ing:", os.path.dirname (output_filename)
            self.mkpath (os.path.dirname (output_filename))
            try:
                self.spawn ([self.link] + ld_args)
            except DistutilsExecError, msg:
                raise LinkError, msg

        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # link_shared_object ()

    def link_executable (self,
                         objects,
                         output_progname,
                         output_dir=None,
                         libraries=None,
                         library_dirs=None,
                         runtime_library_dirs=None,
                         debug=0,
                         extra_preargs=None,
                         extra_postargs=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        (libraries, library_dirs, runtime_library_dirs) = \
            self._fix_lib_args (libraries, library_dirs, runtime_library_dirs)

        if runtime_library_dirs:
            self.warn ("I don't know what to do with 'runtime_library_dirs': "
                       + str (runtime_library_dirs))

        lib_opts = gen_lib_options (self,
                                    library_dirs, runtime_library_dirs,
                                    libraries)
        output_filename = output_progname + self.exe_extension
        if output_dir is not None:
            output_filename = os.path.join (output_dir, output_filename)

        if self._need_link (objects, output_filename):

            if debug:
                ldflags = self.ldflags_shared_debug[1:]
            else:
                ldflags = self.ldflags_shared[1:]

            ld_args = ldflags + lib_opts + \
                      objects + ['/OUT:' + output_filename]

            if extra_preargs:
                ld_args[:0] = extra_preargs
            if extra_postargs:
                ld_args.extend (extra_postargs)

            self.mkpath (os.path.dirname (output_filename))
            try:
                self.spawn ([self.link] + ld_args)
            except DistutilsExecError, msg:
                raise LinkError, msg
        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # -- Miscellaneous methods -----------------------------------------
    # These are all used by the 'gen_lib_options() function, in
    # ccompiler.py.

    def library_dir_option (self, dir):
        return "/LIBPATH:" + dir

    def runtime_library_dir_option (self, dir):
        raise DistutilsPlatformError, \
              "don't know how to set runtime library search path for MSVC++"

    def library_option (self, lib):
        return self.library_filename (lib)

    def find_library_file (self, dirs, lib):

        for dir in dirs:
            libfile = os.path.join (dir, self.library_filename (lib))
            if os.path.exists (libfile):
                return libfile

        else:
            # Oops, didn't find it in *any* of 'dirs'
            return None

    # find_library_file ()

# class MSVCCompiler
<commit_msg>On second thought, first try for _winreg, and then winreg.  Only if both fail do we try for win32api/win32con.  If *those* both fail, then we don't have registry access.  Phew!<commit_after>"""distutils.msvccompiler

Contains MSVCCompiler, an implementation of the abstract CCompiler class
for the Microsoft Visual Studio."""

# created 1999/08/19, Perry Stoll
# hacked by Robin Becker and Thomas Heller to do a better job of
#   finding DevStudio (through the registry)

__revision__ = "$Id$"

import sys, os, string
from types import *
from distutils.errors import \
     DistutilsExecError, DistutilsPlatformError, \
     CompileError, LibError, LinkError
from distutils.ccompiler import \
     CCompiler, gen_preprocess_options, gen_lib_options

_can_read_reg = 0
try:
    try:
        import _winreg
    except ImportError:
        import winreg                   # for pre-2000/06/29 CVS Python

    _can_read_reg = 1
    hkey_mod = _winreg

    RegOpenKeyEx = _winreg.OpenKeyEx
    RegEnumKey = _winreg.EnumKey
    RegEnumValue = _winreg.EnumValue
    RegError = _winreg.error

except ImportError:
    try:
        import win32api
        import win32con
        _can_read_reg = 1
        hkey_mod = win32con

        RegOpenKeyEx = win32api.RegOpenKeyEx
        RegEnumKey = win32api.RegEnumKey
        RegEnumValue = win32api.RegEnumValue
        RegError = win32api.error

    except ImportError:
        pass

if _can_read_reg:
    HKEY_CLASSES_ROOT = hkey_mod.HKEY_CLASSES_ROOT
    HKEY_LOCAL_MACHINE = hkey_mod.HKEY_LOCAL_MACHINE
    HKEY_CURRENT_USER = hkey_mod.HKEY_CURRENT_USER
    HKEY_USERS = hkey_mod.HKEY_USERS

def get_devstudio_versions ():
    """Get list of devstudio versions from the Windows registry.  Return a
       list of strings containing version numbers; the list will be
       empty if we were unable to access the registry (eg. couldn't import
       a registry-access module) or the appropriate registry keys weren't
       found."""

    if not _can_read_reg:
        return []

    K = 'Software\\Microsoft\\Devstudio'
    L = []
    for base in (HKEY_CLASSES_ROOT,
                 HKEY_LOCAL_MACHINE,
                 HKEY_CURRENT_USER,
                 HKEY_USERS):
        try:
            k = RegOpenKeyEx(base,K)
            i = 0
            while 1:
                try:
                    p = RegEnumKey(k,i)
                    if p[0] in '123456789' and p not in L:
                        L.append(p)
                except RegError:
                    break
                i = i + 1
        except RegError:
            pass
    L.sort()
    L.reverse()
    return L

# get_devstudio_versions ()

def get_msvc_paths (path, version='6.0', platform='x86'):
    """Get a list of devstudio directories (include, lib or path).  Return
       a list of strings; will be empty list if unable to access the
       registry or appropriate registry keys not found."""

    if not _can_read_reg:
        return []

    L = []
    if path=='lib':
        path= 'Library'
    path = string.upper(path + ' Dirs')
    K = ('Software\\Microsoft\\Devstudio\\%s\\' +
         'Build System\\Components\\Platforms\\Win32 (%s)\\Directories') % \
        (version,platform)
    for base in (HKEY_CLASSES_ROOT,
                 HKEY_LOCAL_MACHINE,
                 HKEY_CURRENT_USER,
                 HKEY_USERS):
        try:
            k = RegOpenKeyEx(base,K)
            i = 0
            while 1:
                try:
                    (p,v,t) = RegEnumValue(k,i)
                    if string.upper(p) == path:
                        V = string.split(v,';')
                        for v in V:
                            if v == '' or v in L: continue
                            L.append(v)
                        break
                    i = i + 1
                except RegError:
                    break
        except RegError:
            pass
    return L

# get_msvc_paths()

def find_exe (exe, version_number):
    """Try to find an MSVC executable program 'exe' (from version
       'version_number' of MSVC) in several places: first, one of the MSVC
       program search paths from the registry; next, the directories in the
       PATH environment variable.  If any of those work, return an absolute
       path that is known to exist.  If none of them work, just return the
       original program name, 'exe'."""

    for p in get_msvc_paths ('path', version_number):
        fn = os.path.join (os.path.abspath(p), exe)
        if os.path.isfile(fn):
            return fn

    # didn't find it; try existing path
    for p in string.split (os.environ['Path'],';'):
        fn = os.path.join(os.path.abspath(p),exe)
        if os.path.isfile(fn):
            return fn

    return exe                          # last desperate hope

def set_path_env_var (name, version_number):
    """Set environment variable 'name' to an MSVC path type value obtained
       from 'get_msvc_paths()'.  This is equivalent to a SET command prior
       to execution of spawned commands."""

    p = get_msvc_paths (name, version_number)
    if p:
        os.environ[name] = string.join (p,';')

class MSVCCompiler (CCompiler) :
    """Concrete class that implements an interface to Microsoft Visual C++,
       as defined by the CCompiler abstract class."""

    compiler_type = 'msvc'

    # Just set this so CCompiler's constructor doesn't barf.  We currently
    # don't use the 'set_executables()' bureaucracy provided by CCompiler,
    # as it really isn't necessary for this sort of single-compiler class.
    # Would be nice to have a consistent interface with UnixCCompiler,
    # though, so it's worth thinking about.
    executables = {}

    # Private class data (need to distinguish C from C++ source for compiler)
    _c_extensions = ['.c']
    _cpp_extensions = ['.cc','.cpp']

    # Needed for the filename generation methods provided by the
    # base class, CCompiler.
    src_extensions = _c_extensions + _cpp_extensions
    obj_extension = '.obj'
    static_lib_extension = '.lib'
    shared_lib_extension = '.dll'
    static_lib_format = shared_lib_format = '%s%s'
    exe_extension = '.exe'

    def __init__ (self,
                  verbose=0,
                  dry_run=0,
                  force=0):

        CCompiler.__init__ (self, verbose, dry_run, force)
        versions = get_devstudio_versions ()

        if versions:
            version = versions[0]  # highest version

            self.cc   = find_exe("cl.exe", version)
            self.link = find_exe("link.exe", version)
            self.lib  = find_exe("lib.exe", version)
            set_path_env_var ('lib', version)
            set_path_env_var ('include', version)
            path=get_msvc_paths('path', version)
            try:
                for p in string.split(os.environ['path'],';'):
                    path.append(p)
            except KeyError:
                pass
            os.environ['path'] = string.join(path,';')
        else:
            # devstudio not found in the registry
            self.cc = "cl.exe"
            self.link = "link.exe"
            self.lib = "lib.exe"

        self.preprocess_options = None
        self.compile_options = [ '/nologo', '/Ox', '/MD', '/W3' ]
        self.compile_options_debug = ['/nologo', '/Od', '/MDd', '/W3', '/Z7', '/D_DEBUG']

        self.ldflags_shared = ['/DLL', '/nologo', '/INCREMENTAL:NO']
        self.ldflags_shared_debug = [
            '/DLL', '/nologo', '/INCREMENTAL:no', '/pdb:None', '/DEBUG'
            ]
        self.ldflags_static = [ '/nologo']

    # -- Worker methods ------------------------------------------------

    def compile (self,
                 sources,
                 output_dir=None,
                 macros=None,
                 include_dirs=None,
                 debug=0,
                 extra_preargs=None,
                 extra_postargs=None):

        (output_dir, macros, include_dirs) = \
            self._fix_compile_args (output_dir, macros, include_dirs)
        (objects, skip_sources) = self._prep_compile (sources, output_dir)

        if extra_postargs is None:
            extra_postargs = []

        pp_opts = gen_preprocess_options (macros, include_dirs)
        compile_opts = extra_preargs or []
        compile_opts.append ('/c')
        if debug:
            compile_opts.extend (self.compile_options_debug)
        else:
            compile_opts.extend (self.compile_options)

        for i in range (len (sources)):
            src = sources[i] ; obj = objects[i]
            ext = (os.path.splitext (src))[1]

            if skip_sources[src]:
                self.announce ("skipping %s (%s up-to-date)" % (src, obj))
            else:
                if ext in self._c_extensions:
                    input_opt = "/Tc" + src
                elif ext in self._cpp_extensions:
                    input_opt = "/Tp" + src

                output_opt = "/Fo" + obj

                self.mkpath (os.path.dirname (obj))
                try:
                    self.spawn ([self.cc] + compile_opts + pp_opts +
                                [input_opt, output_opt] +
                                extra_postargs)
                except DistutilsExecError, msg:
                    raise CompileError, msg

        return objects

    # compile ()

    def create_static_lib (self,
                           objects,
                           output_libname,
                           output_dir=None,
                           debug=0,
                           extra_preargs=None,
                           extra_postargs=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        output_filename = \
            self.library_filename (output_libname, output_dir=output_dir)

        if self._need_link (objects, output_filename):
            lib_args = objects + ['/OUT:' + output_filename]
            if debug:
                pass                    # XXX what goes here?
            if extra_preargs:
                lib_args[:0] = extra_preargs
            if extra_postargs:
                lib_args.extend (extra_postargs)
            try:
                self.spawn ([self.lib] + lib_args)
            except DistutilsExecError, msg:
                raise LibError, msg

        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # create_static_lib ()

    def link_shared_lib (self,
                         objects,
                         output_libname,
                         output_dir=None,
                         libraries=None,
                         library_dirs=None,
                         runtime_library_dirs=None,
                         export_symbols=None,
                         debug=0,
                         extra_preargs=None,
                         extra_postargs=None,
                         build_temp=None):

        self.link_shared_object (objects,
                                 self.shared_library_name(output_libname),
                                 output_dir=output_dir,
                                 libraries=libraries,
                                 library_dirs=library_dirs,
                                 runtime_library_dirs=runtime_library_dirs,
                                 export_symbols=export_symbols,
                                 debug=debug,
                                 extra_preargs=extra_preargs,
                                 extra_postargs=extra_postargs,
                                 build_temp=build_temp)

    def link_shared_object (self,
                            objects,
                            output_filename,
                            output_dir=None,
                            libraries=None,
                            library_dirs=None,
                            runtime_library_dirs=None,
                            export_symbols=None,
                            debug=0,
                            extra_preargs=None,
                            extra_postargs=None,
                            build_temp=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        (libraries, library_dirs, runtime_library_dirs) = \
            self._fix_lib_args (libraries, library_dirs, runtime_library_dirs)

        if runtime_library_dirs:
            self.warn ("I don't know what to do with 'runtime_library_dirs': "
                       + str (runtime_library_dirs))

        lib_opts = gen_lib_options (self,
                                    library_dirs, runtime_library_dirs,
                                    libraries)
        if output_dir is not None:
            output_filename = os.path.join (output_dir, output_filename)

        if self._need_link (objects, output_filename):

            if debug:
                ldflags = self.ldflags_shared_debug
            else:
                ldflags = self.ldflags_shared

            export_opts = []
            for sym in (export_symbols or []):
                export_opts.append("/EXPORT:" + sym)

            ld_args = (ldflags + lib_opts + export_opts +
                       objects + ['/OUT:' + output_filename])

            if extra_preargs:
                ld_args[:0] = extra_preargs
            if extra_postargs:
                ld_args.extend (extra_postargs)

            print "link_shared_object():"
            print "  output_filename =", output_filename
            print "  mkpath'ing:", os.path.dirname (output_filename)
            self.mkpath (os.path.dirname (output_filename))
            try:
                self.spawn ([self.link] + ld_args)
            except DistutilsExecError, msg:
                raise LinkError, msg

        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # link_shared_object ()

    def link_executable (self,
                         objects,
                         output_progname,
                         output_dir=None,
                         libraries=None,
                         library_dirs=None,
                         runtime_library_dirs=None,
                         debug=0,
                         extra_preargs=None,
                         extra_postargs=None):

        (objects, output_dir) = self._fix_object_args (objects, output_dir)
        (libraries, library_dirs, runtime_library_dirs) = \
            self._fix_lib_args (libraries, library_dirs, runtime_library_dirs)

        if runtime_library_dirs:
            self.warn ("I don't know what to do with 'runtime_library_dirs': "
                       + str (runtime_library_dirs))

        lib_opts = gen_lib_options (self,
                                    library_dirs, runtime_library_dirs,
                                    libraries)
        output_filename = output_progname + self.exe_extension
        if output_dir is not None:
            output_filename = os.path.join (output_dir, output_filename)

        if self._need_link (objects, output_filename):

            if debug:
                ldflags = self.ldflags_shared_debug[1:]
            else:
                ldflags = self.ldflags_shared[1:]

            ld_args = ldflags + lib_opts + \
                      objects + ['/OUT:' + output_filename]

            if extra_preargs:
                ld_args[:0] = extra_preargs
            if extra_postargs:
                ld_args.extend (extra_postargs)

            self.mkpath (os.path.dirname (output_filename))
            try:
                self.spawn ([self.link] + ld_args)
            except DistutilsExecError, msg:
                raise LinkError, msg
        else:
            self.announce ("skipping %s (up-to-date)" % output_filename)

    # -- Miscellaneous methods -----------------------------------------
    # These are all used by the 'gen_lib_options() function, in
    # ccompiler.py.

    def library_dir_option (self, dir):
        return "/LIBPATH:" + dir

    def runtime_library_dir_option (self, dir):
        raise DistutilsPlatformError, \
              "don't know how to set runtime library search path for MSVC++"

    def library_option (self, lib):
        return self.library_filename (lib)

    def find_library_file (self, dirs, lib):

        for dir in dirs:
            libfile = os.path.join (dir, self.library_filename (lib))
            if os.path.exists (libfile):
                return libfile

        else:
            # Oops, didn't find it in *any* of 'dirs'
            return None

    # find_library_file ()

# class MSVCCompiler
<|endoftext|>


<commit_before>import numpy as np
import unittest

import HARK.distribution as distribution

from HARK.distribution import (
    Bernoulli,
    DiscreteDistribution,
    Lognormal,
    MeanOneLogNormal,
    Normal,
    Uniform,
    Weibull,
    calcExpectation,
    combineIndepDstns
)

class DiscreteDistributionTests(unittest.TestCase):
    """
    Tests for distribution.py sampling distributions
    with default seed.
    """

    def test_draw(self):
        self.assertEqual(
            DiscreteDistribution(
                np.ones(1),
                np.zeros(1)).draw(1)[
                0
            ],
            0,
        )

    def test_calcExpectation(self):
        dd_0_1_20 = Normal().approx(20)
        dd_1_1_40 = Normal(mu = 1).approx(40)
        dd_10_10_100 = Normal(mu = 10, sigma = 10).approx(100)

        ce1 = calcExpectation(dd_0_1_20)
        ce2 = calcExpectation(dd_1_1_40)
        ce3 = calcExpectation(dd_10_10_100)

        self.assertAlmostEqual(ce1, 0.0)
        self.assertAlmostEqual(ce2, 1.0)
        self.assertAlmostEqual(ce3, 10.0)

        ce4= calcExpectation(
            dd_0_1_20,
            lambda x: 2 ** x
        )

        self.assertAlmostEqual(ce4, 1.27153712)

        ce5 = calcExpectation(
            dd_1_1_40,
            lambda x: 2 * x
        )

        self.assertAlmostEqual(ce5, 2.0)

        ce6 = calcExpectation(
            dd_10_10_100,
            lambda x, y: 2 * x + y,
            20
        )

        self.assertAlmostEqual(ce6, 40.0)

        ce7 = calcExpectation(
            dd_0_1_20,
            lambda x, y: x + y,
            np.hstack(np.array([0,1,2,3,4,5]))
        )

        self.assertAlmostEqual(ce7.flat[3], 3.0)

        PermShkDstn = MeanOneLogNormal().approx(200)
        TranShkDstn = MeanOneLogNormal().approx(200)
        IncShkDstn = combineIndepDstns(PermShkDstn, TranShkDstn)

        ce8 = calcExpectation(
            IncShkDstn,
            lambda X: X[0] + X[1]
        )

        self.assertAlmostEqual(ce8, 2.0)

        ce9 = calcExpectation(
            IncShkDstn,
            lambda X, a, r: r / X[0] * a + X[1],
            np.array([0,1,2,3,4,5]), # an aNrmNow grid?
            1.05 # an interest rate?
        )

        self.assertAlmostEqual(
            ce9[3],
            9.518015322143837
        )

class DistributionClassTests(unittest.TestCase):
    """
    Tests for distribution.py sampling distributions
    with default seed.
    """

    def test_drawMeanOneLognormal(self):
        self.assertEqual(MeanOneLogNormal().draw(1)[0], 3.5397367004222002)

    def test_Lognormal(self):

        dist = Lognormal()

        self.assertEqual(dist.draw(1)[0], 5.836039190663969)

        dist.draw(100)
        dist.reset()

        self.assertEqual(dist.draw(1)[0], 5.836039190663969)

    def test_Normal(self):
        dist = Normal()

        self.assertEqual(dist.draw(1)[0], 1.764052345967664)

        dist.draw(100)
        dist.reset()

        self.assertEqual(dist.draw(1)[0], 1.764052345967664)

    def test_Weibull(self):
        self.assertEqual(
            Weibull().draw(1)[0],
            0.79587450816311)

    def test_Uniform(self):
        uni = Uniform()

        self.assertEqual(
            Uniform().draw(1)[0],
            0.5488135039273248)

        self.assertEqual(
            calcExpectation(uni.approx(10)),
            0.5
        )

    def test_Bernoulli(self):
        self.assertEqual(
            Bernoulli().draw(1)[0], False
        )

class MarkovProcessTests(unittest.TestCase):
    """
    Tests for MarkovProcess class.
    """

    def test_draw(self):

        mrkv_array = np.array(
            [[.75, .25],[0.1, 0.9]]
        )

        mp = distribution.MarkovProcess(mrkv_array)

        new_state = mp.draw(np.zeros(100).astype(int))

        self.assertEqual(new_state.sum(), 20)

        new_state = mp.draw(new_state)

        self.assertEqual(new_state.sum(), 39)
<commit_msg>add test and lint<commit_after>import numpy as np
import unittest

import HARK.distribution as distribution

from HARK.distribution import (
    Bernoulli,
    DiscreteDistribution,
    Lognormal,
    MeanOneLogNormal,
    Normal,
    MVNormal,
    Uniform,
    Weibull,
    calcExpectation,
    combineIndepDstns,
)

class DiscreteDistributionTests(unittest.TestCase):
    """
    Tests for distribution.py sampling distributions
    with default seed.
    """

    def test_draw(self):
        self.assertEqual(
            DiscreteDistribution(np.ones(1), np.zeros(1)).draw(1)[0], 0,
        )

    def test_calcExpectation(self):
        dd_0_1_20 = Normal().approx(20)
        dd_1_1_40 = Normal(mu=1).approx(40)
        dd_10_10_100 = Normal(mu=10, sigma=10).approx(100)

        ce1 = calcExpectation(dd_0_1_20)
        ce2 = calcExpectation(dd_1_1_40)
        ce3 = calcExpectation(dd_10_10_100)

        self.assertAlmostEqual(ce1, 0.0)
        self.assertAlmostEqual(ce2, 1.0)
        self.assertAlmostEqual(ce3, 10.0)

        ce4 = calcExpectation(dd_0_1_20, lambda x: 2 ** x)

        self.assertAlmostEqual(ce4, 1.27153712)

        ce5 = calcExpectation(dd_1_1_40, lambda x: 2 * x)

        self.assertAlmostEqual(ce5, 2.0)

        ce6 = calcExpectation(dd_10_10_100, lambda x, y: 2 * x + y, 20)

        self.assertAlmostEqual(ce6, 40.0)

        ce7 = calcExpectation(
            dd_0_1_20, lambda x, y: x + y, np.hstack(np.array([0, 1, 2, 3, 4, 5]))
        )

        self.assertAlmostEqual(ce7.flat[3], 3.0)

        PermShkDstn = MeanOneLogNormal().approx(200)
        TranShkDstn = MeanOneLogNormal().approx(200)
        IncShkDstn = combineIndepDstns(PermShkDstn, TranShkDstn)

        ce8 = calcExpectation(IncShkDstn, lambda X: X[0] + X[1])

        self.assertAlmostEqual(ce8, 2.0)

        ce9 = calcExpectation(
            IncShkDstn,
            lambda X, a, r: r / X[0] * a + X[1],
            np.array([0, 1, 2, 3, 4, 5]),  # an aNrmNow grid?
            1.05,  # an interest rate?
        )

        self.assertAlmostEqual(ce9[3], 9.518015322143837)

class DistributionClassTests(unittest.TestCase):
    """
    Tests for distribution.py sampling distributions
    with default seed.
    """

    def test_drawMeanOneLognormal(self):
        self.assertEqual(MeanOneLogNormal().draw(1)[0], 3.5397367004222002)

    def test_Lognormal(self):

        dist = Lognormal()

        self.assertEqual(dist.draw(1)[0], 5.836039190663969)

        dist.draw(100)
        dist.reset()

        self.assertEqual(dist.draw(1)[0], 5.836039190663969)

    def test_Normal(self):
        dist = Normal()

        self.assertEqual(dist.draw(1)[0], 1.764052345967664)

        dist.draw(100)
        dist.reset()

        self.assertEqual(dist.draw(1)[0], 1.764052345967664)

    def test_MVNormal(self):
        dist = MVNormal()

        self.assertTrue(
            np.allclose(dist.draw(1)[0], np.array([2.76405235, 1.40015721]))
        )

        dist.draw(100)
        dist.reset()

        self.assertTrue(
            np.allclose(dist.draw(1)[0], np.array([2.76405235, 1.40015721]))
        )

    def test_Weibull(self):
        self.assertEqual(Weibull().draw(1)[0], 0.79587450816311)

    def test_Uniform(self):
        uni = Uniform()

        self.assertEqual(Uniform().draw(1)[0], 0.5488135039273248)

        self.assertEqual(calcExpectation(uni.approx(10)), 0.5)

    def test_Bernoulli(self):
        self.assertEqual(Bernoulli().draw(1)[0], False)

class MarkovProcessTests(unittest.TestCase):
    """
    Tests for MarkovProcess class.
    """

    def test_draw(self):

        mrkv_array = np.array([[0.75, 0.25], [0.1, 0.9]])

        mp = distribution.MarkovProcess(mrkv_array)

        new_state = mp.draw(np.zeros(100).astype(int))

        self.assertEqual(new_state.sum(), 20)

        new_state = mp.draw(new_state)

        self.assertEqual(new_state.sum(), 39)
<|endoftext|>


<commit_before>import unittest
from typing import List

import utils

# O(n!) time. O(n) space. Partition problem, backtracking.
class Solution:
    def makesquare(self, nums: List[int]) -> bool:
        if len(nums) < 4:
            return False

        s = sum(nums)
        side, r = divmod(s, 4)

        if r:
            return False

        def dfs(start, target):
            if start == len(nums):
                return target == 0

            if target == 0:
                target = side

            for i in range(start, len(nums)):
                new_target = target - nums[i]
                if new_target >= 0:
                    nums[start], nums[i] = nums[i], nums[start]

                    if dfs(start + 1, new_target):
                        return True

                    nums[start], nums[i] = nums[i], nums[start]

            return False

        return dfs(0, side)

class Test(unittest.TestCase):
    def test(self):
        cases = utils.load_test_json(__file__).test_cases

        for case in cases:
            args = str(case.args)
            actual = Solution().makesquare(**case.args.__dict__)
            self.assertEqual(case.expected, actual, msg=args)

if __name__ == '__main__':
    unittest.main()
<commit_msg>473. Matchsticks to Square<commit_after>import unittest
from typing import List

import utils

# O(4^n) time. O(n) space. Partition problem, backtracking.
class Solution:
    def makesquare(self, nums: List[int]) -> bool:
        if len(nums) < 4:
            return False

        s = sum(nums)
        side, r = divmod(s, 4)

        if r:
            return False

        sums = [0] * 4

        def dfs(index):
            if index == len(nums):
                return True

            for i in range(4):
                if sums[i] + nums[index] <= side:
                    sums[i] += nums[index]
                    if dfs(index + 1):
                        return True
                    sums[i] -= nums[index]

            return False

        return dfs(0)

class Test(unittest.TestCase):
    def test(self):
        cases = utils.load_test_json(__file__).test_cases

        for case in cases:
            args = str(case.args)
            actual = Solution().makesquare(**case.args.__dict__)
            self.assertEqual(case.expected, actual, msg=args)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# (C) Copyright 2014 Voyager Search
#
# 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 os
import re
import sys
import tempfile
import urllib2
import xml.etree.cElementTree as eTree
import shutil
import requests
from utils import status
from utils import task_utils

status_writer = status.Writer()
import arcpy
if arcpy.GetInstallInfo()['Version'] == '10.0':
    raise ImportError('write_metadata not available with ArcGIS 10.0.')

errors_reasons = {}
skipped_reasons = {}

def index_item(id):
    """Re-indexes an item.
    :param id: Item's index ID
    """
    solr_url = "{0}/flags?op=add&flag=__to_extract&fq=id:({1})&fl=*,[true]".format(sys.argv[2].split('=')[1], id)
    requests.post(solr_url)

def execute(request):
    """Writes existing metadata for summary, description and tags.
    If overwrite is false, existing metadata is untouched unless any
    field is empty or does not exist, then it is created.

    :param request: json as a dict.
    """
    updated = 0
    errors = 0
    skipped = 0
    parameters = request['params']
    summary = task_utils.get_parameter_value(parameters, 'summary', 'value')
    description = task_utils.get_parameter_value(parameters, 'description', 'value')
    tags = task_utils.get_parameter_value(parameters, 'tags', 'value')
    # Handle commas, spaces, and/or new line separators.
    tags = [tag for tag in re.split(' |,|\n', tags) if not tag == '']
    overwrite = task_utils.get_parameter_value(parameters, 'overwrite', 'value')
    if not overwrite:
        overwrite = False
    if not os.path.exists(request['folder']):
        os.makedirs(request['folder'])

    # Stylesheet
    xslt_file = os.path.join(arcpy.GetInstallInfo()['InstallDir'], 'Metadata/Stylesheets/gpTools/exact copy of.xslt')

    # Template metadata file.
    template_xml = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'supportfiles', 'metadata_template.xml')

    num_results, response_index = task_utils.get_result_count(parameters)
    if num_results > task_utils.CHUNK_SIZE:
        # Query the index for results in groups of 25.
        query_index = task_utils.QueryIndex(parameters[response_index])
        fl = query_index.fl
        query = '{0}{1}{2}'.format(sys.argv[2].split('=')[1], '/select?&wt=json', fl)
        fq = query_index.get_fq()
        if fq:
            groups = task_utils.grouper(range(0, num_results), task_utils.CHUNK_SIZE, '')
            query += fq
        else:
            groups = task_utils.grouper(list(parameters[response_index]['ids']), task_utils.CHUNK_SIZE, '')

        status_writer.send_percent(0.0, _('Starting to process...'), 'write_metadata')
        i = 0.
        for group in groups:
            i += len(group) - group.count('')
            if fq:
                results = urllib2.urlopen(query + "&rows={0}&start={1}".format(task_utils.CHUNK_SIZE, group[0]))
            else:
                results = urllib2.urlopen(query + '{0}&ids={1}'.format(fl, ','.join(group)))

            input_items = task_utils.get_input_items(eval(results.read().replace('false', 'False').replace('true', 'True'))['response']['docs'], True)
            result = write_metadata(input_items, template_xml, xslt_file, summary, description, tags, overwrite)
            updated += result[0]
            errors += result[1]
            skipped += result[2]
            status_writer.send_percent(i / num_results, '{0}: {1:%}'.format("Processed", i / num_results), 'write_metadata')
    else:
        input_items = task_utils.get_input_items(parameters[response_index]['response']['docs'], True)
        updated, errors, skipped = write_metadata(input_items, template_xml, xslt_file,
                                                  summary, description, tags, overwrite, True)

    try:
        shutil.copy2(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'supportfiles', '_thumb.png'), request['folder'])
    except IOError:
        pass

    # Report state.
    if skipped > 0 or errors > 0:
        status_writer.send_state(status.STAT_WARNING, _('{0} results could not be processed').format(skipped + errors))
    else:
        status_writer.send_state(status.STAT_SUCCESS)
    task_utils.report(os.path.join(request['folder'], '_report.md'), updated, skipped, errors, errors_reasons, skipped_reasons)

def write_metadata(input_items, template_xml, xslt_file, summary, description, tags, overwrite, show_progress=False):
    """Writes metadata."""
    updated = 0
    errors = 0
    skipped = 0
    if show_progress:
        item_count = len(input_items)
        i = 1.

    for item in input_items:
        try:
            # Temporary XML file
            temp_xml = tempfile.NamedTemporaryFile(suffix='.xml', delete=True).name

            # Export xml
            try:
                arcpy.XSLTransform_conversion(item, xslt_file, temp_xml)
            except arcpy.ExecuteError:
                src_xml = os.path.join(arcpy.Describe(item).path, '{0}.xml'.format(os.path.basename(item)))
                shutil.copyfile(template_xml, src_xml)
                arcpy.XSLTransform_conversion(src_xml, xslt_file, temp_xml)

            # Read in XML
            tree = eTree.parse(temp_xml)
            root = tree.getroot()
            changes = 0

            # ISO allows many dataIdInfo groups; ArcGIS generally supports only one.
            data_id_elements = root.findall(".//dataIdInfo")
            if not data_id_elements:
                data_id_elements = [eTree.SubElement(root, 'dataIdInfo')]

            for data_id_element in data_id_elements:

                # Write summary.
                summary_element = root.findall(".//idPurp")
                if not summary_element:
                    summary_element = eTree.SubElement(data_id_element, 'idPurp')
                    summary_element.text = summary
                    changes += 1
                else:
                    for element in summary_element:
                        if summary and (overwrite or element.text is None):
                            element.text = summary
                            changes += 1

                # Write description.
                description_element = root.findall(".//idAbs")
                if not description_element:
                    description_element = eTree.SubElement(data_id_element, 'idAbs')
                    description_element.text = description
                    changes += 1
                else:
                    for element in description_element:
                        if description and (overwrite or element.text is None):
                            element.text = description
                            changes += 1

                # Write tags.
                tags = task_utils.get_unique_strings(tags)
                search_keys = root.findall(".//searchKeys")
                if not search_keys:
                    search_element = eTree.SubElement(data_id_element, 'searchKeys')
                    for tag in tags:
                        new_tag = eTree.SubElement(search_element, "keyword")
                        new_tag.text = tag
                        changes += 1
                elif not overwrite:
                    # Still add any new tags.
                    for search_element in search_keys:
                        if tags:
                            for tag in tags:
                                if tag.lower() not in [se.text.lower() for se in search_element.findall('.//keyword')]:
                                    new_tag = eTree.SubElement(search_element, "keyword")
                                    new_tag.text = tag
                                    changes += 1
                else:
                    if tags:
                        for search_element in search_keys:
                            [search_element.remove(e) for e in search_element.findall('.//keyword')]
                            for tag in tags:
                                new_tag = eTree.SubElement(search_element, "keyword")
                                new_tag.text = tag
                                changes += 1

            if changes > 0:
                # Save modifications to the temporary XML file.
                tree.write(temp_xml)
                # Import the XML file to the item; existing metadata is replaced.
                arcpy.MetadataImporter_conversion(temp_xml, item)
                if show_progress:
                    status_writer.send_percent(i / item_count, _('Metadata updated for: {0}').format(item), 'write_metadata')
                    i += 1
                else:
                    status_writer.send_status(_('Metadata updated for: {0}').format(item))

                try:
                    index_item(input_items[item][1])
                except (IndexError, urllib2.HTTPError, urllib2.URLError) as e:
                    status_writer.send_status(e)
                    pass
                updated += 1
            else:
                if show_progress:
                    status_writer.send_percent(i / item_count, _('No metadata changes for: {0}').format(item), 'write_metadata')
                    i += 1
                else:
                    status_writer.send_status(_('No metadata changes for: {0}').format(item))
                skipped_reasons[item] = _('No metadata changes for: {0}').format(item)
                skipped += 1
        except Exception as ex:
            if show_progress:
                status_writer.send_percent(i / item_count, _('Skipped: {0}').format(item), 'write_metadata')
                i += 1
            status_writer.send_status(_('FAIL: {0}').format(repr(ex)))
            errors_reasons[item] = repr(ex)
            errors += 1
            pass

    return updated, errors, skipped
<commit_msg>VG-2978: catch requests exceptions<commit_after># -*- coding: utf-8 -*-
# (C) Copyright 2014 Voyager Search
#
# 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 os
import re
import sys
import tempfile
import urllib2
import xml.etree.cElementTree as eTree
import shutil
import requests
from utils import status
from utils import task_utils

status_writer = status.Writer()
import arcpy
if arcpy.GetInstallInfo()['Version'] == '10.0':
    raise ImportError('write_metadata not available with ArcGIS 10.0.')

errors_reasons = {}
skipped_reasons = {}

def index_item(id):
    """Re-indexes an item.
    :param id: Item's index ID
    """
    solr_url = "{0}/flags?op=add&flag=__to_extract&fq=id:({1})&fl=*,[true]".format(sys.argv[2].split('=')[1], id)
    requests.post(solr_url)

def execute(request):
    """Writes existing metadata for summary, description and tags.
    If overwrite is false, existing metadata is untouched unless any
    field is empty or does not exist, then it is created.

    :param request: json as a dict.
    """
    updated = 0
    errors = 0
    skipped = 0
    parameters = request['params']
    summary = task_utils.get_parameter_value(parameters, 'summary', 'value')
    description = task_utils.get_parameter_value(parameters, 'description', 'value')
    tags = task_utils.get_parameter_value(parameters, 'tags', 'value')
    # Handle commas, spaces, and/or new line separators.
    tags = [tag for tag in re.split(' |,|\n', tags) if not tag == '']
    overwrite = task_utils.get_parameter_value(parameters, 'overwrite', 'value')
    if not overwrite:
        overwrite = False
    if not os.path.exists(request['folder']):
        os.makedirs(request['folder'])

    # Stylesheet
    xslt_file = os.path.join(arcpy.GetInstallInfo()['InstallDir'], 'Metadata/Stylesheets/gpTools/exact copy of.xslt')

    # Template metadata file.
    template_xml = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'supportfiles', 'metadata_template.xml')

    num_results, response_index = task_utils.get_result_count(parameters)
    if num_results > task_utils.CHUNK_SIZE:
        # Query the index for results in groups of 25.
        query_index = task_utils.QueryIndex(parameters[response_index])
        fl = query_index.fl
        query = '{0}{1}{2}'.format(sys.argv[2].split('=')[1], '/select?&wt=json', fl)
        fq = query_index.get_fq()
        if fq:
            groups = task_utils.grouper(range(0, num_results), task_utils.CHUNK_SIZE, '')
            query += fq
        else:
            groups = task_utils.grouper(list(parameters[response_index]['ids']), task_utils.CHUNK_SIZE, '')

        status_writer.send_percent(0.0, _('Starting to process...'), 'write_metadata')
        i = 0.
        for group in groups:
            i += len(group) - group.count('')
            if fq:
                results = urllib2.urlopen(query + "&rows={0}&start={1}".format(task_utils.CHUNK_SIZE, group[0]))
            else:
                results = urllib2.urlopen(query + '{0}&ids={1}'.format(fl, ','.join(group)))

            input_items = task_utils.get_input_items(eval(results.read().replace('false', 'False').replace('true', 'True'))['response']['docs'], True)
            result = write_metadata(input_items, template_xml, xslt_file, summary, description, tags, overwrite)
            updated += result[0]
            errors += result[1]
            skipped += result[2]
            status_writer.send_percent(i / num_results, '{0}: {1:%}'.format("Processed", i / num_results), 'write_metadata')
    else:
        input_items = task_utils.get_input_items(parameters[response_index]['response']['docs'], True)
        updated, errors, skipped = write_metadata(input_items, template_xml, xslt_file,
                                                  summary, description, tags, overwrite, True)

    try:
        shutil.copy2(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'supportfiles', '_thumb.png'), request['folder'])
    except IOError:
        pass

    # Report state.
    if skipped > 0 or errors > 0:
        status_writer.send_state(status.STAT_WARNING, _('{0} results could not be processed').format(skipped + errors))
    else:
        status_writer.send_state(status.STAT_SUCCESS)
    task_utils.report(os.path.join(request['folder'], '_report.md'), updated, skipped, errors, errors_reasons, skipped_reasons)

def write_metadata(input_items, template_xml, xslt_file, summary, description, tags, overwrite, show_progress=False):
    """Writes metadata."""
    updated = 0
    errors = 0
    skipped = 0
    if show_progress:
        item_count = len(input_items)
        i = 1.

    for item in input_items:
        try:
            # Temporary XML file
            temp_xml = tempfile.NamedTemporaryFile(suffix='.xml', delete=True).name

            # Export xml
            try:
                arcpy.XSLTransform_conversion(item, xslt_file, temp_xml)
            except arcpy.ExecuteError:
                src_xml = os.path.join(arcpy.Describe(item).path, '{0}.xml'.format(os.path.basename(item)))
                shutil.copyfile(template_xml, src_xml)
                arcpy.XSLTransform_conversion(src_xml, xslt_file, temp_xml)

            # Read in XML
            tree = eTree.parse(temp_xml)
            root = tree.getroot()
            changes = 0

            # ISO allows many dataIdInfo groups; ArcGIS generally supports only one.
            data_id_elements = root.findall(".//dataIdInfo")
            if not data_id_elements:
                data_id_elements = [eTree.SubElement(root, 'dataIdInfo')]

            for data_id_element in data_id_elements:

                # Write summary.
                summary_element = root.findall(".//idPurp")
                if not summary_element:
                    summary_element = eTree.SubElement(data_id_element, 'idPurp')
                    summary_element.text = summary
                    changes += 1
                else:
                    for element in summary_element:
                        if summary and (overwrite or element.text is None):
                            element.text = summary
                            changes += 1

                # Write description.
                description_element = root.findall(".//idAbs")
                if not description_element:
                    description_element = eTree.SubElement(data_id_element, 'idAbs')
                    description_element.text = description
                    changes += 1
                else:
                    for element in description_element:
                        if description and (overwrite or element.text is None):
                            element.text = description
                            changes += 1

                # Write tags.
                tags = task_utils.get_unique_strings(tags)
                search_keys = root.findall(".//searchKeys")
                if not search_keys:
                    search_element = eTree.SubElement(data_id_element, 'searchKeys')
                    for tag in tags:
                        new_tag = eTree.SubElement(search_element, "keyword")
                        new_tag.text = tag
                        changes += 1
                elif not overwrite:
                    # Still add any new tags.
                    for search_element in search_keys:
                        if tags:
                            for tag in tags:
                                if tag.lower() not in [se.text.lower() for se in search_element.findall('.//keyword')]:
                                    new_tag = eTree.SubElement(search_element, "keyword")
                                    new_tag.text = tag
                                    changes += 1
                else:
                    if tags:
                        for search_element in search_keys:
                            [search_element.remove(e) for e in search_element.findall('.//keyword')]
                            for tag in tags:
                                new_tag = eTree.SubElement(search_element, "keyword")
                                new_tag.text = tag
                                changes += 1

            if changes > 0:
                # Save modifications to the temporary XML file.
                tree.write(temp_xml)
                # Import the XML file to the item; existing metadata is replaced.
                arcpy.MetadataImporter_conversion(temp_xml, item)
                if show_progress:
                    status_writer.send_percent(i / item_count, _('Metadata updated for: {0}').format(item), 'write_metadata')
                    i += 1
                else:
                    status_writer.send_status(_('Metadata updated for: {0}').format(item))

                try:
                    index_item(input_items[item][1])
                except (IndexError, requests.HTTPError, requests.RequestException) as e:
                    status_writer.send_status(e.message)
                    pass
                updated += 1
            else:
                if show_progress:
                    status_writer.send_percent(i / item_count, _('No metadata changes for: {0}').format(item), 'write_metadata')
                    i += 1
                else:
                    status_writer.send_status(_('No metadata changes for: {0}').format(item))
                skipped_reasons[item] = _('No metadata changes for: {0}').format(item)
                skipped += 1
        except Exception as ex:
            if show_progress:
                status_writer.send_percent(i / item_count, _('Skipped: {0}').format(item), 'write_metadata')
                i += 1
            status_writer.send_status(_('FAIL: {0}').format(repr(ex)))
            errors_reasons[item] = repr(ex)
            errors += 1
            pass

    return updated, errors, skipped
<|endoftext|>


<commit_before>from decimal import Decimal
from django.contrib.auth.models import User, Group
from wellsfargo.tests.base import BaseTest
from wellsfargo.core.constants import TRANS_TYPE_AUTH, TRANS_APPROVED
from wellsfargo.models import APICredentials, TransferMetadata
import uuid

class APICredentialsTest(BaseTest):
    def test_selection_no_user(self):
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=None,
            priority=2)
        self.assertEqual(APICredentials.get_credentials().username, 'credsB')

    def test_selection_user_no_group(self):
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=None,
            priority=2)
        user = User.objects.create_user(username='bill')
        self.assertEqual(APICredentials.get_credentials(user).username, 'credsB')

    def test_selection_user_group(self):
        group = Group.objects.create(name='Special Group')
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=group,
            priority=2)
        user = User.objects.create_user(username='bill')
        self.assertEqual(APICredentials.get_credentials(user).username, 'credsA')
        user.groups.add(group)
        self.assertEqual(APICredentials.get_credentials(user).username, 'credsB')
        user.groups.remove(group)
        self.assertEqual(APICredentials.get_credentials(user).username, 'credsA')

class TransferMetadataTest(BaseTest):

    def test_account_number(self):
        transfer = TransferMetadata()
        transfer.user = self.joe
        transfer.credentials = self.credentials
        transfer.merchant_reference = uuid.uuid1()
        transfer.amount = Decimal('10.00')
        transfer.type_code = TRANS_TYPE_AUTH
        transfer.ticket_number = '123'
        transfer.status = TRANS_APPROVED
        transfer.message = 'message'
        transfer.disclosure = 'disclosure'
        transfer.save()

        # No account number is set
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxxxxxx')
        self.assertEqual(transfer.account_number, 'xxxxxxxxxxxxxxxx')

        # Set an account number
        transfer.account_number = '9999999999999991'
        transfer.save()

        # Retrieve account number via decryption
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '9991')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxx9991')
        self.assertEqual(transfer.account_number, '9999999999999991')

        # Purge encrypted copy of account number, leaving on the last 4 digits around
        transfer.purge_encrypted_account_number()

        # Make sure only the last 4 digits still exist
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '9991')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxx9991')
        self.assertEqual(transfer.account_number, 'xxxxxxxxxxxx9991')
<commit_msg>Fix conflicting usernames in test<commit_after>from decimal import Decimal
from django.contrib.auth.models import Group
from wellsfargo.tests.base import BaseTest
from wellsfargo.core.constants import TRANS_TYPE_AUTH, TRANS_APPROVED
from wellsfargo.models import APICredentials, TransferMetadata
import uuid

class APICredentialsTest(BaseTest):
    def test_selection_no_user(self):
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=None,
            priority=2)
        self.assertEqual(APICredentials.get_credentials().username, 'credsB')

    def test_selection_user_no_group(self):
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=None,
            priority=2)
        self.assertEqual(APICredentials.get_credentials(self.joe).username, 'credsB')

    def test_selection_user_group(self):
        group = Group.objects.create(name='Special Group')
        APICredentials.objects.create(
            username='credsA',
            password='',
            merchant_num='',
            user_group=None,
            priority=1)
        APICredentials.objects.create(
            username='credsB',
            password='',
            merchant_num='',
            user_group=group,
            priority=2)
        self.assertEqual(APICredentials.get_credentials(self.joe).username, 'credsA')
        self.joe.groups.add(group)
        self.assertEqual(APICredentials.get_credentials(self.joe).username, 'credsB')
        self.joe.groups.remove(group)
        self.assertEqual(APICredentials.get_credentials(self.joe).username, 'credsA')

class TransferMetadataTest(BaseTest):

    def test_account_number(self):
        transfer = TransferMetadata()
        transfer.user = self.joe
        transfer.credentials = self.credentials
        transfer.merchant_reference = uuid.uuid1()
        transfer.amount = Decimal('10.00')
        transfer.type_code = TRANS_TYPE_AUTH
        transfer.ticket_number = '123'
        transfer.status = TRANS_APPROVED
        transfer.message = 'message'
        transfer.disclosure = 'disclosure'
        transfer.save()

        # No account number is set
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxxxxxx')
        self.assertEqual(transfer.account_number, 'xxxxxxxxxxxxxxxx')

        # Set an account number
        transfer.account_number = '9999999999999991'
        transfer.save()

        # Retrieve account number via decryption
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '9991')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxx9991')
        self.assertEqual(transfer.account_number, '9999999999999991')

        # Purge encrypted copy of account number, leaving on the last 4 digits around
        transfer.purge_encrypted_account_number()

        # Make sure only the last 4 digits still exist
        transfer = TransferMetadata.objects.get(pk=transfer.pk)
        self.assertEqual(transfer.last4_account_number, '9991')
        self.assertEqual(transfer.masked_account_number, 'xxxxxxxxxxxx9991')
        self.assertEqual(transfer.account_number, 'xxxxxxxxxxxx9991')
<|endoftext|>


<commit_before>from django.contrib.gis.db import models
import os
from phonenumber_field.modelfields import PhoneNumberField

class Image(models.Model):
    """
    The Image model holds an image and related data.

    The Created and Modified time fields are created automatically by
    Django when the object is created or modified, and can not be altered.

    This model uses Django's built-ins for holding the image location and
    data in the database, as well as for keeping created and modified
    timestamps.
    """

    image = models.ImageField(upload_to='%Y/%m/%d')
    caption = models.TextField()

    created = models.DateTimeField(auto_now_add=True)
    modified = models.DateTimeField(auto_now=True)

class Vendor(models.Model):
    """
    The Vendor model holds the information for a vendor, including the
    geographic location as a pair of latitudinal/logitudinal coordinates,
    a street address, and an optional text description of their location
    (in case the address/coordinates are of, say, a dock instead of a shop).
    """
    pass

class Product(models.Model):
    """
    The Product model holds the information for a product, including the
    origin, season, market price, and availability.

    In addition, it holds a foreign key to the image and story related to the
    product.
    """
    name = models.TextField()
    variety = models.TextField()
    alt_name = models.TextField()
    description = models.TextField()
    origin = models.TextField()

    season = models.TextField()
    available = models.NullBooleanField()
    market_price = models.TextField()
    link = models.URLField()

    image_id = models.ForeignKey('Image')
    story_id = models.ForeignKey('Story')

    created = models.DateTimeField(auto_now_add=True)
    modified = models.DateTimeField(auto_now=True)

class Story(models.Model):
    """
    The story model holds the stories for products and vendors
    """
    pass
<commit_msg>Create Preparation model unit tests<commit_after>from django.contrib.gis.db import models
import os
from phonenumber_field.modelfields import PhoneNumberField

class Image(models.Model):
    """
    The Image model holds an image and related data.

    The Created and Modified time fields are created automatically by
    Django when the object is created or modified, and can not be altered.

    This model uses Django's built-ins for holding the image location and
    data in the database, as well as for keeping created and modified
    timestamps.
    """

    image = models.ImageField(upload_to='%Y/%m/%d')
    caption = models.TextField()

    created = models.DateTimeField(auto_now_add=True)
    modified = models.DateTimeField(auto_now=True)

class Vendor(models.Model):
    """
    The Vendor model holds the information for a vendor, including the
    geographic location as a pair of latitudinal/logitudinal coordinates,
    a street address, and an optional text description of their location
    (in case the address/coordinates are of, say, a dock instead of a shop).
    """
    pass

class Product(models.Model):
    """
    The Product model holds the information for a product, including the
    origin, season, market price, and availability.

    In addition, it holds a foreign key to the image and story related to the
    product.
    """
    name = models.TextField()
    variety = models.TextField()
    alt_name = models.TextField()
    description = models.TextField()
    origin = models.TextField()

    season = models.TextField()
    available = models.NullBooleanField()
    market_price = models.TextField()
    link = models.URLField()

    image_id = models.ForeignKey('Image')
    story_id = models.ForeignKey('Story')

    created = models.DateTimeField(auto_now_add=True)
    modified = models.DateTimeField(auto_now=True)

class Story(models.Model):
    """
    The story model holds the stories for products and vendors
    """
    pass

class Preparation(models.Model):
    """
    The Preparation model contains possible preparations of product, to be
    associated many-to-many with product (a product can have one or more
    preparations, preparations apply to many products). Preparations may be
    things like 'frozen', 'dried', 'fresh', 'live', etc, to be defined by
    Sea Grant data input.
    """

    pass
<|endoftext|>