<commit_before>#
#

# Copyright (C) 2007, 2011, 2012, 2013 Google Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

"""QA configuration.

"""

import os

from ganeti import constants
from ganeti import utils
from ganeti import serializer
from ganeti import compat

import qa_error

_INSTANCE_CHECK_KEY = "instance-check"
_ENABLED_HV_KEY = "enabled-hypervisors"

#: QA configuration (L{_QaConfig})
_config = None

class _QaInstance(object):
  __slots__ = [
    "name",
    "nicmac",
    "used",
    "disk_template",
    ]

  def __init__(self, name, nicmac):
    """Initializes instances of this class.

    """
    self.name = name
    self.nicmac = nicmac
    self.used = None
    self.disk_template = None

  @classmethod
  def FromDict(cls, data):
    """Creates instance object from JSON dictionary.

    """
    nicmac = []

    macaddr = data.get("nic.mac/0")
    if macaddr:
      nicmac.append(macaddr)

    return cls(name=data["name"], nicmac=nicmac)

  def __getitem__(self, key):
    """Legacy dict-like interface.

    """
    if key == "name":
      return self.name
    else:
      raise KeyError(key)

  def get(self, key, default):
    """Legacy dict-like interface.

    """
    try:
      return self[key]
    except KeyError:
      return default

  def Release(self):
    """Releases instance and makes it available again.

    """
    assert self.used, \
      ("Instance '%s' was never acquired or released more than once" %
       self.name)

    self.used = False
    self.disk_template = None

  def GetNicMacAddr(self, idx, default):
    """Returns MAC address for NIC.

    @type idx: int
    @param idx: NIC index
    @param default: Default value

    """
    if len(self.nicmac) > idx:
      return self.nicmac[idx]
    else:
      return default

class _QaNode(object):
  __slots__ = [
    "primary",
    "secondary",
    "_added",
    "_use_count",
    ]

  def __init__(self, primary, secondary):
    """Initializes instances of this class.

    """
    self.primary = primary
    self.secondary = secondary
    self._added = False
    self._use_count = 0

  @classmethod
  def FromDict(cls, data):
    """Creates node object from JSON dictionary.

    """
    return cls(primary=data["primary"], secondary=data.get("secondary"))

  def __getitem__(self, key):
    """Legacy dict-like interface.

    """
    if key == "primary":
      return self.primary
    elif key == "secondary":
      return self.secondary
    else:
      raise KeyError(key)

  def get(self, key, default):
    """Legacy dict-like interface.

    """
    try:
      return self[key]
    except KeyError:
      return default

  def Use(self):
    """Marks a node as being in use.

    """
    assert self._use_count >= 0

    self._use_count += 1

    return self

  def Release(self):
    """Release a node (opposite of L{Use}).

    """
    assert self.use_count > 0

    self._use_count -= 1

  def MarkAdded(self):
    """Marks node as having been added to a cluster.

    """
    assert not self._added
    self._added = True

  def MarkRemoved(self):
    """Marks node as having been removed from a cluster.

    """
    assert self._added
    self._added = False

  @property
  def added(self):
    """Returns whether a node is part of a cluster.

    """
    return self._added

  @property
  def use_count(self):
    """Returns number of current uses (controlled by L{Use} and L{Release}).

    """
    return self._use_count

_RESOURCE_CONVERTER = {
  "instances": _QaInstance.FromDict,
  "nodes": _QaNode.FromDict,
  }

def _ConvertResources((key, value)):
  """Converts cluster resources in configuration to Python objects.

  """
  fn = _RESOURCE_CONVERTER.get(key, None)
  if fn:
    return (key, map(fn, value))
  else:
    return (key, value)

class _QaConfig(object):
  def __init__(self, data):
    """Initializes instances of this class.

    """
    self._data = data

    #: Cluster-wide run-time value of the exclusive storage flag
    self._exclusive_storage = None

  @classmethod
  def Load(cls, filename):
    """Loads a configuration file and produces a configuration object.

    @type filename: string
    @param filename: Path to configuration file
    @rtype: L{_QaConfig}

    """
    data = serializer.LoadJson(utils.ReadFile(filename))

    result = cls(dict(map(_ConvertResources,
                          data.items()))) # pylint: disable=E1103
    result.Validate()

    return result

  def Validate(self):
    """Validates loaded configuration data.

    """
    if not self.get("nodes"):
      raise qa_error.Error("Need at least one node")

    if not self.get("instances"):
      raise qa_error.Error("Need at least one instance")

    if (self.get("disk") is None or
        self.get("disk-growth") is None or
        len(self.get("disk")) != len(self.get("disk-growth"))):
      raise qa_error.Error("Config options 'disk' and 'disk-growth' must exist"
                           " and have the same number of items")

    check = self.GetInstanceCheckScript()
    if check:
      try:
        os.stat(check)
      except EnvironmentError, err:
        raise qa_error.Error("Can't find instance check script '%s': %s" %
                             (check, err))

    enabled_hv = frozenset(self.GetEnabledHypervisors())
    if not enabled_hv:
      raise qa_error.Error("No hypervisor is enabled")

    difference = enabled_hv - constants.HYPER_TYPES
    if difference:
      raise qa_error.Error("Unknown hypervisor(s) enabled: %s" %
                           utils.CommaJoin(difference))

  def __getitem__(self, name):
    """Returns configuration value.

    @type name: string
    @param name: Name of configuration entry

    """
    return self._data[name]

  def get(self, name, default=None):
    """Returns configuration value.

    @type name: string
    @param name: Name of configuration entry
    @param default: Default value

    """
    return self._data.get(name, default)

  def GetMasterNode(self):
    """Returns the default master node for the cluster.

    """
    return self["nodes"][0]

  def GetInstanceCheckScript(self):
    """Returns path to instance check script or C{None}.

    """
    return self._data.get(_INSTANCE_CHECK_KEY, None)

  def GetEnabledHypervisors(self):
    """Returns list of enabled hypervisors.

    @rtype: list

    """
    try:
      value = self._data[_ENABLED_HV_KEY]
    except KeyError:
      return [constants.DEFAULT_ENABLED_HYPERVISOR]
    else:
      if value is None:
        return []
      elif isinstance(value, basestring):
        # The configuration key ("enabled-hypervisors") implies there can be
        # multiple values. Multiple hypervisors are comma-separated on the
        # command line option to "gnt-cluster init", so we need to handle them
        # equally here.
        return value.split(",")
      else:
        return value

  def GetDefaultHypervisor(self):
    """Returns the default hypervisor to be used.

    """
    return self.GetEnabledHypervisors()[0]

  def SetExclusiveStorage(self, value):
    """Set the expected value of the C{exclusive_storage} flag for the cluster.

    """
    self._exclusive_storage = bool(value)

  def GetExclusiveStorage(self):
    """Get the expected value of the C{exclusive_storage} flag for the cluster.

    """
    value = self._exclusive_storage
    assert value is not None
    return value

  def IsTemplateSupported(self, templ):
    """Is the given disk template supported by the current configuration?

    """
    return (not self.GetExclusiveStorage() or
            templ in constants.DTS_EXCL_STORAGE)

def Load(path):
  """Loads the passed configuration file.

  """
  global _config # pylint: disable=W0603

  _config = _QaConfig.Load(path)

def GetConfig():
  """Returns the configuration object.

  """
  if _config is None:
    raise RuntimeError("Configuration not yet loaded")

  return _config

def get(name, default=None):
  """Wrapper for L{_QaConfig.get}.

  """
  return GetConfig().get(name, default=default)

class Either:
  def __init__(self, tests):
    """Initializes this class.

    @type tests: list or string
    @param tests: List of test names
    @see: L{TestEnabled} for details

    """
    self.tests = tests

def _MakeSequence(value):
  """Make sequence of single argument.

  If the single argument is not already a list or tuple, a list with the
  argument as a single item is returned.

  """
  if isinstance(value, (list, tuple)):
    return value
  else:
    return [value]

def _TestEnabledInner(check_fn, names, fn):
  """Evaluate test conditions.

  @type check_fn: callable
  @param check_fn: Callback to check whether a test is enabled
  @type names: sequence or string
  @param names: Test name(s)
  @type fn: callable
  @param fn: Aggregation function
  @rtype: bool
  @return: Whether test is enabled

  """
  names = _MakeSequence(names)

  result = []

  for name in names:
    if isinstance(name, Either):
      value = _TestEnabledInner(check_fn, name.tests, compat.any)
    elif isinstance(name, (list, tuple)):
      value = _TestEnabledInner(check_fn, name, compat.all)
    else:
      value = check_fn(name)

    result.append(value)

  return fn(result)

def TestEnabled(tests, _cfg=None):
  """Returns True if the given tests are enabled.

  @param tests: A single test as a string, or a list of tests to check; can
    contain L{Either} for OR conditions, AND is default

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  # Get settings for all tests
  cfg_tests = cfg.get("tests", {})

  # Get default setting
  default = cfg_tests.get("default", True)

  return _TestEnabledInner(lambda name: cfg_tests.get(name, default),
                           tests, compat.all)

def GetInstanceCheckScript(*args):
  """Wrapper for L{_QaConfig.GetInstanceCheckScript}.

  """
  return GetConfig().GetInstanceCheckScript(*args)

def GetEnabledHypervisors(*args):
  """Wrapper for L{_QaConfig.GetEnabledHypervisors}.

  """
  return GetConfig().GetEnabledHypervisors(*args)

def GetDefaultHypervisor(*args):
  """Wrapper for L{_QaConfig.GetDefaultHypervisor}.

  """
  return GetConfig().GetDefaultHypervisor(*args)

def GetInstanceNicMac(inst, default=None):
  """Returns MAC address for instance's network interface.

  """
  return inst.GetNicMacAddr(0, default)

def GetMasterNode():
  """Wrapper for L{_QaConfig.GetMasterNode}.

  """
  return GetConfig().GetMasterNode()

def AcquireInstance(_cfg=None):
  """Returns an instance which isn't in use.

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  # Filter out unwanted instances
  instances = filter(lambda inst: not inst.used, cfg["instances"])

  if not instances:
    raise qa_error.OutOfInstancesError("No instances left")

  inst = instances[0]

  assert not inst.used
  assert inst.disk_template is None

  inst.used = True

  return inst

def GetInstanceTemplate(inst):
  """Return the disk template of an instance.

  """
  templ = inst.disk_template
  assert templ is not None
  return templ

def SetInstanceTemplate(inst, template):
  """Set the disk template for an instance.

  """
  inst.disk_template = template

def SetExclusiveStorage(value):
  """Wrapper for L{_QaConfig.SetExclusiveStorage}.

  """
  return GetConfig().SetExclusiveStorage(value)

