<commit_before>#!/usr/bin/python -tt

import sys,getopt

import logging
import commands
import datetime
import json
import time
import pprint

logger = logging.getLogger('stencil')
hdlr = logging.StreamHandler(sys.stdout)
#hdlr = logging.FileHandler('stencil.log')
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)
logger.setLevel(logging.INFO) #logging.DEBUG

def help():

  print """

  Usage: ami-to-autoscale.py

  --instance-id YOUR-INSTANCE-ID
  --description "LIVE"
  --autoscale-group LIVEWEB
  --security-groups "sg-1 sg-2 sg-3"
  [--max N (default 20)]
  [--min N default 2]
  [--instance-size m3.medium]
  [--help]

  """

def Run(cmd):

  print "** Running:", cmd,

  (status,output) = commands.getstatusoutput(cmd)

  if status > 0:
    logger.error(output)
    sys.exit(2)

  print '[OK]'
  return output

def main(argv):

  instance=None
  description=None
  securityGroups=None
  autoscaleGroup=None
  maxInstances=10
  minInstances=2
  instanceSize="t2.micro"

  # make sure command line arguments are valid
  try:
    options, args = getopt.getopt(

       argv,
      'hv',
      [
        'help',
        'verbose',
        'instance-id=',
        'description=',
        'security-groups=',
        'autoscale-group=',
        'min=',
        'max=',
        'instance-size='

      ])

  except getopt.GetoptError:
    logging.fatal("Bad options!")
    help()
    sys.exit(2)

  # handle command line arugments
  for opt, arg in options:
    if opt in ('-h', '--help'):
      help()
      sys.exit(2)
    elif opt in ('-v', '--verbose'):
      logger.setLevel(logging.DEBUG)
    elif opt in ('', '--description'):
      description=arg
    elif opt in ('', '--instance-id'):
      instance=arg
    elif opt in ('', '--security-groups'):
      securityGroups=arg
    elif opt in ('', '--autoscale-group'):
      autoscaleGroup=arg
    elif opt in ('', '--instance-size'):
      instanceSize=arg
    elif opt in ('', '--min'):
      minInstances=int(arg)
    elif opt in ('', '--max'):
     maxInstances=int(arg)

  if None in [instance,description,securityGroups,autoscaleGroup]:
     help()
     sys.exit(2)

  ###################################
  # main code starts here
  ###################################

  # get current number of instances so we can double

  cmd = "aws autoscaling describe-auto-scaling-groups --auto-scaling-group-name %s" % autoscaleGroup
  output = Run(cmd)

  scaling_group_details = json.loads(output)

  try:
    instances = scaling_group_details['AutoScalingGroups'][0]['Instances']
  except:
    logger.error("FATAL: Could not get current instances count from auto scale group %s", )

  currentInstanceCount = len(instances)

  if currentInstanceCount < 1:
    currentInstanceCount = 1

  today = datetime.datetime.now().strftime("%Y-%m-%d-%s")

  tag = "%s-%s" % ( description, today )

  # create the new AMI

  amiImageName = "%s-AMI" % ( tag )

  cmd = """aws ec2 create-image --name "%s" --instance-id %s --description "%s" """ % (

      amiImageName,
      instance,
      amiImageName

  	)

  logger.info("Create AMI %s" % amiImageName )
  output = Run(cmd)

  # grab the json

  try:
    j = json.loads(output)
    ami = j['ImageId']
  except:
    logger.error("FATAL: Could not parse ImageId from JSON")
    logger.error("CMD: %s" % cmd)
    logger.error("JSON: %s" % j)
    sys.exit(2)

  # check status of AMI snapshot
  # wait for snapshot to finish

  iteration = 0
  limit=25

  while True:

    iteration = iteration + 1

    if iteration > limit:
      logger.error("%s snapshot timeout" % ami)

    cmd = "aws ec2 describe-images --image-ids %s " %  ami

    logger.info("Checking %s" % ami)

    output = Run(cmd)

    try:
      j = json.loads(output)
      state = j['Images'][0]['State']
    except:
      logger.error("Could not parse state")
      sys.exit(2)

    if state == "available":
      break
    if state == "failed":
      logger.error("AMI Creation failed")
      logger.error("OUTPUT: %s" % output)
      sys.exit(2)
    else:
      logger.info(state)

    time.sleep(30)

  # create launch configuration

  launchConfigName = "%s-LAUNCH-CONFIG" % ( tag )

  logger.info("Creating new launch config and wait 30 seconds...")

  cmd = """aws autoscaling create-launch-configuration --launch-configuration-name %s --image-id %s --instance-type %s --security-groups %s --instance-monitoring Enabled=true --associate-public-ip-address""" % (

    launchConfigName,
    ami,
    instanceSize,
    securityGroups

  )

  Run(cmd)

  # wait for launch configuration to be available
  time.sleep(30)

  logger.info("Update autoscale group and bump max size")

  # # switch the scaling group to use this AMI
  # # double the size of the scaling group

  cmd = """aws autoscaling update-auto-scaling-group --auto-scaling-group-name "%s" --launch-configuration-name %s --min-size %d --max-size %d""" % (

     autoscaleGroup,
     launchConfigName,
     currentInstanceCount*2, # doubles amount of instances so we can clear old ones when we scale down
     (currentInstanceCount*2)*2 # bump high during transition

  	)

  Run(cmd)

  # loop over all instances waiting for them to be available before descaling
  logger.info("Wait for scaleup instances to finishing launching...")

  # sleep 30 seconds for instances to show up
  time.sleep(60)

  # loop over all instances and check status
  healthy = True
  timeout = 700
  start = time.time()

  while True:

    logger.info("Checking new instances status...")

    for instance in instances:

      if instance['LifecycleState'] != 'InService':
        healthy = False

      if instance['HealthStatus'] != 'Healthy':
        healthy=False

    # all instances are good -- lets break
    if healthy:
      break

    # we have timed out
    if (time.time()-start) > timeout:
      logger.error("FATAL: Instances scale up timed out on health check")
      sys.exit(2)

    # dont spam
    time.sleep(15)

  # descale and set to desired counts (normal size)
  logger.info("Descale and set to desired counts")
  cmd = """aws autoscaling update-auto-scaling-group --auto-scaling-group-name "%s" --min-size %d --max-size %d""" % (

     autoscaleGroup,
     minInstances,
     maxInstances

  	)

  Run(cmd)

  logger.info("Success")

if __name__ == "__main__":
  main(sys.argv[1:])
<commit_msg>Bumped AMI creation timeout to 15 minutes<commit_after>#!/usr/bin/python -tt

import sys,getopt

import logging
import commands
import datetime
import json
import time
import pprint

logger = logging.getLogger('stencil')
hdlr = logging.StreamHandler(sys.stdout)
#hdlr = logging.FileHandler('stencil.log')
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)
logger.setLevel(logging.INFO) #logging.DEBUG

def help():

  print """

  Usage: ami-to-autoscale.py

  --instance-id YOUR-INSTANCE-ID
  --description "LIVE"
  --autoscale-group LIVEWEB
  --security-groups "sg-1 sg-2 sg-3"
  [--max N (default 20)]
  [--min N default 2]
  [--instance-size m3.medium]
  [--help]

  """

def Run(cmd):

  print "** Running:", cmd,

  (status,output) = commands.getstatusoutput(cmd)

  if status > 0:
    logger.error(output)
    sys.exit(2)

  print '[OK]'
  return output

def main(argv):

  instance=None
  description=None
  securityGroups=None
  autoscaleGroup=None
  maxInstances=10
  minInstances=2
  instanceSize="t2.micro"

  # make sure command line arguments are valid
  try:
    options, args = getopt.getopt(

       argv,
      'hv',
      [
        'help',
        'verbose',
        'instance-id=',
        'description=',
        'security-groups=',
        'autoscale-group=',
        'min=',
        'max=',
        'instance-size='

      ])

  except getopt.GetoptError:
    logging.fatal("Bad options!")
    help()
    sys.exit(2)

  # handle command line arugments
  for opt, arg in options:
    if opt in ('-h', '--help'):
      help()
      sys.exit(2)
    elif opt in ('-v', '--verbose'):
      logger.setLevel(logging.DEBUG)
    elif opt in ('', '--description'):
      description=arg
    elif opt in ('', '--instance-id'):
      instance=arg
    elif opt in ('', '--security-groups'):
      securityGroups=arg
    elif opt in ('', '--autoscale-group'):
      autoscaleGroup=arg
    elif opt in ('', '--instance-size'):
      instanceSize=arg
    elif opt in ('', '--min'):
      minInstances=int(arg)
    elif opt in ('', '--max'):
     maxInstances=int(arg)

  if None in [instance,description,securityGroups,autoscaleGroup]:
     help()
     sys.exit(2)

  ###################################
  # main code starts here
  ###################################

  # get current number of instances so we can double

  cmd = "aws autoscaling describe-auto-scaling-groups --auto-scaling-group-name %s" % autoscaleGroup
  output = Run(cmd)

  scaling_group_details = json.loads(output)

  try:
    instances = scaling_group_details['AutoScalingGroups'][0]['Instances']
  except:
    logger.error("FATAL: Could not get current instances count from auto scale group %s", )

  currentInstanceCount = len(instances)

  if currentInstanceCount < 1:
    currentInstanceCount = 1

  today = datetime.datetime.now().strftime("%Y-%m-%d-%s")

  tag = "%s-%s" % ( description, today )

  # create the new AMI

  amiImageName = "%s-AMI" % ( tag )

  cmd = """aws ec2 create-image --name "%s" --instance-id %s --description "%s" """ % (

      amiImageName,
      instance,
      amiImageName

  	)

  logger.info("Create AMI %s" % amiImageName )
  output = Run(cmd)

  # grab the json

  try:
    j = json.loads(output)
    ami = j['ImageId']
  except:
    logger.error("FATAL: Could not parse ImageId from JSON")
    logger.error("CMD: %s" % cmd)
    logger.error("JSON: %s" % j)
    sys.exit(2)

  # check status of AMI snapshot
  # wait for snapshot to finish

  iteration = 0
  limit=30 # 15 minutes 30 @ 30 sec intervals

  while True:

    iteration = iteration + 1

    if iteration > limit:
      logger.error("%s snapshot timeout" % ami)

    cmd = "aws ec2 describe-images --image-ids %s " %  ami

    logger.info("Checking %s" % ami)

    output = Run(cmd)

    try:
      j = json.loads(output)
      state = j['Images'][0]['State']
    except:
      logger.error("Could not parse state")
      sys.exit(2)

    if state == "available":
      break
    if state == "failed":
      logger.error("AMI Creation failed")
      logger.error("OUTPUT: %s" % output)
      sys.exit(2)
    else:
      logger.info(state)

    time.sleep(30)

  # create launch configuration

  launchConfigName = "%s-LAUNCH-CONFIG" % ( tag )

  logger.info("Creating new launch config and wait 30 seconds...")

  cmd = """aws autoscaling create-launch-configuration --launch-configuration-name %s --image-id %s --instance-type %s --security-groups %s --instance-monitoring Enabled=true --associate-public-ip-address""" % (

    launchConfigName,
    ami,
    instanceSize,
    securityGroups

  )

  Run(cmd)

  # wait for launch configuration to be available
  time.sleep(30)

  logger.info("Update autoscale group and bump max size")

  # # switch the scaling group to use this AMI
  # # double the size of the scaling group

  cmd = """aws autoscaling update-auto-scaling-group --auto-scaling-group-name "%s" --launch-configuration-name %s --min-size %d --max-size %d""" % (

     autoscaleGroup,
     launchConfigName,
     currentInstanceCount*2, # doubles amount of instances so we can clear old ones when we scale down
     (currentInstanceCount*2)*2 # bump high during transition

  	)

  Run(cmd)

  # loop over all instances waiting for them to be available before descaling
  logger.info("Wait for scaleup instances to finishing launching...")

  # sleep 30 seconds for instances to show up
  time.sleep(60)

  # loop over all instances and check status
  healthy = True
  timeout = 700
  start = time.time()

  while True:

    logger.info("Checking new instances status...")

    for instance in instances:

      if instance['LifecycleState'] != 'InService':
        healthy = False

      if instance['HealthStatus'] != 'Healthy':
        healthy=False

    # all instances are good -- lets break
    if healthy:
      break

    # we have timed out
    if (time.time()-start) > timeout:
      logger.error("FATAL: Instances scale up timed out on health check")
      sys.exit(2)

    # dont spam
    time.sleep(15)

  # descale and set to desired counts (normal size)
  logger.info("Descale and set to desired counts")
  cmd = """aws autoscaling update-auto-scaling-group --auto-scaling-group-name "%s" --min-size %d --max-size %d""" % (

     autoscaleGroup,
     minInstances,
     maxInstances

  	)

  Run(cmd)

  logger.info("Success")

if __name__ == "__main__":
  main(sys.argv[1:])
<|endoftext|>


<commit_before>##
# Copyright (c) 2005-2014 Apple 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.
##

import os
import sys

from twext.python.filepath import CachingFilePath as FilePath
from twisted.internet import reactor
from twisted.internet.defer import inlineCallbacks, Deferred, returnValue

from twistedcaldav.config import config
from twistedcaldav.directory.directory import DirectoryError
from twistedcaldav.directory import calendaruserproxy

from twistedcaldav.test.util import TestCase, CapturingProcessProtocol, \
    ErrorOutput

from calendarserver.tap.util import directoryFromConfig
from calendarserver.tools.principals import (parseCreationArgs, matchStrings,
    updateRecord, principalForPrincipalID, getProxies, setProxies)