def GetExclusiveStorage():
  """Wrapper for L{_QaConfig.GetExclusiveStorage}.

  """
  return GetConfig().GetExclusiveStorage()

def IsTemplateSupported(templ):
  """Wrapper for L{_QaConfig.GetExclusiveStorage}.

  """
  return GetConfig().IsTemplateSupported(templ)

def AcquireNode(exclude=None, _cfg=None):
  """Returns the least used node.

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  master = cfg.GetMasterNode()

  # Filter out unwanted nodes
  # TODO: Maybe combine filters
  if exclude is None:
    nodes = cfg["nodes"][:]
  elif isinstance(exclude, (list, tuple)):
    nodes = filter(lambda node: node not in exclude, cfg["nodes"])
  else:
    nodes = filter(lambda node: node != exclude, cfg["nodes"])

  nodes = filter(lambda node: node.added or node == master, nodes)

  if not nodes:
    raise qa_error.OutOfNodesError("No nodes left")

  # Get node with least number of uses
  # TODO: Switch to computing sort key instead of comparing directly
  def compare(a, b):
    result = cmp(a.use_count, b.use_count)
    if result == 0:
      result = cmp(a.primary, b.primary)
    return result

  nodes.sort(cmp=compare)

  return nodes[0].Use()

def AcquireManyNodes(num, exclude=None):
  """Return the least used nodes.

  @type num: int
  @param num: Number of nodes; can be 0.
  @type exclude: list of nodes or C{None}
  @param exclude: nodes to be excluded from the choice
  @rtype: list of nodes
  @return: C{num} different nodes

  """
  nodes = []
  if exclude is None:
    exclude = []
  elif isinstance(exclude, (list, tuple)):
    # Don't modify the incoming argument
    exclude = list(exclude)
  else:
    exclude = [exclude]

  try:
    for _ in range(0, num):
      n = AcquireNode(exclude=exclude)
      nodes.append(n)
      exclude.append(n)
  except qa_error.OutOfNodesError:
    ReleaseManyNodes(nodes)
    raise
  return nodes

def ReleaseManyNodes(nodes):
  for node in nodes:
    node.Release()
<commit_msg>QA: Remove dictionary-like interface from config objects<commit_after>#
#

# Copyright (C) 2007, 2011, 2012, 2013 Google Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

"""QA configuration.

"""

import os

from ganeti import constants
from ganeti import utils
from ganeti import serializer
from ganeti import compat

import qa_error

_INSTANCE_CHECK_KEY = "instance-check"
_ENABLED_HV_KEY = "enabled-hypervisors"

#: QA configuration (L{_QaConfig})
_config = None

class _QaInstance(object):
  __slots__ = [
    "name",
    "nicmac",
    "used",
    "disk_template",
    ]

  def __init__(self, name, nicmac):
    """Initializes instances of this class.

    """
    self.name = name
    self.nicmac = nicmac
    self.used = None
    self.disk_template = None

  @classmethod
  def FromDict(cls, data):
    """Creates instance object from JSON dictionary.

    """
    nicmac = []

    macaddr = data.get("nic.mac/0")
    if macaddr:
      nicmac.append(macaddr)

    return cls(name=data["name"], nicmac=nicmac)

  def Release(self):
    """Releases instance and makes it available again.

    """
    assert self.used, \
      ("Instance '%s' was never acquired or released more than once" %
       self.name)

    self.used = False
    self.disk_template = None

  def GetNicMacAddr(self, idx, default):
    """Returns MAC address for NIC.

    @type idx: int
    @param idx: NIC index
    @param default: Default value

    """
    if len(self.nicmac) > idx:
      return self.nicmac[idx]
    else:
      return default

class _QaNode(object):
  __slots__ = [
    "primary",
    "secondary",
    "_added",
    "_use_count",
    ]

  def __init__(self, primary, secondary):
    """Initializes instances of this class.

    """
    self.primary = primary
    self.secondary = secondary
    self._added = False
    self._use_count = 0

  @classmethod
  def FromDict(cls, data):
    """Creates node object from JSON dictionary.

    """
    return cls(primary=data["primary"], secondary=data.get("secondary"))

  def Use(self):
    """Marks a node as being in use.

    """
    assert self._use_count >= 0

    self._use_count += 1

    return self

  def Release(self):
    """Release a node (opposite of L{Use}).

    """
    assert self.use_count > 0

    self._use_count -= 1

  def MarkAdded(self):
    """Marks node as having been added to a cluster.

    """
    assert not self._added
    self._added = True

  def MarkRemoved(self):
    """Marks node as having been removed from a cluster.

    """
    assert self._added
    self._added = False

  @property
  def added(self):
    """Returns whether a node is part of a cluster.

    """
    return self._added

  @property
  def use_count(self):
    """Returns number of current uses (controlled by L{Use} and L{Release}).

    """
    return self._use_count

_RESOURCE_CONVERTER = {
  "instances": _QaInstance.FromDict,
  "nodes": _QaNode.FromDict,
  }

def _ConvertResources((key, value)):
  """Converts cluster resources in configuration to Python objects.

  """
  fn = _RESOURCE_CONVERTER.get(key, None)
  if fn:
    return (key, map(fn, value))
  else:
    return (key, value)

class _QaConfig(object):
  def __init__(self, data):
    """Initializes instances of this class.

    """
    self._data = data

    #: Cluster-wide run-time value of the exclusive storage flag
    self._exclusive_storage = None

  @classmethod
  def Load(cls, filename):
    """Loads a configuration file and produces a configuration object.

    @type filename: string
    @param filename: Path to configuration file
    @rtype: L{_QaConfig}

    """
    data = serializer.LoadJson(utils.ReadFile(filename))

    result = cls(dict(map(_ConvertResources,
                          data.items()))) # pylint: disable=E1103
    result.Validate()

    return result

  def Validate(self):
    """Validates loaded configuration data.

    """
    if not self.get("nodes"):
      raise qa_error.Error("Need at least one node")

    if not self.get("instances"):
      raise qa_error.Error("Need at least one instance")

    if (self.get("disk") is None or
        self.get("disk-growth") is None or
        len(self.get("disk")) != len(self.get("disk-growth"))):
      raise qa_error.Error("Config options 'disk' and 'disk-growth' must exist"
                           " and have the same number of items")

    check = self.GetInstanceCheckScript()
    if check:
      try:
        os.stat(check)
      except EnvironmentError, err:
        raise qa_error.Error("Can't find instance check script '%s': %s" %
                             (check, err))

    enabled_hv = frozenset(self.GetEnabledHypervisors())
    if not enabled_hv:
      raise qa_error.Error("No hypervisor is enabled")

    difference = enabled_hv - constants.HYPER_TYPES
    if difference:
      raise qa_error.Error("Unknown hypervisor(s) enabled: %s" %
                           utils.CommaJoin(difference))

  def __getitem__(self, name):
    """Returns configuration value.

    @type name: string
    @param name: Name of configuration entry

    """
    return self._data[name]

  def get(self, name, default=None):
    """Returns configuration value.

    @type name: string
    @param name: Name of configuration entry
    @param default: Default value

    """
    return self._data.get(name, default)

  def GetMasterNode(self):
    """Returns the default master node for the cluster.

    """
    return self["nodes"][0]

  def GetInstanceCheckScript(self):
    """Returns path to instance check script or C{None}.

    """
    return self._data.get(_INSTANCE_CHECK_KEY, None)

  def GetEnabledHypervisors(self):
    """Returns list of enabled hypervisors.

    @rtype: list

    """
    try:
      value = self._data[_ENABLED_HV_KEY]
    except KeyError:
      return [constants.DEFAULT_ENABLED_HYPERVISOR]
    else:
      if value is None:
        return []
      elif isinstance(value, basestring):
        # The configuration key ("enabled-hypervisors") implies there can be
        # multiple values. Multiple hypervisors are comma-separated on the
        # command line option to "gnt-cluster init", so we need to handle them
        # equally here.
        return value.split(",")
      else:
        return value

  def GetDefaultHypervisor(self):
    """Returns the default hypervisor to be used.

    """
    return self.GetEnabledHypervisors()[0]

  def SetExclusiveStorage(self, value):
    """Set the expected value of the C{exclusive_storage} flag for the cluster.

    """
    self._exclusive_storage = bool(value)

  def GetExclusiveStorage(self):
    """Get the expected value of the C{exclusive_storage} flag for the cluster.

    """
    value = self._exclusive_storage
    assert value is not None
    return value

  def IsTemplateSupported(self, templ):
    """Is the given disk template supported by the current configuration?

    """
    return (not self.GetExclusiveStorage() or
            templ in constants.DTS_EXCL_STORAGE)

def Load(path):
  """Loads the passed configuration file.

  """
  global _config # pylint: disable=W0603

  _config = _QaConfig.Load(path)

def GetConfig():
  """Returns the configuration object.

  """
  if _config is None:
    raise RuntimeError("Configuration not yet loaded")

  return _config

def get(name, default=None):
  """Wrapper for L{_QaConfig.get}.

  """
  return GetConfig().get(name, default=default)

class Either:
  def __init__(self, tests):
    """Initializes this class.

    @type tests: list or string
    @param tests: List of test names
    @see: L{TestEnabled} for details

    """
    self.tests = tests

def _MakeSequence(value):
  """Make sequence of single argument.

  If the single argument is not already a list or tuple, a list with the
  argument as a single item is returned.

  """
  if isinstance(value, (list, tuple)):
    return value
  else:
    return [value]

def _TestEnabledInner(check_fn, names, fn):
  """Evaluate test conditions.

  @type check_fn: callable
  @param check_fn: Callback to check whether a test is enabled
  @type names: sequence or string
  @param names: Test name(s)
  @type fn: callable
  @param fn: Aggregation function
  @rtype: bool
  @return: Whether test is enabled

  """
  names = _MakeSequence(names)

  result = []

  for name in names:
    if isinstance(name, Either):
      value = _TestEnabledInner(check_fn, name.tests, compat.any)
    elif isinstance(name, (list, tuple)):
      value = _TestEnabledInner(check_fn, name, compat.all)
    else:
      value = check_fn(name)

    result.append(value)

  return fn(result)

def TestEnabled(tests, _cfg=None):
  """Returns True if the given tests are enabled.

  @param tests: A single test as a string, or a list of tests to check; can
    contain L{Either} for OR conditions, AND is default

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  # Get settings for all tests
  cfg_tests = cfg.get("tests", {})

  # Get default setting
  default = cfg_tests.get("default", True)

  return _TestEnabledInner(lambda name: cfg_tests.get(name, default),
                           tests, compat.all)

def GetInstanceCheckScript(*args):
  """Wrapper for L{_QaConfig.GetInstanceCheckScript}.

  """
  return GetConfig().GetInstanceCheckScript(*args)

def GetEnabledHypervisors(*args):
  """Wrapper for L{_QaConfig.GetEnabledHypervisors}.

  """
  return GetConfig().GetEnabledHypervisors(*args)

def GetDefaultHypervisor(*args):
  """Wrapper for L{_QaConfig.GetDefaultHypervisor}.

  """
  return GetConfig().GetDefaultHypervisor(*args)

def GetInstanceNicMac(inst, default=None):
  """Returns MAC address for instance's network interface.

  """
  return inst.GetNicMacAddr(0, default)

def GetMasterNode():
  """Wrapper for L{_QaConfig.GetMasterNode}.

  """
  return GetConfig().GetMasterNode()

def AcquireInstance(_cfg=None):
  """Returns an instance which isn't in use.

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  # Filter out unwanted instances
  instances = filter(lambda inst: not inst.used, cfg["instances"])

  if not instances:
    raise qa_error.OutOfInstancesError("No instances left")

  inst = instances[0]

  assert not inst.used
  assert inst.disk_template is None

  inst.used = True

  return inst

def GetInstanceTemplate(inst):
  """Return the disk template of an instance.

  """
  templ = inst.disk_template
  assert templ is not None
  return templ

def SetInstanceTemplate(inst, template):
  """Set the disk template for an instance.

  """
  inst.disk_template = template

def SetExclusiveStorage(value):
  """Wrapper for L{_QaConfig.SetExclusiveStorage}.

  """
  return GetConfig().SetExclusiveStorage(value)

def GetExclusiveStorage():
  """Wrapper for L{_QaConfig.GetExclusiveStorage}.

  """
  return GetConfig().GetExclusiveStorage()

def IsTemplateSupported(templ):
  """Wrapper for L{_QaConfig.GetExclusiveStorage}.

  """
  return GetConfig().IsTemplateSupported(templ)

def AcquireNode(exclude=None, _cfg=None):
  """Returns the least used node.

  """
  if _cfg is None:
    cfg = GetConfig()
  else:
    cfg = _cfg

  master = cfg.GetMasterNode()

  # Filter out unwanted nodes
  # TODO: Maybe combine filters
  if exclude is None:
    nodes = cfg["nodes"][:]
  elif isinstance(exclude, (list, tuple)):
    nodes = filter(lambda node: node not in exclude, cfg["nodes"])
  else:
    nodes = filter(lambda node: node != exclude, cfg["nodes"])

  nodes = filter(lambda node: node.added or node == master, nodes)

  if not nodes:
    raise qa_error.OutOfNodesError("No nodes left")

  # Get node with least number of uses
  # TODO: Switch to computing sort key instead of comparing directly
  def compare(a, b):
    result = cmp(a.use_count, b.use_count)
    if result == 0:
      result = cmp(a.primary, b.primary)
    return result

  nodes.sort(cmp=compare)

  return nodes[0].Use()

def AcquireManyNodes(num, exclude=None):
  """Return the least used nodes.

  @type num: int
  @param num: Number of nodes; can be 0.
  @type exclude: list of nodes or C{None}
  @param exclude: nodes to be excluded from the choice
  @rtype: list of nodes
  @return: C{num} different nodes

  """
  nodes = []
  if exclude is None:
    exclude = []
  elif isinstance(exclude, (list, tuple)):
    # Don't modify the incoming argument
    exclude = list(exclude)
  else:
    exclude = [exclude]

  try:
    for _ in range(0, num):
      n = AcquireNode(exclude=exclude)
      nodes.append(n)
      exclude.append(n)
  except qa_error.OutOfNodesError:
    ReleaseManyNodes(nodes)
    raise
  return nodes

def ReleaseManyNodes(nodes):
  for node in nodes:
    node.Release()
<|endoftext|>


<commit_before>#coding=utf-8
from flask import Flask, jsonify, abort, g, make_response, request , url_for
from flask.ext.restful import Api, Resource
from flask.ext.httpauth import HTTPBasicAuth
from sqlalchemy.orm import relationship, backref
from sqlalchemy import ForeignKey
from app import app,db,api,auth
import json
import datetime
from simple_result import SimpleResult
from user import User

class Comment(db.Model):
    __tablename__ = 't_comment'
    id = db.Column(db.Integer, primary_key=True)
    content = db.Column(db.Text)
    createdAt = db.Column(db.DateTime)
    updatedAt = db.Column(db.DateTime)
    deletedAt = db.Column(db.DateTime)
    topicId = db.Column(db.Integer, ForeignKey('t_topic.id'))
    userId = db.Column(db.Integer, ForeignKey('t_user.id'))
    parentId = db.Column(db.Integer, ForeignKey('t_comment.id'))
    rootId = db.Column(db.Integer, ForeignKey('t_comment.id'))

    def replyCount(self):
        return 0

    def author(self):
        author = User.query.get(self.userId)
        return author

    def json(self):
        author = self.author()
        return {
            "id":self.id,
            "content":self.content,
            "replyCount":self.replyCount(),
            "parentId":self.parentId,
            "rootId":self.rootId,
            "replyTime":self.createdAt.strftime('%Y-%m-%d %H:%M:%S'),
            "authorName":author.nickname,
            "authorIcon":author.icon,
            "authorId":self.userId
        }

@app.route('/api/user/comment/postComment', methods=['POST'])
@auth.login_required
def postComment():
    content = request.json.get("content")
    replyId = request.json.get("replyId")
    topicId = request.json.get("topicId")
    user = g.user
    if not topicId or not content:
        abort(400)
    comment = Comment()
    comment.content = content
    comment.parentId = replyId
    comment.topicId = topicId
    comment.createdAt = datetime.datetime.now()
    comment.updatedAt = comment.createdAt
    comment.rootId = findCommentRootId(replyId)
    comment.userId = user.id
    db.session.add(comment)
    db.session.commit()
    return (jsonify(SimpleResult(0,"发布成功").json()),200)

@app.route('/api/web/comment/getCommentList', methods=['GET'])
def getCommentList():
    page = request.args.get("page")
    limit = request.args.get("limit")
    comment_id = request.args.get("commentId")
    topic_id = request.args.get("topicId")
    if not page or not topic_id or not limit:
        abort(400)
    comments = []
    page = int(page)
    limit = int(limit)
    if comment_id:
        comments = db.session.query(Comment).filter(Comment.rootId == comment_id,Comment.topicId == topic_id).limit(limit).offset(limit * (page-1))
    else:
        comments = db.session.query(Comment).filter(Comment.topicId == topic_id,Comment.rootId == None).limit(limit).offset(limit * (page-1))
    commentsJsonList = []
    for comment in comments:
        commentsJsonList.append(comment.json())
    return (json.dumps(commentsJsonList),200)

def findCommentRootId(replyId):
    if (replyId):
        #寻找该评论
        comment = Comment.query.get(replyId)
        # 如果该评论有父节点，寻找评论父节点
        if comment.parentId:
            return findCommentRootId(comment.parentId)
        else:
            return replyId
    return None

@app.errorhandler(400)
def bad_request(error):
    return make_response(jsonify(SimpleResult(-1,"缺少参数").json()), 400)<commit_msg>comment的 replyCount逻辑补上<commit_after>#coding=utf-8
from flask import Flask, jsonify, abort, g, make_response, request , url_for
from flask.ext.restful import Api, Resource
from flask.ext.httpauth import HTTPBasicAuth
from sqlalchemy.orm import relationship, backref
from sqlalchemy import ForeignKey
from app import app,db,api,auth
import json
import datetime
from simple_result import SimpleResult
from user import User

class Comment(db.Model):
    __tablename__ = 't_comment'
    id = db.Column(db.Integer, primary_key=True)
    content = db.Column(db.Text)
    createdAt = db.Column(db.DateTime)
    updatedAt = db.Column(db.DateTime)
    deletedAt = db.Column(db.DateTime)
    topicId = db.Column(db.Integer, ForeignKey('t_topic.id'))
    userId = db.Column(db.Integer, ForeignKey('t_user.id'))
    parentId = db.Column(db.Integer, ForeignKey('t_comment.id'))
    rootId = db.Column(db.Integer, ForeignKey('t_comment.id'))

    def replyCount(self):
        count = Comment.query.filter(Comment.rootId == self.id).count()
        return count

    def author(self):
        author = User.query.get(self.userId)
        return author

    def json(self):
        author = self.author()
        return {
            "id":self.id,
            "content":self.content,
            "replyCount":self.replyCount(),
            "parentId":self.parentId,
            "rootId":self.rootId,
            "replyTime":self.createdAt.strftime('%Y-%m-%d %H:%M:%S'),
            "authorName":author.nickname,
            "authorIcon":author.icon,
            "authorId":self.userId
        }

@app.route('/api/user/comment/postComment', methods=['POST'])
@auth.login_required
def postComment():
    content = request.json.get("content")
    replyId = request.json.get("replyId")
    topicId = request.json.get("topicId")
    user = g.user
    if not topicId or not content:
        abort(400)
    comment = Comment()
    comment.content = content
    comment.parentId = replyId
    comment.topicId = topicId
    comment.createdAt = datetime.datetime.now()
    comment.updatedAt = comment.createdAt
    comment.rootId = findCommentRootId(replyId)
    comment.userId = user.id
    db.session.add(comment)
    db.session.commit()
    return (jsonify(SimpleResult(0,"发布成功").json()),200)

@app.route('/api/web/comment/getCommentList', methods=['GET'])
def getCommentList():
    page = request.args.get("page")
    limit = request.args.get("limit")
    comment_id = request.args.get("commentId")
    topic_id = request.args.get("topicId")
    if not page or not topic_id or not limit:
        abort(400)
    comments = []
    page = int(page)
    limit = int(limit)
    if comment_id:
        comments = db.session.query(Comment).filter(Comment.rootId == comment_id,Comment.topicId == topic_id).limit(limit).offset(limit * (page-1))
    else:
        comments = db.session.query(Comment).filter(Comment.topicId == topic_id,Comment.rootId == None).limit(limit).offset(limit * (page-1))
    commentsJsonList = []
    for comment in comments:
        commentsJsonList.append(comment.json())
    return (json.dumps(commentsJsonList),200)

def findCommentRootId(replyId):
    if (replyId):
        #寻找该评论
        comment = Comment.query.get(replyId)
        # 如果该评论有父节点，寻找评论父节点
        if comment.parentId:
            return findCommentRootId(comment.parentId)
        else:
            return replyId
    return None

@app.errorhandler(400)
def bad_request(error):
    return make_response(jsonify(SimpleResult(-1,"缺少参数").json()), 400)<|endoftext|>


<commit_before>import collections
import operator

from .utils import _extract_match
from difflib import SequenceMatcher

class Tree:

    def __init__(self, *args, **kwargs):
        self.instance = {}

        # Sequence end indicator
        self._END = '\\*__ends__*//'

        # Original value indicator
        self._ORIG = '\\*__orig__*//'

    def _list_to_nested_dict(self, lis):
        """
        Turns a list into a nested dict

        E.g.:
            ['a','b','c'] => { 'c' : { 'b' : { 'a' : {} } } }

        extra data:

            _end_ - Marks the end of the list
            E.g.:
                ['a','b'] => { 'b' : { 'a' : {}, '_end_': True}, '_end_': False}

            _orig_ - The original value of the key/list item
            E.g.:
                ['A'] => { 'a' : {}, '_orig_': 'A', '_end_': True}
        """
        d = {}
        for i in range(0, len(lis)):
            d[self._END] = True if i == 0 else False
            d[self._ORIG] = ['/'.join(lis[:i+1])]
            d = {lis[i].lower(): d}
        return d

    def _get_best_match(self, path, possibilities):
        """
        Given a path find how similar it is to all paths in possibilities

        :str: path - A path part E.g.: a/b.py => a
        :list: possibilities - Collected possibilities
        """

        # Map out similarity of possible paths with the path being looked up
        similarity = list(map(lambda x: SequenceMatcher(None, path, x).ratio(), possibilities))
        # Get the index, value of the most similar path
        index, value = max(enumerate(similarity), key=operator.itemgetter(1))

        return possibilities[index]

    def _recursive_lookup(self, d, lis, results, i = 0, end=False):
        """
        Performs a lookup in tree recursively

        :dict: d - tree branch
        :list: lis - list of strings to search for
        :list: results - Collected hit results
        :int: i - Index of lis
        :bool: end - Indicates if last lookup was the end of a sequence

        :returns a list of hit results if path is found in the tree
        """
        key = None

        if i < len(lis):
            key = lis[i].lower()

        if i == 0 and d.get(key):
            results = d.get(key).get(self._ORIG)

        if d.get(key):
            root = d.get(key)
            return self._recursive_lookup(
                root,
                lis,
                results,
                i + 1,
                root.get(self._END)
            )
        else:
            if not end:
                results = []
            return results

    def lookup(self, path):
        """
        Lookup a path in the tree

        :str: path - The path to search for

        :returns The closest matching path in the tree if present else None
        """
        path_hit = None
        path_split = list(reversed(path.split('/')))
        results = self._recursive_lookup(self.instance, path_split, [])

        if not results:
            return None

        if len(results) == 1:
            path_hit = results[0]
        else:
            path_hit = self._get_best_match(path, list(reversed(results)))

        return path_hit

    def _update(self, d, u):
        """
        Update a dictionary
        :dict: d - Dictionary being updated
        :dict: u - Dictionary being merged
        """
        for k, v in u.items():
            if isinstance(v, collections.Mapping):
                r = self._update(d.get(k, {}), v)
                d[k] = r
            else:
                if k == self._END  and d.get(k) == True:
                    pass
                elif k == self._ORIG and d.get(k) and u.get(k):
                    d[k] = d[k] + u[k]
                else:
                    d[k] = u[k]
        return d

    def insert(self, path):
        """
        Insert a path into the tree

        :str: path - The path to insert
        """

        path_split = path.split('/')
        root_key =  path_split[-1].lower()
        root = self.instance.get(root_key)

        if not root:
            u = self._list_to_nested_dict(path_split)
            self.instance.update(u)
        else:
            u = self._list_to_nested_dict(path_split)
            self.instance = self._update(self.instance, u)

    def construct_tree(self, toc):
        """
        Constructs a tree

        :str: toc - The table of contents
        """
        constructing = True
        toc_index    = 1

        while constructing:
            if toc_index < len(toc) - 1:
                path = _extract_match(toc, toc_index)
                if path:
                    self.insert(path)
                    toc_index = toc_index  + len(path) + 1
                else:
                    if toc[toc_index] == ',':
                        toc_index += 1
            else:
                constructing = False
                break
<commit_msg>change order of collected original paths<commit_after>import collections
import operator

from .utils import _extract_match
from difflib import SequenceMatcher

class Tree:

    def __init__(self, *args, **kwargs):
        self.instance = {}

        # Sequence end indicator
        self._END = '\\*__ends__*//'

        # Original value indicator
        self._ORIG = '\\*__orig__*//'

    def _list_to_nested_dict(self, lis):
        """
        Turns a list into a nested dict

        E.g.:
            ['a','b','c'] => { 'c' : { 'b' : { 'a' : {} } } }

        extra data:

            _end_ - Marks the end of the list
            E.g.:
                ['a','b'] => { 'b' : { 'a' : {}, '_end_': True}, '_end_': False}

            _orig_ - The original value of the key/list item
            E.g.:
                ['A'] => { 'a' : {}, '_orig_': 'A', '_end_': True}
        """
        d = {}
        for i in range(0, len(lis)):
            d[self._END] = True if i == 0 else False
            d[self._ORIG] = ['/'.join(lis[i:])]
            d = {lis[i].lower(): d}
        return d

    def _get_best_match(self, path, possibilities):
        """
        Given a path find how similar it is to all paths in possibilities

        :str: path - A path part E.g.: a/b.py => a
        :list: possibilities - Collected possibilities
        """

        # Map out similarity of possible paths with the path being looked up
        similarity = list(map(lambda x: SequenceMatcher(None, path, x).ratio(), possibilities))

        # Get the index, value of the most similar path
        index, value = max(enumerate(similarity), key=operator.itemgetter(1))

        return possibilities[index]

    def _recursive_lookup(self, d, lis, results, i = 0, end=False):
        """
        Performs a lookup in tree recursively

        :dict: d - tree branch
        :list: lis - list of strings to search for
        :list: results - Collected hit results
        :int: i - Index of lis
        :bool: end - Indicates if last lookup was the end of a sequence

        :returns a list of hit results if path is found in the tree
        """
        key = None

        if i < len(lis):
            key = lis[i].lower()

        if d.get(key):
            root = d.get(key)
            results = d.get(key).get(self._ORIG)
            return self._recursive_lookup(
                root,
                lis,
                results,
                i + 1,
                root.get(self._END)
            )
        else:
            if not end:
                results = []
            return results

    def lookup(self, path):
        """
        Lookup a path in the tree

        :str: path - The path to search for

        :returns The closest matching path in the tree if present else None
        """
        path_hit = None
        path_split = list(reversed(path.split('/')))
        results = self._recursive_lookup(self.instance, path_split, [])

        if not results:
            return None

        if len(results) == 1:
            path_hit = results[0]
        else:
            path_hit = self._get_best_match(path, list(reversed(results)))

        return path_hit

    def _update(self, d, u):
        """
        Update a dictionary
        :dict: d - Dictionary being updated
        :dict: u - Dictionary being merged
        """
        for k, v in u.items():
            if isinstance(v, collections.Mapping):
                r = self._update(d.get(k, {}), v)
                d[k] = r
            else:
                if k == self._END  and d.get(k) == True:
                    pass
                elif k == self._ORIG and d.get(k) and u.get(k):
                    if d[k] != u[k]:
                        d[k] = d[k] + u[k]
                else:
                    d[k] = u[k]
        return d

    def insert(self, path):
        """
        Insert a path into the tree

        :str: path - The path to insert
        """

        path_split = path.split('/')
        root_key =  path_split[-1].lower()
        root = self.instance.get(root_key)

        if not root:
            u = self._list_to_nested_dict(path_split)
            self.instance.update(u)
        else:
            u = self._list_to_nested_dict(path_split)
            self.instance = self._update(self.instance, u)

    def construct_tree(self, toc):
        """
        Constructs a tree

        :str: toc - The table of contents
        """
        constructing = True
        toc_index    = 1

        while constructing:
            if toc_index < len(toc) - 1:
                path = _extract_match(toc, toc_index)
                if path:
                    self.insert(path)
                    toc_index = toc_index  + len(path) + 1
                else:
                    if toc[toc_index] == ',':
                        toc_index += 1
            else:
                constructing = False
                break
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# ###
# Copyright (c) 2013-2016, Rice University
# This software is subject to the provisions of the GNU Affero General
# Public License version 3 (AGPLv3).
# See LICENCE.txt for details.
# ###
import unittest
from datetime import datetime

from cnxdb.init import init_db
from pyramid import testing

from .. import use_cases
from ..testing import (
    integration_test_settings,
    db_connection_factory,
    )

# FIXME There is an issue with setting up the celery app more than once.
#       Apparently, creating the app a second time doesn't really create
#       it again. There is some global state hanging around that we can't
#       easily get at. This causes the task results tables used in these
#       views to not exist, because the code believes it's already been
#       initialized.
@unittest.skip("celery is too global")
class PostPublicationsViewsTestCase(unittest.TestCase):
    maxDiff = None

    @classmethod
    def setUpClass(cls):
        cls.settings = integration_test_settings()
        from cnxpublishing.config import CONNECTION_STRING
        cls.db_conn_str = cls.settings[CONNECTION_STRING]
        cls.db_connect = staticmethod(db_connection_factory())

    def setUp(self):
        self.config = testing.setUp(settings=self.settings)
        self.config.include('cnxpublishing.tasks')
        init_db(self.db_conn_str, True)

    def tearDown(self):
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("DROP SCHEMA public CASCADE")
                cursor.execute("CREATE SCHEMA public")
        testing.tearDown()

    @property
    def target(self):
        from ...views.admin import admin_post_publications
        return admin_post_publications

    def test_no_results(self):
        request = testing.DummyRequest()

        resp_data = self.target(request)

        self.assertEqual({'states': []}, resp_data)

    def test(self):
        request = testing.DummyRequest()

        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                # Insert one book into archive.
                book = use_cases.setup_BOOK_in_archive(self, cursor)
                db_conn.commit()

                # Insert some data into the association table.
                cursor.execute("""