class ManagePrincipalsTestCase(TestCase):

    def setUp(self):
        super(ManagePrincipalsTestCase, self).setUp()

        # Since this test operates on proxy db, we need to assign the service:
        calendaruserproxy.ProxyDBService = calendaruserproxy.ProxySqliteDB(os.path.abspath(self.mktemp()))

        testRoot = os.path.join(os.path.dirname(__file__), "principals")
        templateName = os.path.join(testRoot, "caldavd.plist")
        templateFile = open(templateName)
        template = templateFile.read()
        templateFile.close()

        databaseRoot = os.path.abspath("_spawned_scripts_db" + str(os.getpid()))
        newConfig = template % {
            "ServerRoot" : os.path.abspath(config.ServerRoot),
            "DataRoot" : os.path.abspath(config.DataRoot),
            "DatabaseRoot" : databaseRoot,
            "DocumentRoot" : os.path.abspath(config.DocumentRoot),
            "LogRoot" : os.path.abspath(config.LogRoot),
        }
        configFilePath = FilePath(os.path.join(config.ConfigRoot, "caldavd.plist"))
        configFilePath.setContent(newConfig)

        self.configFileName = configFilePath.path
        config.load(self.configFileName)

        origUsersFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "users-groups.xml"))
        copyUsersFile = FilePath(os.path.join(config.DataRoot, "accounts.xml"))
        origUsersFile.copyTo(copyUsersFile)

        origResourcesFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "resources-locations.xml"))
        copyResourcesFile = FilePath(os.path.join(config.DataRoot, "resources.xml"))
        origResourcesFile.copyTo(copyResourcesFile)

        origAugmentFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "augments.xml"))
        copyAugmentFile = FilePath(os.path.join(config.DataRoot, "augments.xml"))
        origAugmentFile.copyTo(copyAugmentFile)

        # Make sure trial puts the reactor in the right state, by letting it
        # run one reactor iteration.  (Ignore me, please.)
        d = Deferred()
        reactor.callLater(0, d.callback, True)
        return d

    @inlineCallbacks
    def runCommand(self, *additional):
        """
        Run calendarserver_manage_principals, passing additional as args.
        """
        sourceRoot = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
        python = sys.executable
        script = os.path.join(sourceRoot, "bin", "calendarserver_manage_principals")

        args = [python, script, "-f", self.configFileName]
        args.extend(additional)
        cwd = sourceRoot

        deferred = Deferred()
        reactor.spawnProcess(CapturingProcessProtocol(deferred, None), python, args, env=os.environ, path=cwd)
        output = yield deferred
        returnValue(output)

    @inlineCallbacks
    def test_help(self):
        results = yield self.runCommand("--help")
        self.assertTrue(results.startswith("usage:"))

    @inlineCallbacks
    def test_principalTypes(self):
        results = yield self.runCommand("--list-principal-types")
        self.assertTrue("groups" in results)
        self.assertTrue("users" in results)
        self.assertTrue("locations" in results)
        self.assertTrue("resources" in results)

    @inlineCallbacks
    def test_listPrincipals(self):
        results = yield self.runCommand("--list-principals=users")
        for i in xrange(1, 10):
            self.assertTrue("user%02d" % (i,) in results)

    @inlineCallbacks
    def test_search(self):
        results = yield self.runCommand("--search=user")
        self.assertTrue("10 matches found" in results)
        for i in xrange(1, 10):
            self.assertTrue("user%02d" % (i,) in results)

    @inlineCallbacks
    def test_addRemove(self):
        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource", "edaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Added 'New Resource'" in results)

        results = yield self.runCommand("--get-auto-schedule",
            "resources:newresource")
        self.assertTrue(results.startswith('Auto-schedule for "New Resource" (resources:newresource) is true'))

        results = yield self.runCommand("--get-auto-schedule-mode",
            "resources:newresource")
        self.assertTrue(results.startswith('Auto-schedule mode for "New Resource" (resources:newresource) is default'))

        results = yield self.runCommand("--list-principals=resources")
        self.assertTrue("newresource" in results)

        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource1", "edaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Duplicate guid" in results)

        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource", "fdaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Duplicate shortName" in results)

        results = yield self.runCommand("--remove", "resources:newresource")
        self.assertTrue("Removed 'New Resource'" in results)

        results = yield self.runCommand("--list-principals=resources")
        self.assertFalse("newresource" in results)

    def test_parseCreationArgs(self):

        self.assertEquals(("full name", None, None),
            parseCreationArgs(("full name",)))

        self.assertEquals(("full name", "short name", None),
            parseCreationArgs(("full name", "short name")))

        guid = "02C3DE93-E655-4856-47B76B8BB1A7BDCE"

        self.assertEquals(("full name", "short name", guid),
            parseCreationArgs(("full name", "short name", guid)))

        self.assertEquals(("full name", "short name", guid),
            parseCreationArgs(("full name", guid, "short name")))

        self.assertEquals(("full name", None, guid),
            parseCreationArgs(("full name", guid)))

        self.assertRaises(
            ValueError,
            parseCreationArgs, ("full name", "non guid", "non guid")
        )

    def test_matchStrings(self):
        self.assertEquals("abc", matchStrings("a", ("abc", "def")))
        self.assertEquals("def", matchStrings("de", ("abc", "def")))
        self.assertRaises(
            ValueError,
            matchStrings, "foo", ("abc", "def")
        )

    @inlineCallbacks
    def test_modifyWriteProxies(self):
        results = yield self.runCommand("--add-write-proxy=users:user01",
            "locations:location01")
        self.assertTrue(results.startswith('Added "Test User 01" (users:user01) as a write proxy for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--list-write-proxies",
            "locations:location01")
        self.assertTrue("Test User 01" in results)

        results = yield self.runCommand("--remove-proxy=users:user01",
            "locations:location01")

        results = yield self.runCommand("--list-write-proxies",
            "locations:location01")
        self.assertTrue('No write proxies for "Room 01" (locations:location01)' in results)

    @inlineCallbacks
    def test_modifyReadProxies(self):
        results = yield self.runCommand("--add-read-proxy=users:user01",
            "locations:location01")
        self.assertTrue(results.startswith('Added "Test User 01" (users:user01) as a read proxy for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--list-read-proxies",
            "locations:location01")
        self.assertTrue("Test User 01" in results)

        results = yield self.runCommand("--remove-proxy=users:user01",
            "locations:location01")

        results = yield self.runCommand("--list-read-proxies",
            "locations:location01")
        self.assertTrue('No read proxies for "Room 01" (locations:location01)' in results)

    @inlineCallbacks
    def test_autoSchedule(self):
        results = yield self.runCommand("--get-auto-schedule",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule for "Room 01" (locations:location01) is false'))

        results = yield self.runCommand("--set-auto-schedule=true",
            "locations:location01")
        self.assertTrue(results.startswith('Setting auto-schedule to true for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--get-auto-schedule",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule for "Room 01" (locations:location01) is true'))

        results = yield self.runCommand("--set-auto-schedule=true",
            "users:user01")
        self.assertTrue(results.startswith('Enabling auto-schedule for (users)user01 is not allowed.'))

    @inlineCallbacks
    def test_autoScheduleMode(self):
        results = yield self.runCommand("--get-auto-schedule-mode",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule mode for "Room 01" (locations:location01) is default'))

        results = yield self.runCommand("--set-auto-schedule-mode=accept-if-free",
            "locations:location01")
        self.assertTrue(results.startswith('Setting auto-schedule mode to accept-if-free for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--get-auto-schedule-mode",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule mode for "Room 01" (locations:location01) is accept-if-free'))

        results = yield self.runCommand("--set-auto-schedule-mode=decline-if-busy",
            "users:user01")
        self.assertTrue(results.startswith('Setting auto-schedule mode for (users)user01 is not allowed.'))

        try:
            results = yield self.runCommand("--set-auto-schedule-mode=bogus",
                "users:user01")
        except ErrorOutput:
            pass
        else:
            self.fail("Expected command failure")

    @inlineCallbacks
    def test_updateRecord(self):
        directory = directoryFromConfig(config)
        guid = "EEE28807-A8C5-46C8-A558-A08281C558A7"

        (yield updateRecord(True, directory, "locations",
            guid=guid, fullName="Test Location", shortNames=["testlocation", ],)
        )
        try:
            (yield updateRecord(True, directory, "locations",
                guid=guid, fullName="Test Location", shortNames=["testlocation", ],)
            )
        except DirectoryError:
            # We're expecting an error for trying to create a record with
            # an existing GUID
            pass
        else:
            raise self.failureException("Duplicate guid expected")

        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test Location")
        self.assertTrue(record.autoSchedule)

        (yield updateRecord(False, directory, "locations",
            guid=guid, fullName="Changed", shortNames=["testlocation", ],)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Changed")

        directory.destroyRecord("locations", guid=guid)
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is None)

        # Create a user, change autoSchedule
        guid = "F0DE73A8-39D4-4830-8D32-1FA03ABA3470"
        (yield updateRecord(True, directory, "users",
            guid=guid, fullName="Test User", shortNames=["testuser", ],
            autoSchedule=True)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test User")
        self.assertTrue(record.autoSchedule)

        (yield updateRecord(False, directory, "users",
            guid=guid, fullName="Test User", shortNames=["testuser", ],
            autoSchedule=False)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test User")
        self.assertFalse(record.autoSchedule)

    @inlineCallbacks
    def test_setProxies(self):
        """
        Read and Write proxies can be set en masse
        """
        directory = directoryFromConfig(config)

        principal = principalForPrincipalID("users:user01", directory=directory)
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(readProxies, []) # initially empty
        self.assertEquals(writeProxies, []) # initially empty

        (yield setProxies(None, principal, ["users:user03", "users:user04"], ["users:user05"], directory=directory))
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(set(readProxies), set(["user03", "user04"]))
        self.assertEquals(set(writeProxies), set(["user05"]))

        # Using None for a proxy list indicates a no-op
        (yield setProxies(None, principal, [], None, directory=directory))
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(readProxies, []) # now empty
        self.assertEquals(set(writeProxies), set(["user05"])) # unchanged
<commit_msg>Don't invoke gateway tool with python<commit_after>##
# Copyright (c) 2005-2014 Apple 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.
##

import os
import sys

from twext.python.filepath import CachingFilePath as FilePath
from twisted.internet import reactor
from twisted.internet.defer import inlineCallbacks, Deferred, returnValue

from twistedcaldav.config import config
from twistedcaldav.directory.directory import DirectoryError
from twistedcaldav.directory import calendaruserproxy

from twistedcaldav.test.util import TestCase, CapturingProcessProtocol, \
    ErrorOutput

from calendarserver.tap.util import directoryFromConfig
from calendarserver.tools.principals import (parseCreationArgs, matchStrings,
    updateRecord, principalForPrincipalID, getProxies, setProxies)

class ManagePrincipalsTestCase(TestCase):

    def setUp(self):
        super(ManagePrincipalsTestCase, self).setUp()

        # Since this test operates on proxy db, we need to assign the service:
        calendaruserproxy.ProxyDBService = calendaruserproxy.ProxySqliteDB(os.path.abspath(self.mktemp()))

        testRoot = os.path.join(os.path.dirname(__file__), "principals")
        templateName = os.path.join(testRoot, "caldavd.plist")
        templateFile = open(templateName)
        template = templateFile.read()
        templateFile.close()

        databaseRoot = os.path.abspath("_spawned_scripts_db" + str(os.getpid()))
        newConfig = template % {
            "ServerRoot" : os.path.abspath(config.ServerRoot),
            "DataRoot" : os.path.abspath(config.DataRoot),
            "DatabaseRoot" : databaseRoot,
            "DocumentRoot" : os.path.abspath(config.DocumentRoot),
            "LogRoot" : os.path.abspath(config.LogRoot),
        }
        configFilePath = FilePath(os.path.join(config.ConfigRoot, "caldavd.plist"))
        configFilePath.setContent(newConfig)

        self.configFileName = configFilePath.path
        config.load(self.configFileName)

        origUsersFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "users-groups.xml"))
        copyUsersFile = FilePath(os.path.join(config.DataRoot, "accounts.xml"))
        origUsersFile.copyTo(copyUsersFile)

        origResourcesFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "resources-locations.xml"))
        copyResourcesFile = FilePath(os.path.join(config.DataRoot, "resources.xml"))
        origResourcesFile.copyTo(copyResourcesFile)

        origAugmentFile = FilePath(os.path.join(os.path.dirname(__file__),
            "principals", "augments.xml"))
        copyAugmentFile = FilePath(os.path.join(config.DataRoot, "augments.xml"))
        origAugmentFile.copyTo(copyAugmentFile)

        # Make sure trial puts the reactor in the right state, by letting it
        # run one reactor iteration.  (Ignore me, please.)
        d = Deferred()
        reactor.callLater(0, d.callback, True)
        return d

    @inlineCallbacks
    def runCommand(self, *additional):
        """
        Run calendarserver_manage_principals, passing additional as args.
        """
        sourceRoot = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
        cmd = os.path.join(sourceRoot, "bin", "calendarserver_manage_principals")

        args = [cmd, "-f", self.configFileName]
        args.extend(additional)
        cwd = sourceRoot

        deferred = Deferred()
        reactor.spawnProcess(CapturingProcessProtocol(deferred, None), cmd, args, env=os.environ, path=cwd)
        output = yield deferred
        returnValue(output)

    @inlineCallbacks
    def test_help(self):
        results = yield self.runCommand("--help")
        self.assertTrue(results.startswith("usage:"))

    @inlineCallbacks
    def test_principalTypes(self):
        results = yield self.runCommand("--list-principal-types")
        self.assertTrue("groups" in results)
        self.assertTrue("users" in results)
        self.assertTrue("locations" in results)
        self.assertTrue("resources" in results)

    @inlineCallbacks
    def test_listPrincipals(self):
        results = yield self.runCommand("--list-principals=users")
        for i in xrange(1, 10):
            self.assertTrue("user%02d" % (i,) in results)

    @inlineCallbacks
    def test_search(self):
        results = yield self.runCommand("--search=user")
        self.assertTrue("10 matches found" in results)
        for i in xrange(1, 10):
            self.assertTrue("user%02d" % (i,) in results)

    @inlineCallbacks
    def test_addRemove(self):
        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource", "edaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Added 'New Resource'" in results)

        results = yield self.runCommand("--get-auto-schedule",
            "resources:newresource")
        self.assertTrue(results.startswith('Auto-schedule for "New Resource" (resources:newresource) is true'))

        results = yield self.runCommand("--get-auto-schedule-mode",
            "resources:newresource")
        self.assertTrue(results.startswith('Auto-schedule mode for "New Resource" (resources:newresource) is default'))

        results = yield self.runCommand("--list-principals=resources")
        self.assertTrue("newresource" in results)

        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource1", "edaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Duplicate guid" in results)

        results = yield self.runCommand("--add", "resources", "New Resource",
            "newresource", "fdaa6ae6-011b-4d89-ace3-6b688cdd91d9")
        self.assertTrue("Duplicate shortName" in results)

        results = yield self.runCommand("--remove", "resources:newresource")
        self.assertTrue("Removed 'New Resource'" in results)

        results = yield self.runCommand("--list-principals=resources")
        self.assertFalse("newresource" in results)

    def test_parseCreationArgs(self):

        self.assertEquals(("full name", None, None),
            parseCreationArgs(("full name",)))

        self.assertEquals(("full name", "short name", None),
            parseCreationArgs(("full name", "short name")))

        guid = "02C3DE93-E655-4856-47B76B8BB1A7BDCE"

        self.assertEquals(("full name", "short name", guid),
            parseCreationArgs(("full name", "short name", guid)))

        self.assertEquals(("full name", "short name", guid),
            parseCreationArgs(("full name", guid, "short name")))

        self.assertEquals(("full name", None, guid),
            parseCreationArgs(("full name", guid)))

        self.assertRaises(
            ValueError,
            parseCreationArgs, ("full name", "non guid", "non guid")
        )

    def test_matchStrings(self):
        self.assertEquals("abc", matchStrings("a", ("abc", "def")))
        self.assertEquals("def", matchStrings("de", ("abc", "def")))
        self.assertRaises(
            ValueError,
            matchStrings, "foo", ("abc", "def")
        )

    @inlineCallbacks
    def test_modifyWriteProxies(self):
        results = yield self.runCommand("--add-write-proxy=users:user01",
            "locations:location01")
        self.assertTrue(results.startswith('Added "Test User 01" (users:user01) as a write proxy for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--list-write-proxies",
            "locations:location01")
        self.assertTrue("Test User 01" in results)

        results = yield self.runCommand("--remove-proxy=users:user01",
            "locations:location01")

        results = yield self.runCommand("--list-write-proxies",
            "locations:location01")
        self.assertTrue('No write proxies for "Room 01" (locations:location01)' in results)

    @inlineCallbacks
    def test_modifyReadProxies(self):
        results = yield self.runCommand("--add-read-proxy=users:user01",
            "locations:location01")
        self.assertTrue(results.startswith('Added "Test User 01" (users:user01) as a read proxy for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--list-read-proxies",
            "locations:location01")
        self.assertTrue("Test User 01" in results)

        results = yield self.runCommand("--remove-proxy=users:user01",
            "locations:location01")

        results = yield self.runCommand("--list-read-proxies",
            "locations:location01")
        self.assertTrue('No read proxies for "Room 01" (locations:location01)' in results)

    @inlineCallbacks
    def test_autoSchedule(self):
        results = yield self.runCommand("--get-auto-schedule",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule for "Room 01" (locations:location01) is false'))

        results = yield self.runCommand("--set-auto-schedule=true",
            "locations:location01")
        self.assertTrue(results.startswith('Setting auto-schedule to true for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--get-auto-schedule",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule for "Room 01" (locations:location01) is true'))

        results = yield self.runCommand("--set-auto-schedule=true",
            "users:user01")
        self.assertTrue(results.startswith('Enabling auto-schedule for (users)user01 is not allowed.'))

    @inlineCallbacks
    def test_autoScheduleMode(self):
        results = yield self.runCommand("--get-auto-schedule-mode",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule mode for "Room 01" (locations:location01) is default'))

        results = yield self.runCommand("--set-auto-schedule-mode=accept-if-free",
            "locations:location01")
        self.assertTrue(results.startswith('Setting auto-schedule mode to accept-if-free for "Room 01" (locations:location01)'))

        results = yield self.runCommand("--get-auto-schedule-mode",
            "locations:location01")
        self.assertTrue(results.startswith('Auto-schedule mode for "Room 01" (locations:location01) is accept-if-free'))

        results = yield self.runCommand("--set-auto-schedule-mode=decline-if-busy",
            "users:user01")
        self.assertTrue(results.startswith('Setting auto-schedule mode for (users)user01 is not allowed.'))

        try:
            results = yield self.runCommand("--set-auto-schedule-mode=bogus",
                "users:user01")
        except ErrorOutput:
            pass
        else:
            self.fail("Expected command failure")

    @inlineCallbacks
    def test_updateRecord(self):
        directory = directoryFromConfig(config)
        guid = "EEE28807-A8C5-46C8-A558-A08281C558A7"

        (yield updateRecord(True, directory, "locations",
            guid=guid, fullName="Test Location", shortNames=["testlocation", ],)
        )
        try:
            (yield updateRecord(True, directory, "locations",
                guid=guid, fullName="Test Location", shortNames=["testlocation", ],)
            )
        except DirectoryError:
            # We're expecting an error for trying to create a record with
            # an existing GUID
            pass
        else:
            raise self.failureException("Duplicate guid expected")

        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test Location")
        self.assertTrue(record.autoSchedule)

        (yield updateRecord(False, directory, "locations",
            guid=guid, fullName="Changed", shortNames=["testlocation", ],)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Changed")

        directory.destroyRecord("locations", guid=guid)
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is None)

        # Create a user, change autoSchedule
        guid = "F0DE73A8-39D4-4830-8D32-1FA03ABA3470"
        (yield updateRecord(True, directory, "users",
            guid=guid, fullName="Test User", shortNames=["testuser", ],
            autoSchedule=True)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test User")
        self.assertTrue(record.autoSchedule)

        (yield updateRecord(False, directory, "users",
            guid=guid, fullName="Test User", shortNames=["testuser", ],
            autoSchedule=False)
        )
        record = directory.recordWithGUID(guid)
        self.assertTrue(record is not None)
        self.assertEquals(record.fullName, "Test User")
        self.assertFalse(record.autoSchedule)

    @inlineCallbacks
    def test_setProxies(self):
        """
        Read and Write proxies can be set en masse
        """
        directory = directoryFromConfig(config)

        principal = principalForPrincipalID("users:user01", directory=directory)
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(readProxies, []) # initially empty
        self.assertEquals(writeProxies, []) # initially empty

        (yield setProxies(None, principal, ["users:user03", "users:user04"], ["users:user05"], directory=directory))
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(set(readProxies), set(["user03", "user04"]))
        self.assertEquals(set(writeProxies), set(["user05"]))

        # Using None for a proxy list indicates a no-op
        (yield setProxies(None, principal, [], None, directory=directory))
        readProxies, writeProxies = (yield getProxies(principal, directory=directory))
        self.assertEquals(readProxies, []) # now empty
        self.assertEquals(set(writeProxies), set(["user05"])) # unchanged
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import
from datetime import datetime
from itertools import product
import logging
import os
import os.path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from ..utils import make_same_length, setdbspl, int2srt

log = logging.getLogger(__name__)

class Experiment(object):
    """
    Performs a speech intelligibility experiment.

    The masker is truncated to the length of the target, or is padded with
    zeros if it is shorter than the target.

    Parameters
    ----------
    models : single model, or list
        List of intelligibility models.
    materials : object
        object that implements a `next` interface that returns the next
        pair of target and maskers.
    snrs : array_like
        List of SNRs.
    name : str,
        name of experiment, will be appended to the date when writing to file.
    write : bool, optional
        Write the result to file, as CSV, the default is True.
    output_path : string, optional.
        Path where the results will be written if `write` is True. The
        default is './Data'.
    timestamp_format : str, optional
        Datetime timestamp format for the CSV file name. The default is of
        the form YYYYMMDD-HHMMSS.

    """

    def __init__(
            self,
            models,
            material,
            snrs,
            distortion=None,
            dist_params=(None,),
            fixed_level=65,
            fixed_target=True,
            name=None,
            write=True,
            output_path='./Data/',
            timestamp_format="%Y%m%d-%H%M%S"
    ):
        """@todo: to be defined1. """
        self.models = models
        self.material = material
        self.snrs = snrs
        self.distortion = distortion
        self.dist_params = dist_params
        self.fixed_level = fixed_level
        self.fixed_target = fixed_target
        self.name = name
        self.timestamp_format = timestamp_format
        self.write = write
        self.output_path = output_path
        self._key_full_pred = 'Full Prediction'
        self._key_value = 'Value'
        self._key_output = 'Output'
        self._key_dist_params = "Distortion params"
        self._key_models = 'Model'
        self._key_snr = 'SNR'
        self._key_sent = 'Sentence number'
        self._all_keys = [
            self._key_full_pred,
            self._key_value,
            self._key_output,
            self._key_dist_params,
            self._key_models,
            self._key_snr,
            self._key_sent
        ]

    def preprocessing(self, target, masker, snr, params):
        """
        Applies preprocessing to the target and masker before setting the
        levels. In this case, the masker is padded with zeros if it is longer
        than the target, or it is truncated to be the same length as the target.

        :param target:
        :param masker:
        :return:
        """

        if params:
            target, masker = self.distortion(target, masker, *params)

        # Make target and masker same length
        if target.shape[-1] != masker.shape[-1]:
            target, masker = make_same_length(target, masker,
                                              extend_first=False)
        if params:
            if isinstance(params, dict):
                target, masker = self.distortion(target, masker, **params)
            else:
                target, masker = self.distortion(target, masker, *params)
        return target, target + masker, masker

    def adjust_levels(self, target, masker, snr):
        """
        Adjusts level of target and maskers.

        Uses the `self.fixed_level` as the reference level for the target and
        masker. If `self.fixed_target` is True, the masker level is varied to
        set the required SNR, otherwise the target level is changed.

        :param target: ndarray
            Target signal.
        :param masker: ndarray
            Masker signal.
        :param snr: float
            SNR at which to set the target and masker.
        :return: tuple
            Level adjusted `target` and `masker`.
        """

        target_level = self.fixed_level
        masker_level = self.fixed_level
        if self.fixed_target:
            masker_level -= snr
        else:
            target_level += snr
        target = setdbspl(target, target_level)
        masker = setdbspl(masker, masker_level)
        return target, masker

    def next_masker(self, target):
        return self.material.ssn(target)

    def append_results(
            self,
            df,
            res,
            model,
            snr,
            i_target,
            params
    ):
        """
        Appends results to a DataFrame

        Parameters
        ----------
        df : dataframe
            DataFrame where the new results will be appended.
        res : dict
            Output dictionary from an intelligibility model.
        model: object
            Intelligibility model. Will use it's `name` attribute,
            if available, to add the source model to the DataFrame. Otherwise,
            the `__class__.__name__` attribute will be used.
        snr : float
            SNR at which the simulation was performed.
        i_target : int
            Number of the target sentence
        params : object
            Parameters that were passed to the distortion process.

        Returns
        -------
        df : dataframe
            DataFrame with new entry appended.
        """
        try:
            model_name = model.name
        except AttributeError:
            model_name = model.__class__.__name__
        d = {
            'SNR': snr
            , 'Model': model_name
            , 'Sentence number': i_target
            , self._key_full_pred: res
        }
        # If the distortion parameters are in a dictionary, put each value in
        # a different column. Otherwise, group everything in a single column.
        if isinstance(params, dict):
            for k, v in params.iteritems():
                d[k] = v
        else:
            # Make sure the values are hashable for later manipulation
            if isinstance(params, list):
                params = tuple(params)
            else:
                pass
            d['Distortion params'] = params

        for name, value in res['p'].iteritems():
            d['Output'] = name
            d['Value'] = value
            df = df.append(d, ignore_index=True)

        return df

    def run(self, n=None, seed=0):
        """ Run the experiment.

        Parameters
        ----------
        n : int
            Number of sentences to process.

        Returns
        -------
        df : pd.Dataframe
            Pandas dataframe with the experimental results.

        """
        if not seed:
            seed = 0
        np.random.seed(seed)

        try:
            self.models = iter(self.models)
        except TypeError:
            self.models = [self.models]

        targets = self.material.load_files(n)

        # Initialize the dataframe in which the results are saved.
        df = pd.DataFrame()

        for ii, ((i_target, target), params, snr, model) \
                in enumerate(product(
                enumerate(targets),
                self.dist_params,
                self.snrs,
                self.models
        )):
            masker = self.next_masker(target)

            target, mix, masker = self.preprocessing(
                target,
                masker,
                snr,
                params
            )

            target, masker = self.adjust_levels(target, masker, snr)

            log.info("Simulation # {}\t SNR: {}, sentence {}", ii, snr, i_target)
            res = self.prediction(model, target, masker)

            df = self.append_results(
                df,
                res,
                model,
                snr,
                i_target,
                params
            )

        if self.write:
            self._write_results(df)
        return df

    def _write_results(self, df):
        """Writes results to CSV file.

        Will drop the column where all the complete model output is stored
        before writing to disk.

        Parameters
        ----------
        df : dataframe

        Returns
        -------
        filepath : str
            Path to the CSV file.

        Raises
        ------
        IOError : Raise if the path where to write the CSV file is not
        accessible. Additionally, the function tries to save the CSV file to
        the current directory, in order not to loose the simulation data.

        """
        timestamp = datetime.now()
        date = timestamp.strftime(self.timestamp_format)
        if self.name:
            name = "-{}".format(self.name)
        else:
            name = ''
        filename = "{date}{name}.csv".format(date=date, name=name)

        output_file = os.path.join(self.output_path, filename)
        try:
            df.drop(self._key_full_pred, axis=1).to_csv(output_file)
            log.info('Saved CSV file to location: %s'.format(output_file))
        except IOError as e:
            try:
                alternate_path = os.path.join(os.getcwd(), filename)
                err_msg = 'Could not write CSV file to path: {}, tried to ' \
                          'save to ' \
                          '{} in order not to loose data.'.format(
                    output_file, alternate_path)
                log.error(err_msg)
                raise
            finally:
                try:
                    df.drop(self._key_full_pred, axis=1).to_csv(alternate_path)
                except:
                    pass
        else:
            return output_file

    @staticmethod
    def prediction(model, target, masker):
        """
        Predicts intelligibility for a target and masker pair. The target and
        masker are simply added together to create the mixture.

        Parameters
        ----------
        model :
        target :
        masker :

        Returns
        -------
        :return:
        """
        return model.predict(target, target + masker, masker)

    def plot_results(self, df):
        df = df.ix[self.df['Target distance'] == 0.5, :]
        for key, grp in df.groupby(['Target distance', 'Masker distance']):
            plt.plot(grp['SNR'].unique(),
                     grp.groupby('SNR')['Intelligibility_L'].mean(),
                     label=str(key) + '_L')
            plt.plot(grp['SNR'].unique(),
                     grp.groupby('SNR')['Intelligibility_R'].mean(),
                     label=str(key) + '_R')
        srt_l = int2srt(grp['SNR'].unique(),
                        grp.groupby('SNR')['Intelligibility_L'].mean().values)
        srt_r = int2srt(grp['SNR'].unique(),
                        grp.groupby('SNR')['Intelligibility_R'].mean().values)
        try:
            print('Positions:%s\tSRTs L:%.1f\t R:%.1f' %
                  (key, srt_l[0], srt_r[0]))
        except:
            pass
        plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left',
                   borderaxespad=0.)
        plt.xlabel('SNR (dB)')
        plt.ylabel('% Intelligibility')

    @staticmethod
    def snr_to_pc(df, col, fc, out_name='Intelligibility'):
        """Converts the data in a given column to percent correct.

        Parameters
        ----------
        col : string
            Name of the column to convert to intelligibility.
        fc : function,
            Takes a float as input and returns a float. The function performs
            the transformation to intelligibility.
        out_name : str
            Name of the output column (default: 'Intelligibility')

        Returns
        -------
        df : dataframe
            Dataframe with the new columns with trans
        """
        df[out_name] = df[col].map(fc)
        return df

    @staticmethod
    def get_srt(df):
        """Converts SRTs to DeltaSRTs.

        :return: tuple, srts and DeltaSRTs
        """

        scores = []
        df_tn = df[df['Target distance'] == 0.5]
        for ii, (key, grp) in enumerate(df_tn.groupby(['Target distance',
                                                       'Masker distance'])):
            scores.append(grp.groupby('SNR')['Intelligibility_L'].mean())

        srts = np.zeros(4)
        for ii, score in enumerate(scores):
            srts[ii] = int2srt(score.index, score)

        print(srts)
        dsrts = srts[1] - srts
        print(dsrts)
        return srts, dsrts

def srt_dict_to_dataframe(d):
    df_srts = pd.DataFrame()
    for k, v in d.iteritems():
        model, material, tdist, mdist = k.split('_')

        try:
            srt = v[0]
        except TypeError:
            srt = np.nan

        df_srts = df_srts.append({'model': model,
                                  'tidst': tdist,
                                  'mdist': mdist,
                                  'srt': srt},
                                 ignore_index=True)
    df_srts = df_srts.convert_objects(convert_numeric=True, )
    return df_srts.sort(['model', 'mdist'])

def plot_srt_dataframe(df):
    for key, grp in df.groupby('model'):
        plt.plot(grp['mdist'], grp['srt'], label=key)
    plt.legend(loc='best')
    plt.xlabel('Masker distance')
    plt.ylabel('SRT (dB)')
<commit_msg>Perform distortion and level adjust in preproc<commit_after># -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import
from datetime import datetime
from itertools import product
import logging
import os
import os.path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from ..utils import make_same_length, setdbspl, int2srt

log = logging.getLogger(__name__)

class Experiment(object):
    """
    Performs a speech intelligibility experiment.

    The masker is truncated to the length of the target, or is padded with
    zeros if it is shorter than the target.

    Parameters
    ----------
    models : single model, or list
        List of intelligibility models.
    materials : object
        object that implements a `next` interface that returns the next
        pair of target and maskers.
    snrs : array_like
        List of SNRs.
    name : str,
        name of experiment, will be appended to the date when writing to file.
    write : bool, optional
        Write the result to file, as CSV, the default is True.
    output_path : string, optional.
        Path where the results will be written if `write` is True. The
        default is './Data'.
    timestamp_format : str, optional
        Datetime timestamp format for the CSV file name. The default is of
        the form YYYYMMDD-HHMMSS.

    """

    def __init__(
            self,
            models,
            material,
            snrs,
            distortion=None,
            dist_params=(None,),
            fixed_level=65,
            fixed_target=True,
            name=None,
            write=True,
            output_path='./Data/',
            timestamp_format="%Y%m%d-%H%M%S"
    ):
        """@todo: to be defined1. """
        self.models = models
        self.material = material
        self.snrs = snrs
        self.distortion = distortion
        self.dist_params = dist_params
        self.fixed_level = fixed_level
        self.fixed_target = fixed_target
        self.name = name
        self.timestamp_format = timestamp_format
        self.write = write
        self.output_path = output_path
        self._key_full_pred = 'Full Prediction'
        self._key_value = 'Value'
        self._key_output = 'Output'
        self._key_dist_params = "Distortion params"
        self._key_models = 'Model'
        self._key_snr = 'SNR'
        self._key_sent = 'Sentence number'
        self._all_keys = [
            self._key_full_pred,
            self._key_value,
            self._key_output,
            self._key_dist_params,
            self._key_models,
            self._key_snr,
            self._key_sent
        ]

    def preprocessing(self, target, masker, snr, params):
        """
        Applies preprocessing to the target and masker before setting the
        levels. In this case, the masker is padded with zeros if it is longer
        than the target, or it is truncated to be the same length as the target.

        :param target:
        :param masker:
        :return:
        """

        # Make target and masker same length
        if target.shape[-1] != masker.shape[-1]:
            target, masker = make_same_length(target, masker,
                                              extend_first=False)

        if params:
            if isinstance(params, dict):
                target, masker = self.distortion(target, masker, **params)
            else:
                target, masker = self.distortion(target, masker, *params)

        target, masker = self.adjust_levels(target, masker, snr)
        return target, target + masker, masker

    def adjust_levels(self, target, masker, snr):
        """
        Adjusts level of target and maskers.

        Uses the `self.fixed_level` as the reference level for the target and
        masker. If `self.fixed_target` is True, the masker level is varied to
        set the required SNR, otherwise the target level is changed.

        :param target: ndarray
            Target signal.
        :param masker: ndarray
            Masker signal.
        :param snr: float
            SNR at which to set the target and masker.
        :return: tuple
            Level adjusted `target` and `masker`.
        """

        target_level = self.fixed_level
        masker_level = self.fixed_level
        if self.fixed_target:
            masker_level -= snr
        else:
            target_level += snr
        target = setdbspl(target, target_level)
        masker = setdbspl(masker, masker_level)
        return target, masker

    def next_masker(self, target):
        return self.material.ssn(target)

    def append_results(
            self,
            df,
            res,
            model,
            snr,
            i_target,
            params
    ):
        """
        Appends results to a DataFrame

        Parameters
        ----------
        df : dataframe
            DataFrame where the new results will be appended.
        res : dict
            Output dictionary from an intelligibility model.
        model: object
            Intelligibility model. Will use it's `name` attribute,
            if available, to add the source model to the DataFrame. Otherwise,
            the `__class__.__name__` attribute will be used.
        snr : float
            SNR at which the simulation was performed.
        i_target : int
            Number of the target sentence
        params : object
            Parameters that were passed to the distortion process.

        Returns
        -------
        df : dataframe
            DataFrame with new entry appended.
        """
        try:
            model_name = model.name
        except AttributeError:
            model_name = model.__class__.__name__
        d = {
            'SNR': snr
            , 'Model': model_name
            , 'Sentence number': i_target
            , self._key_full_pred: res
        }
        # If the distortion parameters are in a dictionary, put each value in
        # a different column. Otherwise, group everything in a single column.
        if isinstance(params, dict):
            for k, v in params.iteritems():
                d[k] = v
        else:
            # Make sure the values are hashable for later manipulation
            if isinstance(params, list):
                params = tuple(params)
            else:
                pass
            d['Distortion params'] = params

        for name, value in res['p'].iteritems():
            d['Output'] = name
            d['Value'] = value
            df = df.append(d, ignore_index=True)

        return df

    def run(self, n=None, seed=0):
        """ Run the experiment.

        Parameters
        ----------
        n : int
            Number of sentences to process.

        Returns
        -------
        df : pd.Dataframe
            Pandas dataframe with the experimental results.

        """
        if not seed:
            seed = 0
        np.random.seed(seed)

        try:
            self.models = iter(self.models)
        except TypeError:
            self.models = [self.models]

        targets = self.material.load_files(n)

        # Initialize the dataframe in which the results are saved.
        df = pd.DataFrame()

        for ii, ((i_target, target), params, snr, model) \
                in enumerate(product(
                enumerate(targets),
                self.dist_params,
                self.snrs,
                self.models
        )):
            masker = self.next_masker(target)

            target, mix, masker = self.preprocessing(
                target,
                masker,
                snr,
                params
            )

            log.info("Simulation # %s\t SNR: %s, sentence %s", ii, snr,
                     i_target)
            res = self.prediction(model, target, mix, masker)

            df = self.append_results(
                df,
                res,
                model,
                snr,
                i_target,
                params
            )

        if self.write:
            self._write_results(df)
        return df

    def _write_results(self, df):
        """Writes results to CSV file.

        Will drop the column where all the complete model output is stored
        before writing to disk.

        Parameters
        ----------
        df : dataframe

        Returns
        -------
        filepath : str
            Path to the CSV file.

        Raises
        ------
        IOError : Raise if the path where to write the CSV file is not
        accessible. Additionally, the function tries to save the CSV file to
        the current directory, in order not to loose the simulation data.

        """
        timestamp = datetime.now()
        date = timestamp.strftime(self.timestamp_format)
        if self.name:
            name = "-{}".format(self.name)
        else:
            name = ''
        filename = "{date}{name}.csv".format(date=date, name=name)

        output_file = os.path.join(self.output_path, filename)
        try:
            df.drop(self._key_full_pred, axis=1).to_csv(output_file)
            log.info('Saved CSV file to location: %s'.format(output_file))
        except IOError as e:
            try:
                alternate_path = os.path.join(os.getcwd(), filename)
                err_msg = 'Could not write CSV file to path: {}, tried to ' \
                          'save to ' \
                          '{} in order not to loose data.'.format(
                    output_file, alternate_path)
                log.error(err_msg)
                raise
            finally:
                try:
                    df.drop(self._key_full_pred, axis=1).to_csv(alternate_path)
                except:
                    pass
        else:
            return output_file

    @staticmethod
    def prediction(model, target, masker):
        """
        Predicts intelligibility for a target and masker pair. The target and
        masker are simply added together to create the mixture.

        Parameters
        ----------
        model :
        target :
        masker :

        Returns
        -------
        :return:
        """
        return model.predict(target, target + masker, masker)

    def plot_results(self, df):
        df = df.ix[self.df['Target distance'] == 0.5, :]
        for key, grp in df.groupby(['Target distance', 'Masker distance']):
            plt.plot(grp['SNR'].unique(),
                     grp.groupby('SNR')['Intelligibility_L'].mean(),
                     label=str(key) + '_L')
            plt.plot(grp['SNR'].unique(),
                     grp.groupby('SNR')['Intelligibility_R'].mean(),
                     label=str(key) + '_R')
        srt_l = int2srt(grp['SNR'].unique(),
                        grp.groupby('SNR')['Intelligibility_L'].mean().values)
        srt_r = int2srt(grp['SNR'].unique(),
                        grp.groupby('SNR')['Intelligibility_R'].mean().values)
        try:
            print('Positions:%s\tSRTs L:%.1f\t R:%.1f' %
                  (key, srt_l[0], srt_r[0]))
        except:
            pass
        plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left',
                   borderaxespad=0.)
        plt.xlabel('SNR (dB)')
        plt.ylabel('% Intelligibility')

    @staticmethod
    def snr_to_pc(df, col, fc, out_name='Intelligibility'):
        """Converts the data in a given column to percent correct.

        Parameters
        ----------
        col : string
            Name of the column to convert to intelligibility.
        fc : function,
            Takes a float as input and returns a float. The function performs
            the transformation to intelligibility.
        out_name : str
            Name of the output column (default: 'Intelligibility')

        Returns
        -------
        df : dataframe
            Dataframe with the new columns with trans
        """
        df[out_name] = df[col].map(fc)
        return df

    @staticmethod
    def get_srt(df):
        """Converts SRTs to DeltaSRTs.

        :return: tuple, srts and DeltaSRTs
        """

        scores = []
        df_tn = df[df['Target distance'] == 0.5]
        for ii, (key, grp) in enumerate(df_tn.groupby(['Target distance',
                                                       'Masker distance'])):
            scores.append(grp.groupby('SNR')['Intelligibility_L'].mean())

        srts = np.zeros(4)
        for ii, score in enumerate(scores):
            srts[ii] = int2srt(score.index, score)

        print(srts)
        dsrts = srts[1] - srts
        print(dsrts)
        return srts, dsrts

def srt_dict_to_dataframe(d):
    df_srts = pd.DataFrame()
    for k, v in d.iteritems():
        model, material, tdist, mdist = k.split('_')

        try:
            srt = v[0]
        except TypeError:
            srt = np.nan

        df_srts = df_srts.append({'model': model,
                                  'tidst': tdist,
                                  'mdist': mdist,
                                  'srt': srt},
                                 ignore_index=True)
    df_srts = df_srts.convert_objects(convert_numeric=True, )
    return df_srts.sort(['model', 'mdist'])

def plot_srt_dataframe(df):
    for key, grp in df.groupby('model'):
        plt.plot(grp['mdist'], grp['srt'], label=key)
    plt.legend(loc='best')
    plt.xlabel('Masker distance')
    plt.ylabel('SRT (dB)')
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
import argparse
import os.path
from zstacklib import *
from datetime import datetime

def add_true_in_command(cmd):
    return "%s || true" % cmd

# create log
logger_dir = "/var/log/zstack/"
create_log(logger_dir)
banner("Starting to deploy zstack network agent")
start_time = datetime.now()
# set default value
pip_url = "https=//pypi.python.org/simple/"
proxy = ""
sproxy = ""
chroot_env = 'false'
zstack_repo = 'false'
current_dir = os.path.dirname(os.path.realpath(__file__))
file_root = "files/zsnagentansible"
pkg_zsn = ""
post_url = ""
chrony_servers = None
fs_rootpath = ""
remote_user = "root"
remote_pass = None
remote_port = None
require_python_env = "false"
tmout = None

# get parameter from shell
parser = argparse.ArgumentParser(description='Deploy zsn-agent to host')
parser.add_argument('-i', type=str, help="""specify inventory host file
                        default=/etc/ansible/hosts""")
parser.add_argument('--private-key', type=str, help='use this file to authenticate the connection')
parser.add_argument('-e', type=str, help='set additional variables as key=value or YAML/JSON')
args = parser.parse_args()
argument_dict = eval(args.e)

# update the variable from shell arguments
locals().update(argument_dict)
zsn_root = "%s/zsn-agent/package" % zstack_root

host_post_info = HostPostInfo()
host_post_info.host_inventory = args.i
host_post_info.host = host
host_post_info.post_url = post_url
host_post_info.chrony_servers = chrony_servers
host_post_info.private_key = args.private_key
host_post_info.remote_user = remote_user
host_post_info.remote_pass = remote_pass
host_post_info.remote_port = remote_port
if remote_pass is not None and remote_user != 'root':
    host_post_info.become = True

# include zstacklib.py
(distro, distro_version, distro_release, _) = get_remote_host_info(host_post_info)
zstacklib_args = ZstackLibArgs()
zstacklib_args.distro = distro
zstacklib_args.distro_release = distro_release
zstacklib_args.distro_version = distro_version
zstacklib_args.zstack_repo = zstack_repo
zstacklib_args.yum_server = yum_server
zstacklib_args.zstack_root = zstack_root
zstacklib_args.host_post_info = host_post_info
zstacklib_args.pip_url = pip_url
zstacklib_args.trusted_host = trusted_host
zstacklib_args.require_python_env = require_python_env
zstacklib = ZstackLib(zstacklib_args)

if distro in RPM_BASED_OS:
    if zstack_repo == 'false':
        yum_install_package("libpcap", host_post_info)
    else:
        command = ("pkg_list=`rpm -q %s | grep \"not installed\" | awk '{ print $2 }'` && for pkg in $pkg_list; do yum "
                   "--disablerepo=* --enablerepo=%s install -y $pkg; done;") % ("libpcap", zstack_repo)
        run_remote_command(command, host_post_info)

elif distro in DEB_BASED_OS:
    apt_install_packages(["libpcap"], host_post_info)

else:
    error("ERROR: Unsupported distribution")

run_remote_command(add_true_in_command("rm -rf %s/*" % zsn_root), host_post_info)
command = 'mkdir -p %s ' % (zsn_root)
run_remote_command(add_true_in_command(command), host_post_info)

# name: copy zsn binary
copy_arg = CopyArg()
dest_pkg = "%s/%s" % (zsn_root, pkg_zsn)
copy_arg.src = "%s/%s" % (file_root, pkg_zsn)
copy_arg.dest = dest_pkg
copy(copy_arg, host_post_info)

# name: install zstack-network
command = "bash %s %s " % (dest_pkg, fs_rootpath)
run_remote_command(add_true_in_command(command), host_post_info)

# integrate zstack-network with systemd
run_remote_command(add_true_in_command("/bin/cp -f /usr/local/zstack/zsn-agent/bin/zstack-network-agent.service /usr/lib/systemd/system/"), host_post_info)
run_remote_command(add_true_in_command("pkill zsn-agent; rm /etc/init.d/zstack-network-agent"), host_post_info)

if tmout is None:
    tmout = 960
service_env = "'ZSNARGS=-log-file /var/log/zstack/zsn-agent/zsn-agent.log -tmout %s'" % int(tmout)
service_env = service_env.replace("/", "\/")
command = "sed -i \"s/.*Environment=.*/Environment=%s/g\" /usr/lib/systemd/system/zstack-network-agent.service" % service_env
run_remote_command(add_true_in_command(command), host_post_info)

command = "systemctl daemon-reload && systemctl enable zstack-network-agent.service && systemctl restart zstack-network-agent.service"
run_remote_command(add_true_in_command(command), host_post_info)

host_post_info.start_time = start_time
handle_ansible_info("SUCC: Deploy zstack network agent successful", host_post_info, "INFO")

sys.exit(0)

<commit_msg>Avoid useless agent restart<commit_after>#!/usr/bin/env python
# encoding: utf-8
import argparse
import commands
import os.path
from zstacklib import *
from datetime import datetime

def add_true_in_command(cmd):
    return "%s || true" % cmd

# create log
logger_dir = "/var/log/zstack/"
create_log(logger_dir)
banner("Starting to deploy zstack network agent")
start_time = datetime.now()
# set default value
pip_url = "https=//pypi.python.org/simple/"
proxy = ""
sproxy = ""
chroot_env = 'false'
zstack_repo = 'false'
current_dir = os.path.dirname(os.path.realpath(__file__))
file_root = "files/zsnagentansible"
pkg_zsn = ""
post_url = ""
chrony_servers = None
fs_rootpath = ""
remote_user = "root"
remote_pass = None
remote_port = None
require_python_env = "false"
tmout = None

# get parameter from shell
parser = argparse.ArgumentParser(description='Deploy zsn-agent to host')
parser.add_argument('-i', type=str, help="""specify inventory host file
                        default=/etc/ansible/hosts""")
parser.add_argument('--private-key', type=str, help='use this file to authenticate the connection')
parser.add_argument('-e', type=str, help='set additional variables as key=value or YAML/JSON')
args = parser.parse_args()
argument_dict = eval(args.e)

# update the variable from shell arguments
locals().update(argument_dict)
zsn_root = "%s/zsn-agent/package" % zstack_root

host_post_info = HostPostInfo()
host_post_info.host_inventory = args.i
host_post_info.host = host
host_post_info.post_url = post_url
host_post_info.chrony_servers = chrony_servers
host_post_info.private_key = args.private_key
host_post_info.remote_user = remote_user
host_post_info.remote_pass = remote_pass
host_post_info.remote_port = remote_port
if remote_pass is not None and remote_user != 'root':
    host_post_info.become = True

# include zstacklib.py
(distro, distro_version, distro_release, _) = get_remote_host_info(host_post_info)
zstacklib_args = ZstackLibArgs()
zstacklib_args.distro = distro
zstacklib_args.distro_release = distro_release
zstacklib_args.distro_version = distro_version
zstacklib_args.zstack_repo = zstack_repo
zstacklib_args.yum_server = yum_server
zstacklib_args.zstack_root = zstack_root
zstacklib_args.host_post_info = host_post_info
zstacklib_args.pip_url = pip_url
zstacklib_args.trusted_host = trusted_host
zstacklib_args.require_python_env = require_python_env
zstacklib = ZstackLib(zstacklib_args)

if distro in RPM_BASED_OS:
    if zstack_repo == 'false':
        yum_install_package("libpcap", host_post_info)
    else:
        command = ("pkg_list=`rpm -q %s | grep \"not installed\" | awk '{ print $2 }'` && for pkg in $pkg_list; do yum "
                   "--disablerepo=* --enablerepo=%s install -y $pkg; done;") % ("libpcap", zstack_repo)
        run_remote_command(command, host_post_info)

elif distro in DEB_BASED_OS:
    apt_install_packages(["libpcap"], host_post_info)

else:
    error("ERROR: Unsupported distribution")

run_remote_command(add_true_in_command("rm -rf %s/*" % zsn_root), host_post_info)
command = 'mkdir -p %s ' % (zsn_root)
run_remote_command(add_true_in_command(command), host_post_info)

# name: copy zsn binary
copy_arg = CopyArg()
dest_pkg = "%s/%s" % (zsn_root, pkg_zsn)
copy_arg.src = "%s/%s" % (file_root, pkg_zsn)
copy_arg.dest = dest_pkg
copy(copy_arg, host_post_info)

# name: install zstack-network
command = "bash %s %s " % (dest_pkg, fs_rootpath)
run_remote_command(add_true_in_command(command), host_post_info)

# integrate zstack-network with systemd
run_remote_command(add_true_in_command("/bin/cp -f /usr/local/zstack/zsn-agent/bin/zstack-network-agent.service /usr/lib/systemd/system/"), host_post_info)

if tmout is None:
    tmout = 960

successTmout, stdoutTmout = run_remote_command(add_true_in_command("grep -- '-tmout %s' /usr/lib/systemd/system/zstack-network-agent.service" % int(tmout)), host_post_info, True, True)
successMd5, stdoutMd5 = run_remote_command(add_true_in_command("md5sum /var/lib/zstack/zsn-agent/package/zsn-agent.bin"), host_post_info, True, True)
successLocalMd5, localMd5 = commands.getstatusoutput("md5sum /usr/local/zstack/ansible/files/zsnagentansible/zsn-agent.bin")

msg = Msg()
post_url = host_post_info.post_url
msg.details = "zsn agent deploying %s, %s, %s, %s, %s, %s" % (successTmout, stdoutTmout, successMd5, stdoutMd5, successLocalMd5, localMd5)
msg.label = "ansible.weiw"
msg.parameters = host_post_info.post_label_param
msg.level = "WARNING"
post_msg(msg, post_url)

if successTmout is True and successMd5 is True and successLocalMd5 == 0 and \
    len(stdoutTmout.strip()) != 0 and stdoutMd5.split(" ")[0] == localMd5.split(" ")[0]:
    host_post_info.start_time = start_time
    handle_ansible_info("SUCC: Deploy zstack network agent successful", host_post_info, "INFO")
    sys.exit(0)

run_remote_command(add_true_in_command("pkill zsn-agent; rm /etc/init.d/zstack-network-agent"), host_post_info)

service_env = "'ZSNARGS=-log-file /var/log/zstack/zsn-agent/zsn-agent.log -tmout %s'" % int(tmout)
service_env = service_env.replace("/", "\/")
command = "sed -i \"s/.*Environment=.*/Environment=%s/g\" /usr/lib/systemd/system/zstack-network-agent.service" % service_env
run_remote_command(add_true_in_command(command), host_post_info)

command = "systemctl daemon-reload && systemctl enable zstack-network-agent.service && systemctl restart zstack-network-agent.service"
run_remote_command(add_true_in_command(command), host_post_info)

host_post_info.start_time = start_time
handle_ansible_info("SUCC: Deploy zstack network agent successful", host_post_info, "INFO")

sys.exit(0)

<|endoftext|>


<commit_before>__version__ = '0.1.0'
<commit_msg>Bump version to 0.3.0<commit_after>__version__ = '0.3.0'
<|endoftext|>


<commit_before>import os
import os.path
import fnmatch
import logging
import ycm_core

BASE_FLAGS = [
    '-Wall',
    '-Wextra',
    '-Werror',
    '-Wc++98-compat',
    '-Wno-long-long',
    '-Wno-variadic-macros',
    '-fexceptions',
    '-DNDEBUG',
    '-std=c++98',
    '-xc++',
    '-I/usr/include/'
]

SOURCE_EXTENSIONS = [
    '.cpp',
    '.cxx',
    '.cc',
    '.c',
    '.m',
    '.mm'
]

HEADER_EXTENSIONS = [
    '.h',
    '.hxx',
    '.hpp',
    '.hh'
]

def IsHeaderFile(filename):
    extension = os.path.splitext(filename)[1]
    return extension in HEADER_EXTENSIONS

def GetCompilationInfoForFile(database, filename):
    if IsHeaderFile(filename):
        basename = os.path.splitext(filename)[0]
        for extension in SOURCE_EXTENSIONS:
            replacement_file = basename + extension
            if os.path.exists(replacement_file):
                compilation_info = database.GetCompilationInfoForFile(replacement_file)
                if compilation_info.compiler_flags_:
                    return compilation_info
        return None
    return database.GetCompilationInfoForFile(filename)

def FindNearest(path, target):
    candidate = os.path.join(path, target)
    if(os.path.isfile(candidate) or os.path.isdir(candidate)):
        logging.info("Found nearest " + target + " at " + candidate)
        return candidate;
    else:
        parent = os.path.dirname(os.path.abspath(path));
        if(parent == path):
            raise RuntimeError("Could not find " + target);
        return FindNearest(parent, target)

def MakeRelativePathsInFlagsAbsolute(flags, working_directory):
    if not working_directory:
        return list(flags)
    new_flags = []
    make_next_absolute = False
    path_flags = [ '-isystem', '-I', '-iquote', '--sysroot=' ]
    for flag in flags:
        new_flag = flag

        if make_next_absolute:
            make_next_absolute = False
            if not flag.startswith('/'):
                new_flag = os.path.join(working_directory, flag)

        for path_flag in path_flags:
            if flag == path_flag:
                make_next_absolute = True
                break

            if flag.startswith(path_flag):
                path = flag[ len(path_flag): ]
                new_flag = path_flag + os.path.join(working_directory, path)
                break

        if new_flag:
            new_flags.append(new_flag)
    return new_flags

def FlagsForClangComplete(root):
    try:
        clang_complete_path = FindNearest(root, '.clang_complete')
        clang_complete_flags = open(clang_complete_path, 'r').read().splitlines()
        return clang_complete_flags
    except:
        return None

def FlagsForInclude(root):
    try:
        include_path = FindNearest(root, 'include')
        flags = []
        for dirroot, dirnames, filenames in os.walk(include_path):
            for dir_path in dirnames:
                real_path = os.path.join(dirroot, dir_path)
                flags = flags + ['-I', real_path]
        return flags
    except:
        return None

def FlagsForCompilationDatabase(root, filename):
    try:
        compilation_db_path = FindNearest(root, 'compile_commands.json')
        compilation_db_dir = os.path.dirname(compilation_db_path)
        logging.info("Set compilation database directory to " + compilation_db_dir)
        compilation_db =  ycm_core.CompilationDatabase(compilation_db_dir)
        if not compilation_db:
            logging.info("Compilation database file found but unable to load")
            return None
        compilation_info = GetCompilationInfoForFile(compilation_db, filename)
        if not compilation_info:
            logging.info("No compilation info for " + filename + " in compilation database")
            return None
        return MakeRelativePathsInFlagsAbsolute(
            compilation_info.compiler_flags_,
            compilation_info.compiler_working_dir_)
    except:
        return None

def FlagsForFile(filename):
    root = os.path.realpath(filename);
    compilation_db_flags = FlagsForCompilationDatabase(root, filename)
    if compilation_db_flags:
        final_flags = compilation_db_flags
    else:
        final_flags = BASE_FLAGS
        clang_flags = FlagsForClangComplete(root)
        if clang_flags:
            final_flags = final_flags + clang_flags
        include_flags = FlagsForInclude(root)
        if include_flags:
            final_flags = final_flags + include_flags
    return {
        'flags': final_flags,
        'do_cache': True
    }
<commit_msg>Show all errors in large projects<commit_after>import os
import os.path
import fnmatch
import logging
import ycm_core

BASE_FLAGS = [
    '-Wall',
    '-Wextra',
    '-Werror',
    '-Wc++98-compat',
    '-Wno-long-long',
    '-Wno-variadic-macros',
    '-fexceptions',
    '-ferror-limit=1000',
    '-DNDEBUG',
    '-std=c++98',
    '-xc++',
    '-I/usr/include/'
]

SOURCE_EXTENSIONS = [
    '.cpp',
    '.cxx',
    '.cc',
    '.c',
    '.m',
    '.mm'
]

HEADER_EXTENSIONS = [
    '.h',
    '.hxx',
    '.hpp',
    '.hh'
]

def IsHeaderFile(filename):
    extension = os.path.splitext(filename)[1]
    return extension in HEADER_EXTENSIONS

def GetCompilationInfoForFile(database, filename):
    if IsHeaderFile(filename):
        basename = os.path.splitext(filename)[0]
        for extension in SOURCE_EXTENSIONS:
            replacement_file = basename + extension
            if os.path.exists(replacement_file):
                compilation_info = database.GetCompilationInfoForFile(replacement_file)
                if compilation_info.compiler_flags_:
                    return compilation_info
        return None
    return database.GetCompilationInfoForFile(filename)

def FindNearest(path, target):
    candidate = os.path.join(path, target)
    if(os.path.isfile(candidate) or os.path.isdir(candidate)):
        logging.info("Found nearest " + target + " at " + candidate)
        return candidate;
    else:
        parent = os.path.dirname(os.path.abspath(path));
        if(parent == path):
            raise RuntimeError("Could not find " + target);
        return FindNearest(parent, target)

def MakeRelativePathsInFlagsAbsolute(flags, working_directory):
    if not working_directory:
        return list(flags)
    new_flags = []
    make_next_absolute = False
    path_flags = [ '-isystem', '-I', '-iquote', '--sysroot=' ]
    for flag in flags:
        new_flag = flag

        if make_next_absolute:
            make_next_absolute = False
            if not flag.startswith('/'):
                new_flag = os.path.join(working_directory, flag)

        for path_flag in path_flags:
            if flag == path_flag:
                make_next_absolute = True
                break

            if flag.startswith(path_flag):
                path = flag[ len(path_flag): ]
                new_flag = path_flag + os.path.join(working_directory, path)
                break

        if new_flag:
            new_flags.append(new_flag)
    return new_flags

def FlagsForClangComplete(root):
    try:
        clang_complete_path = FindNearest(root, '.clang_complete')
        clang_complete_flags = open(clang_complete_path, 'r').read().splitlines()
        return clang_complete_flags
    except:
        return None

def FlagsForInclude(root):
    try:
        include_path = FindNearest(root, 'include')
        flags = []
        for dirroot, dirnames, filenames in os.walk(include_path):
            for dir_path in dirnames:
                real_path = os.path.join(dirroot, dir_path)
                flags = flags + ['-I', real_path]
        return flags
    except:
        return None

def FlagsForCompilationDatabase(root, filename):
    try:
        compilation_db_path = FindNearest(root, 'compile_commands.json')
        compilation_db_dir = os.path.dirname(compilation_db_path)
        logging.info("Set compilation database directory to " + compilation_db_dir)
        compilation_db =  ycm_core.CompilationDatabase(compilation_db_dir)
        if not compilation_db:
            logging.info("Compilation database file found but unable to load")
            return None
        compilation_info = GetCompilationInfoForFile(compilation_db, filename)
        if not compilation_info:
            logging.info("No compilation info for " + filename + " in compilation database")
            return None
        return MakeRelativePathsInFlagsAbsolute(
            compilation_info.compiler_flags_,
            compilation_info.compiler_working_dir_)
    except:
        return None

def FlagsForFile(filename):
    root = os.path.realpath(filename);
    compilation_db_flags = FlagsForCompilationDatabase(root, filename)
    if compilation_db_flags:
        final_flags = compilation_db_flags
    else:
        final_flags = BASE_FLAGS
        clang_flags = FlagsForClangComplete(root)
        if clang_flags:
            final_flags = final_flags + clang_flags
        include_flags = FlagsForInclude(root)
        if include_flags:
            final_flags = final_flags + include_flags
    return {
        'flags': final_flags,
        'do_cache': True
    }
<|endoftext|>