INSERT INTO document_baking_result_associations
  (result_id, module_ident)
SELECT
  uuid_generate_v4(),
  (SELECT module_ident FROM modules ORDER BY module_ident DESC LIMIT 1);""")
                db_conn.commit()

                cursor.execute("""\
INSERT INTO document_baking_result_associations
  (result_id, module_ident)
SELECT
  uuid_generate_v4(),
  (SELECT module_ident FROM modules ORDER BY module_ident DESC LIMIT 1);""")

        resp_data = self.target(request)
        self.assertEqual({
            'states': [
                {'created': resp_data['states'][0]['created'],
                 'ident_hash': book.ident_hash,
                 'state': u'PENDING',
                 'state_message': '',
                 'title': 'Book of Infinity'},
                {'created': resp_data['states'][1]['created'],
                 'ident_hash': book.ident_hash,
                 'state': u'PENDING',
                 'state_message': '',
                 'title': 'Book of Infinity'},
                ]}, resp_data)

class ErrorBannerViewsTestCase(unittest.TestCase):
    maxDiff = None

    @classmethod
    def setUpClass(cls):
        cls.settings = integration_test_settings()
        from cnxpublishing.config import CONNECTION_STRING
        cls.db_conn_str = cls.settings[CONNECTION_STRING]
        cls.db_connect = staticmethod(db_connection_factory())

    def setUp(self):
        self.config = testing.setUp(settings=self.settings)
        self.config.include('cnxpublishing.tasks')
        init_db(self.db_conn_str, True)

    def tearDown(self):
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("DROP SCHEMA public CASCADE")
                cursor.execute("CREATE SCHEMA public")
        testing.tearDown()

    @unittest.skip("celery is too global")
    def test_error_banner_get(self):
        request = testing.DummyRequest()

        from ...views.admin import admin_post_error_banner
        defaults = admin_post_error_banner(request)
        self.assertEqual(set(defaults.keys()),
                         set(['start_date', 'start_time',
                              'end_date', 'end_time']))

    @unittest.skip("celery is too global")
    def test_error_banner_post(self):
        request = testing.DummyRequest()
        request.POST = {
            'message': 'test message',
            'priority': 1,
            'type': 1,
            'start_date': '2017-01-01',
            'start_time': '00:01',
            'end_date': '2017-01-02',
            'end_time': '00:02',
        }
        from ...views.admin import admin_post_error_banner_POST
        results = admin_post_error_banner_POST(request)

        # assert the error message has been added to the table
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("""SELECT message, priority, starts, ends
                                  from service_state_messages
                                  WHERE message='test message';""")
                result = cursor.fetchone()
                self.assertEqual('test message', result[0])
                self.assertEqual(1, result[1])
                self.assertTrue(datetime.strftime(
                    result[2], "%Y-%m-%d %H:%M") == '2017-01-01 00:01')
                self.assertTrue(datetime.strftime(
                    result[3], "%Y-%m-%d %H:%M") == '2017-01-02 00:02')
                cursor.execute("""DELETE from service_state_messages
                                  WHERE message='test message';""")

        # Assert the correct variables were passed to the template
        self.assertEqual('test message', results['message'])
        self.assertEqual(1, results['priority'])
        self.assertEqual(1, results['type'])
        self.assertTrue(datetime.strftime(
            results['starts'], "%Y-%m-%d %H:%M") == '2017-01-01 00:01')
        self.assertTrue(datetime.strftime(
            results['ends'], "%Y-%m-%d %H:%M") == '2017-01-02 00:02')
<commit_msg>skips not needed?<commit_after># -*- coding: utf-8 -*-
# ###
# Copyright (c) 2013-2016, Rice University
# This software is subject to the provisions of the GNU Affero General
# Public License version 3 (AGPLv3).
# See LICENCE.txt for details.
# ###
import unittest
from datetime import datetime

from cnxdb.init import init_db
from pyramid import testing

from .. import use_cases
from ..testing import (
    integration_test_settings,
    db_connection_factory,
    )

# FIXME There is an issue with setting up the celery app more than once.
#       Apparently, creating the app a second time doesn't really create
#       it again. There is some global state hanging around that we can't
#       easily get at. This causes the task results tables used in these
#       views to not exist, because the code believes it's already been
#       initialized.
@unittest.skip("celery is too global")
class PostPublicationsViewsTestCase(unittest.TestCase):
    maxDiff = None

    @classmethod
    def setUpClass(cls):
        cls.settings = integration_test_settings()
        from cnxpublishing.config import CONNECTION_STRING
        cls.db_conn_str = cls.settings[CONNECTION_STRING]
        cls.db_connect = staticmethod(db_connection_factory())

    def setUp(self):
        self.config = testing.setUp(settings=self.settings)
        self.config.include('cnxpublishing.tasks')
        init_db(self.db_conn_str, True)

    def tearDown(self):
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("DROP SCHEMA public CASCADE")
                cursor.execute("CREATE SCHEMA public")
        testing.tearDown()

    @property
    def target(self):
        from ...views.admin import admin_post_publications
        return admin_post_publications

    def test_no_results(self):
        request = testing.DummyRequest()

        resp_data = self.target(request)

        self.assertEqual({'states': []}, resp_data)

    def test(self):
        request = testing.DummyRequest()

        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                # Insert one book into archive.
                book = use_cases.setup_BOOK_in_archive(self, cursor)
                db_conn.commit()

                # Insert some data into the association table.
                cursor.execute("""
INSERT INTO document_baking_result_associations
  (result_id, module_ident)
SELECT
  uuid_generate_v4(),
  (SELECT module_ident FROM modules ORDER BY module_ident DESC LIMIT 1);""")
                db_conn.commit()

                cursor.execute("""\
INSERT INTO document_baking_result_associations
  (result_id, module_ident)
SELECT
  uuid_generate_v4(),
  (SELECT module_ident FROM modules ORDER BY module_ident DESC LIMIT 1);""")

        resp_data = self.target(request)
        self.assertEqual({
            'states': [
                {'created': resp_data['states'][0]['created'],
                 'ident_hash': book.ident_hash,
                 'state': u'PENDING',
                 'state_message': '',
                 'title': 'Book of Infinity'},
                {'created': resp_data['states'][1]['created'],
                 'ident_hash': book.ident_hash,
                 'state': u'PENDING',
                 'state_message': '',
                 'title': 'Book of Infinity'},
                ]}, resp_data)

class ErrorBannerViewsTestCase(unittest.TestCase):
    maxDiff = None

    @classmethod
    def setUpClass(cls):
        cls.settings = integration_test_settings()
        from cnxpublishing.config import CONNECTION_STRING
        cls.db_conn_str = cls.settings[CONNECTION_STRING]
        cls.db_connect = staticmethod(db_connection_factory())

    def setUp(self):
        self.config = testing.setUp(settings=self.settings)
        # self.config.include('cnxpublishing.tasks')
        init_db(self.db_conn_str, True)

    def tearDown(self):
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("DROP SCHEMA public CASCADE")
                cursor.execute("CREATE SCHEMA public")
        testing.tearDown()

    def test_error_banner_get(self):
        request = testing.DummyRequest()

        from ...views.admin import admin_post_error_banner
        defaults = admin_post_error_banner(request)
        self.assertEqual(set(defaults.keys()),
                         set(['start_date', 'start_time',
                              'end_date', 'end_time']))

    def test_error_banner_post(self):
        request = testing.DummyRequest()
        request.POST = {
            'message': 'test message',
            'priority': 1,
            'type': 1,
            'start_date': '2017-01-01',
            'start_time': '00:01',
            'end_date': '2017-01-02',
            'end_time': '00:02',
        }
        from ...views.admin import admin_post_error_banner_POST
        results = admin_post_error_banner_POST(request)

        # assert the error message has been added to the table
        with self.db_connect() as db_conn:
            with db_conn.cursor() as cursor:
                cursor.execute("""SELECT message, priority, starts, ends
                                  from service_state_messages
                                  WHERE message='test message';""")
                result = cursor.fetchone()
                self.assertEqual('test message', result[0])
                self.assertEqual(1, result[1])
                self.assertTrue(datetime.strftime(
                    result[2], "%Y-%m-%d %H:%M") == '2017-01-01 00:01')
                self.assertTrue(datetime.strftime(
                    result[3], "%Y-%m-%d %H:%M") == '2017-01-02 00:02')
                cursor.execute("""DELETE from service_state_messages
                                  WHERE message='test message';""")

        # Assert the correct variables were passed to the template
        self.assertEqual('test message', results['message'])
        self.assertEqual(1, results['priority'])
        self.assertEqual(1, results['type'])
        self.assertTrue(datetime.strftime(
            results['starts'], "%Y-%m-%d %H:%M") == '2017-01-01 00:01')
        self.assertTrue(datetime.strftime(
            results['ends'], "%Y-%m-%d %H:%M") == '2017-01-02 00:02')
<|endoftext|>


<commit_before>import argparse
import dateutil.parser

def get_parser():

    parser = argparse.ArgumentParser(description='Vumi Go Data Tools')
    parser.add_argument(
        '-m', '--msisdn', help='MSISDN to extract messages', required=False)
    parser.add_argument(
        '-d', '--direction', help='Message direction to extract',
        required=False, choices=['inbound', 'outbound', 'all'])
    parser.add_argument(
        '-s', '--start',
        help=('Date time to start from '
              '(as ISO timestamp, e.g. 2013-09-01 01:00:00)'),
        required=False, type=dateutil.parser.parse)
    parser.add_argument(
        '-e', '--end',
        help=('Date time to extract to '
              '(as ISO timestamp, e.g. 2013-09-10 03:00:00)'),
        required=False, type=dateutil.parser.parse)

    return parser
<commit_msg>structuring cmdline<commit_after>import argparse
import dateutil.parser

def get_parser():

    parser = argparse.ArgumentParser(description='Vumi Go Data Tools')

    subparsers = parser.add_subparsers(help='use `command --help`.')

    msisdn_parser = subparsers.add_parser(
        'msisdn', help='Filter on an msisdn')
    msisdn_parser.add_argument(
        '-m', '--msisdn', help='MSISDN to extract messages for',
        required=True)
    msisdn_parser.add_argument(
        '-t', '--addr-type',
        help='Which address type to filter on.',
        choices=['to_addr', 'from_addr'], required=True)

    daterange_parser = subparsers.add_parser(
        'daterange', help='Filter on a date range.')
    daterange_parser.add_argument(
        '-s', '--start',
        help=('Date time to start from '
              '(as ISO timestamp, e.g. 2013-09-01 01:00:00)'),
        required=False, type=dateutil.parser.parse)
    daterange_parser.add_argument(
        '-e', '--end',
        help=('Date time to extract to '
              '(as ISO timestamp, e.g. 2013-09-10 03:00:00)'),
        required=False, type=dateutil.parser.parse)

    direction_parser = subparsers.add_parser(
        'direction', help='Filter on message direction.')
    direction_parser.add_argument(
        '-d', '--direction', help='Message direction to extract',
        required=False, choices=['inbound', 'outbound'])

    return parser
<|endoftext|>


<commit_before>from ..utils import asciireplace, limit_text
from ..exceptions import RowError, FlagError
from ..exceptions.messages import RowObjectMsg as msg
from .formulas import Formula

from tabulate import tabulate
from string import ascii_lowercase, digits

import keyword
accepted_chars = (ascii_lowercase + "_" + digits)

class BaseRow(dict):
    """This Base Class represents a row in a spreadsheet"""

    __slots__ = ["__output__", "__flag__"]

    def __init__(self, data, *args, **kwargs):
        self.__flag__ = 2
        super(BaseRow, self).__init__(data)
        self.__output__ = True
        self(flag=1).construct(*args, **kwargs)
        self.resetflag()

    def __call__(self, flag=1):
        """Changes Flag"""
        # flag 1 = dict_mode
        # flag 2 = psv_mode
        self.__flag__ = flag
        return self

    def __hash__(self):
        return hash((column, self[column]["value"]) for column in sorted(self.keys()))

    def resetflag(self, to=2):
        return self(flag=to)

    def construct(self, *args, **kwargs):
        """This method can be used by inherited objects of :class:`BaseRow` as if it was __init__"""
        pass

    def formula(self, columnname, func, rowref=None, **kwargs):
        """Replaces a cell with a formula

            :param columnname: Name of the column.
            :param func: Function that accepts 1 argument and optionally kwargs.
            :param rowref: The referenced row. Defaults to the current row.
            :param kwargs: Parameters that will be supplied to the function when ran.
            :type columnname: :class:`str`
            :type func: :class:`FunctionType`

            :returns: Nothing
            :rtype: :class:`NoneType`
        """
        if rowref is None:
            rowref = self
        self.setcolumn(columnname, Formula(func, rowref, **kwargs))

    def getformula(self, columnname):
        """Get Formula

            :param columnname: Gets the result of a formula. Used when
            formulas referenced other formulas

        """
        try:
            return self.getcolumn(columnname).call()
        except ValueError:
            raise ValueError(
                msg.getformulamsg.format(columnname))

    @property
    def outputrow(self):
        """Returns a boolean of the current output flag for this row"""
        return self.__output__

    @outputrow.setter
    def outputrow(self, v):
        if not isinstance(v, bool):
            raise TypeError(msg.outputrowmsg.format(bool, type(v)))
        self.__output__ = v

    def __getitem__(self, key):
        if self.__flag__ == 1:
            result = super(BaseRow, self).__getitem__(key)
            self.resetflag()
            return result
        elif self.__flag__ == 2:
            result = self(flag=1)[key]["value"]
            return result
        else:
            raise FlagError(msg.flagmessage)

    def __setitem__(self, key, v):
        if self.__flag__ == 1:
            super(BaseRow, self).__setitem__(key, v)
            self.resetflag()
        elif self.__flag__ == 2:
            result = self(flag=1).update({key:
                                          {"value": v, "org_name": self(flag=1)[key]["org_name"]}})
            self.resetflag()
            return result
        else:
            raise FlagError(msg.flagmessage)

    def __delitem__(self, key):
        if self.__flag__ == 1:
            super(BaseRow, self).__delitem__(key)
            self.resetflag()
        elif self.__flag__ == 2:
            self.__delattr__(key)
        else:
            raise FlagError(msg.flagmessage)

    def getcolumn(self, column):
        """Get a cell by the orginal column name

        :param column: The column name. Can be both long and short form.
        :type column: :class:`str`

        :returns: String of the data, or an int/float if a number/decimal.
        :rtype: :class:`str`, :class:`int`, or :class:`float`
        """
        s = cleanup_name(column)
        if s in self(1).keys():
            self.resetflag()
            return getattr(self, s)
        else:
            raise KeyError("{}".format(column))

    def setcolumn(self, column, value):
        """Set a cell by the orginal column name

            :param column: The column name. Can be both long and short form.
            :param value: The data to be set to the specified column
            :type column: :class:`str`

        """
        s = cleanup_name(column)
        if s in self(1).keys():
            self.resetflag()
            self.__setattr__(s, value)
        else:
            raise Exception("{}".format(column))

    def delcolumn(self, column):
        """Delete a cell by the orginal column name

        :param column: The column name. Can be both long and short form.
        :type column: :class:`str`

        """

        s = cleanup_name(column)
        if s in self.keys():
            self.__delattr__(s)
        else:
            raise KeyError("{}".format(column))

    def __repr__(self):
        return "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self.keys())
        )

    def __str__(self):
        rv = "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self(1).keys())
        )
        self.resetflag()
        return rv

    def __pos__(self):
        self.outputrow = True
        return self

    def __neg__(self):
        self.outputrow = False
        return self

    def __invert__(self):
        self.outputrow = not (self.outputrow)
        return self

    def __getattribute__(self, attr):
        if not super(BaseRow, self).get(attr, False) or self.__flag__ == 1:
            return super(dict, self).__getattribute__(attr)
        else:
            result = (self(flag=1)[attr]["value"])
            return result

    def __getattr__(self, attr):
        s = cleanup_name(attr)
        if s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            raise AttributeError('\'{}\' has no attribute \'{}\''.format(
                type(self), attr))

    def __setattr__(self, attr, v):
        """Allows setting of rows and attributes by using =
            statement"""

        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self(flag=1)[attr]["value"] = v
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            super(BaseRow, self).__setattr__(attr, v)

    def __delattr__(self, attr):
        """Allows deletion of rows and attributes (Makes a row empty) by using
        del statement"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self(flag=1)[attr]["value"] = ""
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            super(BaseRow, self).__delattr__(attr)

    def addcolumn(self, columnname, columndata=""):
        """Adds a column for this row only"""
        short_cn = cleanup_name(columnname)
        if not self.get(short_cn):
            self(flag=1)[short_cn] = {
                "org_name": columnname, "value": columndata}
            self.resetflag()
        else:
            raise Exception("Column already exists.")

    def longcolumn(self, columns=None):
        """
            :params columns: A collection of columns, if supplied the method
                will return only the specified columns.
            :type columns: :class:`tuple`, :class:`list`

            :returns: Generates a :class:`dict` that uses orginal names of
                the column.
            :rtype: :class:`dict`
        """
        newdict = {}
        if columns:
            shortcolumns_check = [cleanup_name(x) for x in columns]
        for k in self(flag=1).keys():
            if columns:
                if not (k in shortcolumns_check):
                    continue
            newdict.update(
                {self(flag=1)[k]["org_name"]: self(flag=1)[k]["value"]})
        self.resetflag()
        return newdict

    def tabulate(self, format="grid", only_ascii=True, columns=None, text_limit=None):
        """Integrates tabulate library with psv

            :param format: A valid format for :class:`tabulate` library.
            :only_ascii: If :data:`True`, only return ascii characters.
            :param columns: Collection of column names that will be included in the
                tabulating.
            :param text_limit: The number of characters to include per cell.
            :type format: :class:`str`

        """
        data = self.longcolumn()
        sortedcolumns = sorted(data) if not columns else columns

        result = tabulate(
            [sortedcolumns] +
            [[limit_text(data[c], text_limit) for c in sortedcolumns]],
            headers="firstrow",
            tablefmt=format)
        if only_ascii:
            return asciireplace(result)
        else:
            return result

#This block was in utils,
# but it relied on a circular reference that re-imported
# a variable everytime this core function was called.
#While less clean, this produces a decent speedup.
non_accepted_key_names = set(tuple(dir(
    BaseRow)) + ("__flag__", "row_obj") + tuple(keyword.kwlist))

def cleanup_name(s):
    result = "".join(filter(lambda x: x in accepted_chars, s.lower()))
    if result in non_accepted_key_names:
        return "psv_" + result
    else:
        return result
    return result<commit_msg>fixed bugged implementation of __hash__ on rowobject<commit_after>from ..utils import asciireplace, limit_text
from ..exceptions import RowError, FlagError
from ..exceptions.messages import RowObjectMsg as msg
from .formulas import Formula

from tabulate import tabulate
from string import ascii_lowercase, digits

import keyword
accepted_chars = (ascii_lowercase + "_" + digits)

class BaseRow(dict):
    """This Base Class represents a row in a spreadsheet"""

    __slots__ = ["__output__", "__flag__"]

    def __init__(self, data, *args, **kwargs):
        self.__flag__ = 2
        super(BaseRow, self).__init__(data)
        self.__output__ = True
        self(flag=1).construct(*args, **kwargs)
        self.resetflag()

    def __call__(self, flag=1):
        """Changes Flag"""
        # flag 1 = dict_mode
        # flag 2 = psv_mode
        self.__flag__ = flag
        return self

    def __hash__(self):
        return hash(tuple((column, self[column]) for column in sorted(self.keys())))

    def resetflag(self, to=2):
        return self(flag=to)

    def construct(self, *args, **kwargs):
        """This method can be used by inherited objects of :class:`BaseRow` as if it was __init__"""
        pass

    def formula(self, columnname, func, rowref=None, **kwargs):
        """Replaces a cell with a formula

            :param columnname: Name of the column.
            :param func: Function that accepts 1 argument and optionally kwargs.
            :param rowref: The referenced row. Defaults to the current row.
            :param kwargs: Parameters that will be supplied to the function when ran.
            :type columnname: :class:`str`
            :type func: :class:`FunctionType`

            :returns: Nothing
            :rtype: :class:`NoneType`
        """
        if rowref is None:
            rowref = self
        self.setcolumn(columnname, Formula(func, rowref, **kwargs))

    def getformula(self, columnname):
        """Get Formula

            :param columnname: Gets the result of a formula. Used when
            formulas referenced other formulas

        """
        try:
            return self.getcolumn(columnname).call()
        except ValueError:
            raise ValueError(
                msg.getformulamsg.format(columnname))

    @property
    def outputrow(self):
        """Returns a boolean of the current output flag for this row"""
        return self.__output__

    @outputrow.setter
    def outputrow(self, v):
        if not isinstance(v, bool):
            raise TypeError(msg.outputrowmsg.format(bool, type(v)))
        self.__output__ = v

    def __getitem__(self, key):
        if self.__flag__ == 1:
            result = super(BaseRow, self).__getitem__(key)
            self.resetflag()
            return result
        elif self.__flag__ == 2:
            result = self(flag=1)[key]["value"]
            return result
        else:
            raise FlagError(msg.flagmessage)

    def __setitem__(self, key, v):
        if self.__flag__ == 1:
            super(BaseRow, self).__setitem__(key, v)
            self.resetflag()
        elif self.__flag__ == 2:
            result = self(flag=1).update({key:
                                          {"value": v, "org_name": self(flag=1)[key]["org_name"]}})
            self.resetflag()
            return result
        else:
            raise FlagError(msg.flagmessage)

    def __delitem__(self, key):
        if self.__flag__ == 1:
            super(BaseRow, self).__delitem__(key)
            self.resetflag()
        elif self.__flag__ == 2:
            self.__delattr__(key)
        else:
            raise FlagError(msg.flagmessage)

    def getcolumn(self, column):
        """Get a cell by the orginal column name

        :param column: The column name. Can be both long and short form.
        :type column: :class:`str`

        :returns: String of the data, or an int/float if a number/decimal.
        :rtype: :class:`str`, :class:`int`, or :class:`float`
        """
        s = cleanup_name(column)
        if s in self(1).keys():
            self.resetflag()
            return getattr(self, s)
        else:
            raise KeyError("{}".format(column))

    def setcolumn(self, column, value):
        """Set a cell by the orginal column name

            :param column: The column name. Can be both long and short form.
            :param value: The data to be set to the specified column
            :type column: :class:`str`

        """
        s = cleanup_name(column)
        if s in self(1).keys():
            self.resetflag()
            self.__setattr__(s, value)
        else:
            raise Exception("{}".format(column))

    def delcolumn(self, column):
        """Delete a cell by the orginal column name

        :param column: The column name. Can be both long and short form.
        :type column: :class:`str`

        """

        s = cleanup_name(column)
        if s in self.keys():
            self.__delattr__(s)
        else:
            raise KeyError("{}".format(column))

    def __repr__(self):
        return "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self.keys())
        )

    def __str__(self):
        rv = "<'{rowname}':{columnamount}>".format(
            rowname=self.__class__.__name__,
            columnamount=len(self(1).keys())
        )
        self.resetflag()
        return rv

    def __pos__(self):
        self.outputrow = True
        return self

    def __neg__(self):
        self.outputrow = False
        return self

    def __invert__(self):
        self.outputrow = not (self.outputrow)
        return self

    def __getattribute__(self, attr):
        if not super(BaseRow, self).get(attr, False) or self.__flag__ == 1:
            return super(dict, self).__getattribute__(attr)
        else:
            result = (self(flag=1)[attr]["value"])
            return result

    def __getattr__(self, attr):
        s = cleanup_name(attr)
        if s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            raise AttributeError('\'{}\' has no attribute \'{}\''.format(
                type(self), attr))

    def __setattr__(self, attr, v):
        """Allows setting of rows and attributes by using =
            statement"""

        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self(flag=1)[attr]["value"] = v
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            super(BaseRow, self).__setattr__(attr, v)

    def __delattr__(self, attr):
        """Allows deletion of rows and attributes (Makes a row empty) by using
        del statement"""
        s = cleanup_name(attr)
        if attr in super(BaseRow, self).keys():
            self(flag=1)[attr]["value"] = ""
        elif s in super(BaseRow, self).keys():
            raise AttributeError((
                "{}{}"
                .format(
                    '\'{}\' has no attribute \'{}\''.format(
                        type(self), attr),
                    ". However, '{s}' is an existing condensed ".format(s=s) +
                    "column name. Only the condensed version is supported."
                    .format(s=s)
                )))
        else:
            super(BaseRow, self).__delattr__(attr)

    def addcolumn(self, columnname, columndata=""):
        """Adds a column for this row only"""
        short_cn = cleanup_name(columnname)
        if not self.get(short_cn):
            self(flag=1)[short_cn] = {
                "org_name": columnname, "value": columndata}
            self.resetflag()
        else:
            raise Exception("Column already exists.")

    def longcolumn(self, columns=None):
        """
            :params columns: A collection of columns, if supplied the method
                will return only the specified columns.
            :type columns: :class:`tuple`, :class:`list`

            :returns: Generates a :class:`dict` that uses orginal names of
                the column.
            :rtype: :class:`dict`
        """
        newdict = {}
        if columns:
            shortcolumns_check = [cleanup_name(x) for x in columns]
        for k in self(flag=1).keys():
            if columns:
                if not (k in shortcolumns_check):
                    continue
            newdict.update(
                {self(flag=1)[k]["org_name"]: self(flag=1)[k]["value"]})
        self.resetflag()
        return newdict

    def tabulate(self, format="grid", only_ascii=True, columns=None, text_limit=None):
        """Integrates tabulate library with psv

            :param format: A valid format for :class:`tabulate` library.
            :only_ascii: If :data:`True`, only return ascii characters.
            :param columns: Collection of column names that will be included in the
                tabulating.
            :param text_limit: The number of characters to include per cell.
            :type format: :class:`str`

        """
        data = self.longcolumn()
        sortedcolumns = sorted(data) if not columns else columns

        result = tabulate(
            [sortedcolumns] +
            [[limit_text(data[c], text_limit) for c in sortedcolumns]],
            headers="firstrow",
            tablefmt=format)
        if only_ascii:
            return asciireplace(result)
        else:
            return result

#This block was in utils,
# but it relied on a circular reference that re-imported
# a variable everytime this core function was called.
#While less clean, this produces a decent speedup.
non_accepted_key_names = set(tuple(dir(
    BaseRow)) + ("__flag__", "row_obj") + tuple(keyword.kwlist))

def cleanup_name(s):
    result = "".join(filter(lambda x: x in accepted_chars, s.lower()))
    if result in non_accepted_key_names:
        return "psv_" + result
    else:
        return result
    return result<|endoftext|>


<commit_before>__author__ = 'Allison MacLeay'

import sys
import os
import argparse
import time
from glob import glob
from tempfile import mkdtemp

"""
Run a command on every file in a directory
"""

def get_cmd(fname, cmd, params):
    """ return command """

    out_name = os.path.join(params['out'], params['name'] + fname)
    cmd = cmd.replace('$INPUT', fname)
    cmd = cmd.replace('$OUT', out_name)

    return cmd

def bsub_cmd(user, log, flags):
    """ return bsub command"""
    if flags != '':  # make sure it is buffered by spaces
        if flags[0] != ' ':
            flags = ' ' + flags
        if flags[-1] != ' ':
            flags = flags + ' '
    return 'bsub -q medium -u ' + user + ' -o ' + flags + os.path.join(log, 'lsf_out.log') + ' -e ' + os.path.join(log,
                                                                                'lsf_err.log') + ' '

def get_names(path, pattern):
    """ Return all unique file prefixes """
    return glob(os.path.join(path, pattern))

def check_done(file_num, path):
    """
    Delay completion of script until all
    files are written
    """
    start = time.time()
    timeout = (24 * 60 * 60)  # 24 hours
    done = 0
    while done == 0:
        files = next(os.walk(path))[2]
        all_closed = 0
        if len(files) == file_num:
            all_closed = 1
            for f in files:
                if f.find('tmp') > 0:
                    all_closed = 0
                    print 'found tmp file ' + f + '.  Waiting...'
                    continue
        if all_closed == 1:
            done = 1
        elif (time.time() - start) > timeout:
            done = 1
            print 'Job timed out'
        else:
            time.sleep(5)
            print 'checking for job completion after waiting %d seconds' % (time.time() - start)
            print 'searching for ' + str(file_num) + ', found ' + str(len(files)) + ' files in ' + path

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog="Batch command helper",
                                     description="Run umitag utility in batches of similarly prefixed names.  "
                                                 "Pass in a command in the format ls $INPUT >> $OUTPUT")
    parser.add_argument('--dir', default='.', help='directory containing input files')
    parser.add_argument('--cmd', default='ls $INPUT',
                        help='command to run on each file in the specified directory')
    parser.add_argument('--out', default='tmp', help='directory to deposit output files')
    parser.add_argument('--batch_name', default='batch', help='name to prepend to processed files')
    parser.add_argument('--pattern', default='*.*', help='match pattern for input files.  default=*.*')
    parser.add_argument('--verbose', default=False, help='verbosity.  default=False')

    # optional bsub options
    parser.add_argument('--bsub', action='store_true', help='use bsub')
    parser.add_argument('--log', help='directory to deposit bsub log files')
    parser.add_argument('--bsub_user', default='am282', help='user name for bsub command. default=am282')
    parser.add_argument('--bsub_mod', default='', help='extra parameters for bsub command')
    parser.add_argument('--output_count', help='number of expected output files per input.  Process will wait to '
                                               'complete until all files are created.  Leave this flag out to avoid '
                                               'the process waiting')
    args = parser.parse_args()

    if args.out == 'tmp':
        print 'WARNING: Output directory was not supplied.  Temporary directory will be used instead.'
    p = {}
    p['path'] = args.dir
    p['name'] = args.batch_name
    if args.out != 'tmp':
        os.system('mkdir -p {}'.format(p['out']))
    else:
        p['out'] = mkdtemp()

    # bsub options
    if args.bsub:
        if not hasattr(args, 'log'):
            raise Exception("ERROR: If using bsub you must specify a directory for the log files.")
        os.system('mkdir -p {}'.format(args.log))
        os.system('ls {} >> {}'.format(p['path'], os.path.join(args.log, 'ls_inputdir.txt')))
        expected_output = int(args.output_count) if hasattr(args, 'output_count') else 0

    files = get_names(args.dir, args.pattern)
    if len(files) < 1:
        print "Error: No file prefixes were found in {}.".format(args.dir)

    count_lsf = 0
    for fname in files:
        if (fname.find('undetermined') > -1):
            # skip undeterminded for now
            cmd = 'echo skipping undetermined files'
        elif not args.bsub:
            cmd = get_cmd(fname, args.cmd, p)
        else:
            cmd = bsub_cmd(args.bsub_user, args.log, args.bsub_mod) + get_cmd(fname, args.cmd, p)
            # Keep track of lsf job for listener
            count_lsf += expected_output  # stays 0 if the process will not wait until completion
        if args.verbose:
            print 'batch command helper running command:\n' + cmd
        os.system(cmd)

    if count_lsf > 0:
        check_done(count_lsf, p['out'])

    if args.verbose:
        print 'batch_process done'
<commit_msg>minor fixes<commit_after>__author__ = 'Allison MacLeay'

import sys
import os
import argparse
import time
from glob import glob
from tempfile import mkdtemp

"""
Run a command on every file in a directory
"""

def get_cmd(fname, cmd, params):
    """ return command """

    out_name = os.path.join(params['out'], params['name'] + fname)
    cmd = cmd.replace('$INPUT', fname)
    cmd = cmd.replace('$OUT', out_name)

    return cmd

def bsub_cmd(user, log, flags):
    """ return bsub command"""
    if flags != '':  # make sure it is buffered by spaces
        if flags[0] != ' ':
            flags = ' ' + flags
        if flags[-1] != ' ':
            flags = flags + ' '
    return 'bsub -q medium -u ' + user + ' -o ' + flags + os.path.join(log, 'lsf_out.log') + ' -e ' + os.path.join(log,
                                                                                'lsf_err.log') + ' '

def get_names(path, pattern):
    """ Return all unique file prefixes """
    return glob(os.path.join(path, pattern))

def check_done(file_num, path):
    """
    Delay completion of script until all
    files are written
    """
    start = time.time()
    timeout = (24 * 60 * 60)  # 24 hours
    done = 0
    while done == 0:
        files = next(os.walk(path))[2]
        all_closed = 0
        if len(files) == file_num:
            all_closed = 1
            for f in files:
                if f.find('tmp') > 0:
                    all_closed = 0
                    print 'found tmp file ' + f + '.  Waiting...'
                    continue
        if all_closed == 1:
            done = 1
        elif (time.time() - start) > timeout:
            done = 1
            print 'Job timed out'
        else:
            time.sleep(5)
            print 'checking for job completion after waiting %d seconds' % (time.time() - start)
            print 'searching for ' + str(file_num) + ', found ' + str(len(files)) + ' files in ' + path

if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog="Batch command helper",
                                     description="Run umitag utility in batches of similarly prefixed names.  "
                                                 "Pass in a command in the format ls $INPUT >> $OUT")
    parser.add_argument('--dir', default='.', help='directory containing input files')
    parser.add_argument('--cmd', default='ls $INPUT',
                        help='command to run on each file in the specified directory')
    parser.add_argument('--out', default='tmp', help='directory to deposit output files')
    parser.add_argument('--batch_name', default='batch', help='name to prepend to processed files')
    parser.add_argument('--pattern', default='*.*', help='match pattern for input files.  default=*.*')
    parser.add_argument('--verbose', action='store_true', help='verbosity.  default=False')

    # optional bsub options
    parser.add_argument('--bsub', action='store_true', help='use bsub')
    parser.add_argument('--log', help='directory to deposit bsub log files')
    parser.add_argument('--bsub_user', default='am282', help='user name for bsub command. default=am282')
    parser.add_argument('--bsub_mod', default='', help='extra parameters for bsub command')
    parser.add_argument('--output_count', help='number of expected output files per input.  Process will wait to '
                                               'complete until all files are created.  Leave this flag out to avoid '
                                               'the process waiting')
    args = parser.parse_args()

    if args.out == 'tmp':
        print 'WARNING: Output directory was not supplied.  Temporary directory will be used instead.'
    p = {}
    p['path'] = args.dir
    p['name'] = args.batch_name
    if args.out != 'tmp':
        os.system('mkdir -p {}'.format(p['out']))
    else:
        p['out'] = mkdtemp()

    # bsub options
    if args.bsub:
        if not hasattr(args, 'log'):
            raise Exception("ERROR: If using bsub you must specify a directory for the log files.")
        os.system('mkdir -p {}'.format(args.log))
        os.system('ls {} >> {}'.format(p['path'], os.path.join(args.log, 'ls_inputdir.txt')))
        expected_output = int(args.output_count) if hasattr(args, 'output_count') else 0

    files = get_names(args.dir, args.pattern)
    if len(files) < 1:
        print "Error: No file prefixes were found in {}.".format(args.dir)

    count_lsf = 0
    for fname in files:
        if (fname.find('undetermined') > -1):
            # skip undeterminded for now
            cmd = 'echo skipping undetermined files'
        elif not args.bsub:
            cmd = get_cmd(fname, args.cmd, p)
        else:
            cmd = bsub_cmd(args.bsub_user, args.log, args.bsub_mod) + get_cmd(fname, args.cmd, p)
            # Keep track of lsf job for listener
            count_lsf += expected_output  # stays 0 if the process will not wait until completion
        if args.verbose:
            print 'batch command helper running command:\n' + cmd
        os.system(cmd)

    if count_lsf > 0:
        check_done(count_lsf, p['out'])

    if args.verbose:
        print 'batch_process done'
<|endoftext|>


<commit_before>import os,sys,inspect
dir_current = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
dir_parent = os.path.dirname(dir_current)
sys.path.insert(0, dir_parent)

import sys
import logging
import argparse
import asyncio
from aiohttp import web
import aiohttp_cors

import utilery
from utilery.config_handler import Configs
from utilery.plugins import Plugins
from utilery.db_handler import DB
from utilery.tile_handler import ServePBF, ServeGeoJSON, ServeJSON, TileJson

# parse arguments from command line
ARGS = argparse.ArgumentParser(description="asyncio python tile server")
ARGS.add_argument('--server_configs',
                  action="store",
                  dest='server_configs',
                  default='../configs_server_example.yaml',
                  help='The YAML database config file')
ARGS.add_argument('--layer_recipes',
                  action="store",
                  dest="layer_recipes",
                  default='../layer_recipe_examples/',
                  help='The YAML layers configs file')
ARGS.add_argument('--host',
                  action="store",
                  dest="host",
                  default='0.0.0.0',
                  help='The host IP address')
ARGS.add_argument('--port',
                  action="store",
                  dest="port",
                  default=8080,
                  help='The host port')
ARGS.add_argument('--version', action='version', version=utilery.VERSION)

args = ARGS.parse_args()

# set the server configs
Configs.init_server_configs(os.path.abspath(args.server_configs))

# setup logging
logging.basicConfig(stream=sys.stdout, level=Configs.server['log_level'].upper())
logger = logging.getLogger(__name__)
logger.info('STARTING ASYNCIO TILE SERVER APP')

# set the layer configs
for file in os.listdir(os.path.abspath(args.layer_recipes)):
    if file.endswith('.yaml') or file.endswith('.yml'):
        logger.info('Adding layer: {0}'.format(file))
        Configs.init_layer_recipes(os.path.join(os.path.abspath(args.layer_recipes), file))

# load plugins
logger.info('Loading plugins')
Plugins.load()

# create server app
logger.info('Creating the server app')
app = web.Application()

# setup url routes and corresponding handlers
async def request_pbf(request):
    content_type, body = await ServePBF.serve(request)
    return web.Response(content_type=content_type, body=body)
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.pbf', request_pbf)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.pbf', request_pbf)
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.mvt', request_pbf)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.mvt', request_pbf)

async def request_geojson(request):
    content_type, body = await ServeGeoJSON.serve(request)
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.geojson', request_geojson)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.geojson', request_geojson)

async def request_json(request):
    content_type, body = await ServeJSON.serve(request)
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.json', request_json)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.json', request_json)

async def request_tilejson(request):
    content_type, body = TileJson.get()
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/tilejson/mvt.json', request_tilejson)

# configure CORS TODO: look at how this interacts with plugin?
cors = aiohttp_cors.setup(app, defaults={
    Configs.server['CORS']: aiohttp_cors.ResourceOptions(
        allow_credentials=True,
        expose_headers='*',
        allow_headers='*'
    )
})
for route in list(app.router.routes()):
    cors.add(route)

# start the database pool
logger.info('Creating the database pool')
Configs.DB = DB() # create the DB instance
loop = asyncio.get_event_loop()
for db_name, dsn_string in Configs.server['databases'].items():
    loop.run_until_complete(DB.connect(db_name, dsn_string)) # this is an awaitable async so use run_until_complete

# start the server
logger.info('Starting the server app')
web.run_app(app, host=args.host, port=args.port)
<commit_msg>added args for host and port<commit_after>import os,sys,inspect
dir_current = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
dir_parent = os.path.dirname(dir_current)
sys.path.insert(0, dir_parent)

import sys
import logging
import argparse
import asyncio
from aiohttp import web
import aiohttp_cors

import utilery
from utilery.config_handler import Configs
from utilery.plugins import Plugins
from utilery.db_handler import DB
from utilery.tile_handler import ServePBF, ServeGeoJSON, ServeJSON, TileJson

# parse arguments from command line
ARGS = argparse.ArgumentParser(description="asyncio python tile server")
ARGS.add_argument('--server_configs',
                  action="store",
                  dest='server_configs',
                  default='../configs_server_example.yaml',
                  help='The YAML database config file')
ARGS.add_argument('--layer_recipes',
                  action="store",
                  dest="layer_recipes",
                  default='../layer_recipe_examples/',
                  help='The YAML layers configs file')
ARGS.add_argument('--host',
                  action="store",
                  dest="host",
                  default='0.0.0.0',
                  help='The host IP address')
ARGS.add_argument('--port',
                  action="store",
                  dest="port",
                  default=8080,
                  type=int,
                  help='The host port')
ARGS.add_argument('--version', action='version', version=utilery.VERSION)

args = ARGS.parse_args()

# set the server configs
Configs.init_server_configs(os.path.abspath(args.server_configs))

# setup logging
logging.basicConfig(stream=sys.stdout, level=Configs.server['log_level'].upper())
logger = logging.getLogger(__name__)
logger.info('STARTING ASYNCIO TILE SERVER APP')

# set the layer configs
for file in os.listdir(os.path.abspath(args.layer_recipes)):
    if file.endswith('.yaml') or file.endswith('.yml'):
        logger.info('Adding layer: {0}'.format(file))
        Configs.init_layer_recipes(os.path.join(os.path.abspath(args.layer_recipes), file))

# load plugins
logger.info('Loading plugins')
Plugins.load()

# create server app
logger.info('Creating the server app')
app = web.Application()

# setup url routes and corresponding handlers
async def request_pbf(request):
    content_type, body = await ServePBF.serve(request)
    return web.Response(content_type=content_type, body=body)
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.pbf', request_pbf)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.pbf', request_pbf)
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.mvt', request_pbf)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.mvt', request_pbf)

async def request_geojson(request):
    content_type, body = await ServeGeoJSON.serve(request)
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.geojson', request_geojson)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.geojson', request_geojson)

async def request_json(request):
    content_type, body = await ServeJSON.serve(request)
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/{layers}/{z}/{x}/{y}.json', request_json)
app.router.add_route('GET', '/{recipe}/{layers}/{z}/{x}/{y}.json', request_json)

async def request_tilejson(request):
    content_type, body = TileJson.get()
    return web.Response(content_type=content_type, body=body.encode())
app.router.add_route('GET', '/tilejson/mvt.json', request_tilejson)

# configure CORS TODO: look at how this interacts with plugin?
cors = aiohttp_cors.setup(app, defaults={
    Configs.server['CORS']: aiohttp_cors.ResourceOptions(
        allow_credentials=True,
        expose_headers='*',
        allow_headers='*'
    )
})
for route in list(app.router.routes()):
    cors.add(route)

# start the database pool
logger.info('Creating the database pool')
Configs.DB = DB() # create the DB instance
loop = asyncio.get_event_loop()
for db_name, dsn_string in Configs.server['databases'].items():
    loop.run_until_complete(DB.connect(db_name, dsn_string)) # this is an awaitable async so use run_until_complete

# start the server
logger.info('Starting the server app')
web.run_app(app, host=args.host, port=args.port)
<|endoftext|>