<commit_before>#!/usr/bin/python3

# Copyright 2017-2020 Yury Gribov
#
# The MIT License (MIT)
#
# Use of this source code is governed by MIT license that can be
# found in the LICENSE.txt file.

import sys
import os.path
import re
import subprocess
import argparse
import string
import configparser

me = os.path.basename(__file__)
root = os.path.dirname(__file__)

def warn(msg):
  """Emits a nicely-decorated warning."""
  sys.stderr.write('%s: warning: %s\n' % (me, msg))

def error(msg):
  """Emits a nicely-decorated error and exits."""
  sys.stderr.write('%s: error: %s\n' % (me, msg))
  sys.exit(1)

def run(args, input=''):
  """Runs external program and aborts on error."""
  p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE)
  out, err = p.communicate(input=input.encode('utf-8'))
  out = out.decode('utf-8')
  err = err.decode('utf-8')
  if p.returncode != 0 or err:
    error("%s failed with retcode %d:\n%s" % (args[0], p.returncode, err))
  return out, err

def make_toc(words):
  return {i: n for i, n in enumerate(words)}

def parse_row(words, toc, hex_keys):
  vals = {k: (words[i] if i < len(words) else '') for i, k in toc.items()}
  for k in hex_keys:
    if vals[k]:
      vals[k] = int(vals[k], 16)
  return vals

def collect_syms(f):
  """Collect ELF dynamic symtab."""

  out, err = run(["readelf", "-W", "--dyn-syms", f])

  toc = None
  syms = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    words = re.split(r' +', line)
    if line.startswith('Num'):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      # Colons are different across readelf versions so get rid of them.
      toc = make_toc(map(lambda n: n.replace(':', ''), words))
    elif toc is not None:
      sym = parse_row(words, toc, ['Value'])
      name = sym['Name']
      if '@' in name:
        sym['Default'] = '@@' in name
        name, ver = re.split(r'@+', name)
        sym['Name'] = name
        sym['Version'] = ver
      else:
        sym['Default'] = True
        sym['Version'] = None
      syms.append(sym)

  if toc is None:
    error("failed to analyze %s" % f)

  # Also collected demangled names
  if syms:
    out, err = run(['c++filt'], '\n'.join((sym['Name'] for sym in syms)))
    for i, name in enumerate(out.split("\n")):
      syms[i]['Demangled Name'] = name

  return syms

def collect_relocs(f):
  """Collect ELF dynamic relocs."""

  out, err = run(["readelf", "-r", f])

  toc = None
  rels = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    if re.match(r'^\s*Offset', line):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      words = re.split(r'\s\s+', line)  # "Sym. Name + Addend"
      toc = make_toc(words)
    elif toc is not None:
      line = re.sub(r' \+ ', '+', line)
      words = re.split(r'\s+', line)
      rel = parse_row(words, toc, ['Offset', 'Info'])
      rels.append(rel)
      # Split symbolic representation
      sym_name = 'Sym. Name + Addend'
      if rel[sym_name]:
        p = rel[sym_name].split('+')
        if len(p) == 1:
          p = ['', p[0]]
        rel[sym_name] = (p[0], int(p[1], 16))

  if toc is None:
    error("failed to analyze %s" % f)

  return rels

def collect_sections(f):
  """Collect section info from ELF."""

  out, err = run(["readelf", "-SW", f])

  toc = None
  sections = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    line = re.sub(r'\[\s+', '[', line)
    words = re.split(r' +', line)
    if line.startswith('[Nr]'):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      toc = make_toc(words)
    elif line.startswith('[') and toc is not None:
      sec = parse_row(words, toc, ['Address', 'Off', 'Size'])
      if 'A' in sec['Flg']:  # Allocatable section?
        sections.append(sec)

  if toc is None:
    error("failed to analyze %s" % f)

  return sections

def read_unrelocate_data(f, chunks):
  """Collect unrelocated data from ELF."""

  secs = collect_sections(f)
  secs.sort(key=lambda s: s['Address'])

  # TODO: read bytes for each chunk (via -x)
  data = []

  return data

def main():
  parser = argparse.ArgumentParser(description="Generate wrappers for shared library functions.",
                                   formatter_class=argparse.RawDescriptionHelpFormatter,
                                   epilog="""\
Examples:
  $ python3 {0} /usr/lib/x86_64-linux-gnu/libaccountsservice.so.0
  Generating libaccountsservice.so.0.tramp.S...
  Generating libaccountsservice.so.0.init.c...
""".format(me))

  parser.add_argument('library',
                      metavar='LIB',
                      help="Library to be wrapped.")
  parser.add_argument('--verbose', '-v',
                      help="Print diagnostic info",
                      action='count',
                      default=0)
  parser.add_argument('--dlopen-callback',
                      help="Call user-provided custom callback to load library instead of dlopen",
                      default='')
  parser.add_argument('--dlopen',
                      help="Emit dlopen call (default)",
                      dest='dlopen', action='store_true', default=True)
  parser.add_argument('--no-dlopen',
                      help="Do not emit dlopen call (user must load library himself)",
                      dest='dlopen', action='store_false')
  parser.add_argument('--library-load-name',
                      help="Use custom name for dlopened library (default is LIB)")
  parser.add_argument('--lazy-load',
                      help="Load library lazily on first call to one of it's functions (default)",
                      dest='lazy_load', action='store_true', default=True)
  parser.add_argument('--no-lazy-load',
                      help="Load library eagerly at program start",
                      dest='lazy_load', action='store_false')
  parser.add_argument('--vtables',
                      help="Intercept virtual tables (EXPERIMENTAL)",
                      dest='vtables', action='store_true', default=False)
  parser.add_argument('--no-vtables',
                      help="Do not intercept virtual tables (default)",
                      dest='vtables', action='store_false')
  parser.add_argument('--target',
                      help="Target platform triple e.g. x86_64-unknown-linux-gnu or arm-none-eabi (atm x86_64, i[0-9]86, arm/armhf and aarch64 are supported)",
                      default='x86_64')
  parser.add_argument('--symbol-list',
                      help="Path to file with symbols that should be present in wrapper (all by default)")
  parser.add_argument('--symbol-prefix',
                      metavar='PFX',
                      help="Prefix wrapper symbols with PFX",
                      default='')
  parser.add_argument('-q', '--quiet',
                      help="Do not print progress info",
                      action='store_true')
  parser.add_argument('--outdir', '-o',
                      help="Path to create wrapper at",
                      default='./')

  args = parser.parse_args()

  input_name = args.library
  verbose = args.verbose
  dlopen_callback = args.dlopen_callback
  dlopen = args.dlopen
  lazy_load = args.lazy_load
  load_name = args.library_load_name if args.library_load_name is not None else os.path.basename(input_name)
  if args.target.startswith('arm'):
    target = 'arm'  # Handle armhf-...
  elif re.match(r'^i[0-9]86', args.target):
      target = 'i386'
  else:
    target = args.target.split('-')[0]
  quiet = args.quiet
  outdir = args.outdir

  if args.symbol_list is None:
    funs = None
  else:
    with open(args.symbol_list, 'r') as f:
      funs = []
      for line in re.split(r'\r?\n', f.read()):
        line = re.sub(r'#.*', '', line)
        line = line.strip()
        if line:
          funs.append(line)

  # Collect target info

  target_dir = os.path.join(root, 'arch', target)

  if not os.path.exists(target_dir):
    error("unknown architecture '%s'" % target)

  cfg = configparser.ConfigParser(inline_comment_prefixes=';')
  cfg.read(target_dir + '/config.ini')

  ptr_size = int(cfg['Arch']['PointerSize'])

  def is_exported(s):
    return (s['Type'] != 'LOCAL'
            and s['Ndx'] != 'UND'
            and s['Name'] not in ['', '_init', '_fini'])

  syms = list(filter(is_exported, collect_syms(input_name)))

  exported_data = [s['Name'] for s in syms if s['Type'] == 'OBJECT' and not s['Name'].startswith('_')]
  if exported_data:
    warn("library '%s' contains data symbols which won't be intercepted: %s" % (input_name, ', '.join(exported_data)))

  # Collect functions
  # TODO: warn if user-specified functions are missing

  if funs is None:
    orig_funs = filter(lambda s: s['Type'] == 'FUNC', syms)

    funs = []
    warn_versioned = False
    for s in orig_funs:
      if s['Version'] is not None:
        # TODO: support versions
        if not warn_versioned:
          warn("library %s contains versioned symbols which are NYI" % input_name)
          warn_versioned = True
        if verbose:
          print("Skipping versioned symbol %s" % s['Name'])
        continue
      funs.append(s['Name'])

    if not funs and not quiet:
      print("no public functions were found in %s" % input_name)

  if verbose:
    print("Exported functions:")
    for i, fun in enumerate(funs):
      print("  {0}: {1}".format(i, fun))

  # Collect vtables

  if args.vtables:
    vtabs = {}

    for s in syms:
      m = re.match(r'^(vtable|typeinfo|typeinfo name) for (.*)', s['Demangled Name'])
      if m is not None and is_exported(s):
        typ, cls = m.groups()
        vtabs.setdefault(cls, {})[typ] = s

    rels = collect_relocs(input_name)
    secs = collect_sections(input_name)

    # TODO: collect vtable raw contents and relocations

    if verbose:
      print("Exported vtables:")
      for i, (cls, _) in enumerate(sorted(vtabs.items())):
        print("  {0}: {1}".format(i, cls))
      print("Relocs:")
      for rel in rels:
        print("  {0}: {1}".format(rel['Offset'], rel['Sym. Name + Addend']))
      print("Sections:")
      for sec in secs:
        print("  {}: [{:x}, {:x}), at {:x}".format(sec['Name'], sec['Address'], sec['Address'] + sec['Size'], sec['Off']))

  # Generate assembly code

  suffix = os.path.basename(load_name)
  lib_suffix = re.sub(r'[^a-zA-Z_0-9]+', '_', suffix)

  tramp_file = '%s.tramp.S' % suffix
  with open(os.path.join(outdir, tramp_file), 'w') as f:
    if not quiet:
      print("Generating %s..." % tramp_file)
    with open(target_dir + '/table.S.tpl', 'r') as t:
      table_text = string.Template(t.read()).substitute(
        lib_suffix=lib_suffix,
        table_size=ptr_size*(len(funs) + 1))
    f.write(table_text)

    with open(target_dir + '/trampoline.S.tpl', 'r') as t:
      tramp_tpl = string.Template(t.read())

    for i, name in enumerate(funs):
      tramp_text = tramp_tpl.substitute(
        lib_suffix=lib_suffix,
        sym=args.symbol_prefix + name,
        offset=i*ptr_size,
        number=i)
      f.write(tramp_text)

  # Generate C code

  init_file = '%s.init.c' % suffix
  with open(os.path.join(outdir, init_file), 'w') as f:
    if not quiet:
      print("Generating %s..." % init_file)
    with open(os.path.join(root, 'arch/common/init.c.tpl'), 'r') as t:
      init_text = string.Template(t.read()).substitute(
        lib_suffix=lib_suffix,
        load_name=load_name,
        dlopen_callback=dlopen_callback,
        has_dlopen_callback=int(bool(dlopen_callback)),
        no_dlopen=not int(dlopen),
        lazy_load=int(lazy_load),
        sym_names=',\n  '.join('"%s"' % name for name in funs))
    f.write(init_text)

if __name__ == '__main__':
  main()
<commit_msg>Fixed warning for data symbols.<commit_after>#!/usr/bin/python3

# Copyright 2017-2020 Yury Gribov
#
# The MIT License (MIT)
#
# Use of this source code is governed by MIT license that can be
# found in the LICENSE.txt file.

import sys
import os.path
import re
import subprocess
import argparse
import string
import configparser

me = os.path.basename(__file__)
root = os.path.dirname(__file__)

def warn(msg):
  """Emits a nicely-decorated warning."""
  sys.stderr.write('%s: warning: %s\n' % (me, msg))

def error(msg):
  """Emits a nicely-decorated error and exits."""
  sys.stderr.write('%s: error: %s\n' % (me, msg))
  sys.exit(1)

def run(args, input=''):
  """Runs external program and aborts on error."""
  p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE)
  out, err = p.communicate(input=input.encode('utf-8'))
  out = out.decode('utf-8')
  err = err.decode('utf-8')
  if p.returncode != 0 or err:
    error("%s failed with retcode %d:\n%s" % (args[0], p.returncode, err))
  return out, err

def make_toc(words):
  return {i: n for i, n in enumerate(words)}

def parse_row(words, toc, hex_keys):
  vals = {k: (words[i] if i < len(words) else '') for i, k in toc.items()}
  for k in hex_keys:
    if vals[k]:
      vals[k] = int(vals[k], 16)
  return vals

def collect_syms(f):
  """Collect ELF dynamic symtab."""

  out, err = run(["readelf", "-W", "--dyn-syms", f])

  toc = None
  syms = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    words = re.split(r' +', line)
    if line.startswith('Num'):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      # Colons are different across readelf versions so get rid of them.
      toc = make_toc(map(lambda n: n.replace(':', ''), words))
    elif toc is not None:
      sym = parse_row(words, toc, ['Value'])
      name = sym['Name']
      if '@' in name:
        sym['Default'] = '@@' in name
        name, ver = re.split(r'@+', name)
        sym['Name'] = name
        sym['Version'] = ver
      else:
        sym['Default'] = True
        sym['Version'] = None
      syms.append(sym)

  if toc is None:
    error("failed to analyze %s" % f)

  # Also collected demangled names
  if syms:
    out, err = run(['c++filt'], '\n'.join((sym['Name'] for sym in syms)))
    for i, name in enumerate(out.split("\n")):
      syms[i]['Demangled Name'] = name

  return syms

def collect_relocs(f):
  """Collect ELF dynamic relocs."""

  out, err = run(["readelf", "-r", f])

  toc = None
  rels = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    if re.match(r'^\s*Offset', line):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      words = re.split(r'\s\s+', line)  # "Sym. Name + Addend"
      toc = make_toc(words)
    elif toc is not None:
      line = re.sub(r' \+ ', '+', line)
      words = re.split(r'\s+', line)
      rel = parse_row(words, toc, ['Offset', 'Info'])
      rels.append(rel)
      # Split symbolic representation
      sym_name = 'Sym. Name + Addend'
      if rel[sym_name]:
        p = rel[sym_name].split('+')
        if len(p) == 1:
          p = ['', p[0]]
        rel[sym_name] = (p[0], int(p[1], 16))

  if toc is None:
    error("failed to analyze %s" % f)

  return rels

def collect_sections(f):
  """Collect section info from ELF."""

  out, err = run(["readelf", "-SW", f])

  toc = None
  sections = []
  for line in out.splitlines():
    line = line.strip()
    if not line:
      continue
    line = re.sub(r'\[\s+', '[', line)
    words = re.split(r' +', line)
    if line.startswith('[Nr]'):  # Header?
      if toc is not None:
        error("multiple headers in output of readelf")
      toc = make_toc(words)
    elif line.startswith('[') and toc is not None:
      sec = parse_row(words, toc, ['Address', 'Off', 'Size'])
      if 'A' in sec['Flg']:  # Allocatable section?
        sections.append(sec)

  if toc is None:
    error("failed to analyze %s" % f)

  return sections

def read_unrelocate_data(f, chunks):
  """Collect unrelocated data from ELF."""

  secs = collect_sections(f)
  secs.sort(key=lambda s: s['Address'])

  # TODO: read bytes for each chunk (via -x)
  data = []

  return data

def main():
  parser = argparse.ArgumentParser(description="Generate wrappers for shared library functions.",
                                   formatter_class=argparse.RawDescriptionHelpFormatter,
                                   epilog="""\
Examples:
  $ python3 {0} /usr/lib/x86_64-linux-gnu/libaccountsservice.so.0
  Generating libaccountsservice.so.0.tramp.S...
  Generating libaccountsservice.so.0.init.c...
""".format(me))

  parser.add_argument('library',
                      metavar='LIB',
                      help="Library to be wrapped.")
  parser.add_argument('--verbose', '-v',
                      help="Print diagnostic info",
                      action='count',
                      default=0)
  parser.add_argument('--dlopen-callback',
                      help="Call user-provided custom callback to load library instead of dlopen",
                      default='')
  parser.add_argument('--dlopen',
                      help="Emit dlopen call (default)",
                      dest='dlopen', action='store_true', default=True)
  parser.add_argument('--no-dlopen',
                      help="Do not emit dlopen call (user must load library himself)",
                      dest='dlopen', action='store_false')
  parser.add_argument('--library-load-name',
                      help="Use custom name for dlopened library (default is LIB)")
  parser.add_argument('--lazy-load',
                      help="Load library lazily on first call to one of it's functions (default)",
                      dest='lazy_load', action='store_true', default=True)
  parser.add_argument('--no-lazy-load',
                      help="Load library eagerly at program start",
                      dest='lazy_load', action='store_false')
  parser.add_argument('--vtables',
                      help="Intercept virtual tables (EXPERIMENTAL)",
                      dest='vtables', action='store_true', default=False)
  parser.add_argument('--no-vtables',
                      help="Do not intercept virtual tables (default)",
                      dest='vtables', action='store_false')
  parser.add_argument('--target',
                      help="Target platform triple e.g. x86_64-unknown-linux-gnu or arm-none-eabi (atm x86_64, i[0-9]86, arm/armhf and aarch64 are supported)",
                      default='x86_64')
  parser.add_argument('--symbol-list',
                      help="Path to file with symbols that should be present in wrapper (all by default)")
  parser.add_argument('--symbol-prefix',
                      metavar='PFX',
                      help="Prefix wrapper symbols with PFX",
                      default='')
  parser.add_argument('-q', '--quiet',
                      help="Do not print progress info",
                      action='store_true')
  parser.add_argument('--outdir', '-o',
                      help="Path to create wrapper at",
                      default='./')

  args = parser.parse_args()

  input_name = args.library
  verbose = args.verbose
  dlopen_callback = args.dlopen_callback
  dlopen = args.dlopen
  lazy_load = args.lazy_load
  load_name = args.library_load_name if args.library_load_name is not None else os.path.basename(input_name)
  if args.target.startswith('arm'):
    target = 'arm'  # Handle armhf-...
  elif re.match(r'^i[0-9]86', args.target):
      target = 'i386'
  else:
    target = args.target.split('-')[0]
  quiet = args.quiet
  outdir = args.outdir

  if args.symbol_list is None:
    funs = None
  else:
    with open(args.symbol_list, 'r') as f:
      funs = []
      for line in re.split(r'\r?\n', f.read()):
        line = re.sub(r'#.*', '', line)
        line = line.strip()
        if line:
          funs.append(line)

  # Collect target info

  target_dir = os.path.join(root, 'arch', target)

  if not os.path.exists(target_dir):
    error("unknown architecture '%s'" % target)

  cfg = configparser.ConfigParser(inline_comment_prefixes=';')
  cfg.read(target_dir + '/config.ini')

  ptr_size = int(cfg['Arch']['PointerSize'])

  def is_exported(s):
    return (s['Type'] != 'LOCAL'
            and s['Ndx'] != 'UND'
            and s['Name'] not in ['', '_init', '_fini'])

  syms = list(filter(is_exported, collect_syms(input_name)))

  def is_data_symbol(s):
    return (s['Type'] == 'OBJECT'
            # Allow vtables if --vtables is on
            and not (' for ' in s['Demangled Name'] and args.vtables))

  exported_data = [s['Name'] for s in syms if is_data_symbol(s)]
  if exported_data:
    warn("library '%s' contains data symbols which won't be intercepted: %s" % (input_name, ', '.join(exported_data)))

  # Collect functions
  # TODO: warn if user-specified functions are missing

  if funs is None:
    orig_funs = filter(lambda s: s['Type'] == 'FUNC', syms)

    funs = []
    warn_versioned = False
    for s in orig_funs:
      if s['Version'] is not None:
        # TODO: support versions
        if not warn_versioned:
          warn("library %s contains versioned symbols which are NYI" % input_name)
          warn_versioned = True
        if verbose:
          print("Skipping versioned symbol %s" % s['Name'])
        continue
      funs.append(s['Name'])

    if not funs and not quiet:
      print("no public functions were found in %s" % input_name)

  if verbose:
    print("Exported functions:")
    for i, fun in enumerate(funs):
      print("  {0}: {1}".format(i, fun))

  # Collect vtables

  if args.vtables:
    vtabs = {}

    for s in syms:
      m = re.match(r'^(vtable|typeinfo|typeinfo name) for (.*)', s['Demangled Name'])
      if m is not None and is_exported(s):
        typ, cls = m.groups()
        vtabs.setdefault(cls, {})[typ] = s

    rels = collect_relocs(input_name)
    secs = collect_sections(input_name)

    # TODO: collect vtable raw contents and relocations

    if verbose:
      print("Exported vtables:")
      for i, (cls, _) in enumerate(sorted(vtabs.items())):
        print("  {0}: {1}".format(i, cls))
      print("Relocs:")
      for rel in rels:
        print("  {0}: {1}".format(rel['Offset'], rel['Sym. Name + Addend']))
      print("Sections:")
      for sec in secs:
        print("  {}: [{:x}, {:x}), at {:x}".format(sec['Name'], sec['Address'], sec['Address'] + sec['Size'], sec['Off']))

  # Generate assembly code

  suffix = os.path.basename(load_name)
  lib_suffix = re.sub(r'[^a-zA-Z_0-9]+', '_', suffix)

  tramp_file = '%s.tramp.S' % suffix
  with open(os.path.join(outdir, tramp_file), 'w') as f:
    if not quiet:
      print("Generating %s..." % tramp_file)
    with open(target_dir + '/table.S.tpl', 'r') as t:
      table_text = string.Template(t.read()).substitute(
        lib_suffix=lib_suffix,
        table_size=ptr_size*(len(funs) + 1))
    f.write(table_text)

    with open(target_dir + '/trampoline.S.tpl', 'r') as t:
      tramp_tpl = string.Template(t.read())

    for i, name in enumerate(funs):
      tramp_text = tramp_tpl.substitute(
        lib_suffix=lib_suffix,
        sym=args.symbol_prefix + name,
        offset=i*ptr_size,
        number=i)
      f.write(tramp_text)

  # Generate C code

  init_file = '%s.init.c' % suffix
  with open(os.path.join(outdir, init_file), 'w') as f:
    if not quiet:
      print("Generating %s..." % init_file)
    with open(os.path.join(root, 'arch/common/init.c.tpl'), 'r') as t:
      init_text = string.Template(t.read()).substitute(
        lib_suffix=lib_suffix,
        load_name=load_name,
        dlopen_callback=dlopen_callback,
        has_dlopen_callback=int(bool(dlopen_callback)),
        no_dlopen=not int(dlopen),
        lazy_load=int(lazy_load),
        sym_names=',\n  '.join('"%s"' % name for name in funs))
    f.write(init_text)

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


<commit_before>import os

from macholib.MachOGraph import MachOGraph, MissingMachO
from macholib.util import iter_platform_files, in_system_path, mergecopy, \
    mergetree, flipwritable, has_filename_filter
from macholib.dyld import framework_info
from collections import deque

class ExcludedMachO(MissingMachO):
    pass

class FilteredMachOGraph(MachOGraph):
    def __init__(self, delegate, *args, **kwargs):
        super(FilteredMachOGraph, self).__init__(*args, **kwargs)
        self.delegate = delegate

    def createNode(self, cls, name):
        cls = self.delegate.getClass(name, cls)
        res = super(FilteredMachOGraph, self).createNode(cls, name)
        return res

    def locate(self, filename):
        newname = super(FilteredMachOGraph, self).locate(filename)
        if newname is None:
            return None
        return self.delegate.locate(newname)

class MachOStandalone(object):
    def __init__(self, base, dest=None, graph=None, env=None,
            executable_path=None):
        self.base = os.path.join(os.path.abspath(base), '')
        if dest is None:
            dest = os.path.join(self.base, 'Contents', 'Frameworks')
        self.dest = dest
        self.mm = FilteredMachOGraph(self, graph=graph, env=env,
            executable_path=executable_path)
        self.changemap = {}
        self.excludes = []
        self.pending = deque()

    def getClass(self, name, cls):
        if in_system_path(name):
            return ExcludedMachO
        for base in self.excludes:
            if name.startswith(base):
                return ExcludedMachO
        return cls

    def locate(self, filename):
        if in_system_path(filename):
            return filename
        if filename.startswith(self.base):
            return filename
        for base in self.excludes:
            if filename.startswith(base):
                return filename
        if filename in self.changemap:
            return self.changemap[filename]
        info = framework_info(filename)
        if info is None:
            res = self.copy_dylib(filename)
            self.changemap[filename] = res
            return res
        else:
            res = self.copy_framework(info)
            self.changemap[filename] = res
            return res

    def copy_dylib(self, filename):
        dest = os.path.join(self.dest, os.path.basename(filename))
        if not os.path.exists(dest):
            self.mergecopy(filename, dest)
        return dest

    def mergecopy(self, src, dest):
        return mergecopy(src, dest)

    def mergetree(self, src, dest):
        return mergetree(src, dest)

    def copy_framework(self, info):
        dest = os.path.join(self.dest, info['shortname'] + '.framework')
        destfn = os.path.join(self.dest, info['name'])
        src = os.path.join(info['location'], info['shortname'] + '.framework')
        if not os.path.exists(dest):
            self.mergetree(src, dest)
            self.pending.append((destfn, iter_platform_files(dest)))
        return destfn

    def run(self, platfiles=None, contents=None):
        mm = self.mm
        if contents is None:
            contents = '@executable_path/..'
        if platfiles is None:
            platfiles = iter_platform_files(self.base)

        for fn in platfiles:
            mm.run_file(fn)

        while self.pending:
            fmwk, files = self.pending.popleft()
            ref = mm.findNode(fmwk)
            for fn in files:
                mm.run_file(fn, caller=ref)

        changemap = {}
        skipcontents = os.path.join(os.path.dirname(self.dest), '')
        machfiles = []

        for node in mm.flatten(has_filename_filter):
            machfiles.append(node)
            dest = os.path.join(contents, node.filename[len(skipcontents):])
            changemap[node.filename] = dest

        def changefunc(path):
            res = mm.locate(path)
            return changemap.get(res)

        for node in machfiles:
            fn = mm.locate(node.filename)
            if fn is None:
                continue
            rewroteAny = False
            for header in node.headers:
                if node.rewriteLoadCommands(changefunc):
                    rewroteAny = True
            if rewroteAny:
                old_mode = flipwritable(fn)
                try:
                    with open(fn, 'rb+') as f:
                        for header in node.headers:
                            f.seek(0)
                            node.write(f)
                        f.seek(0, 2)
                        f.flush()
                finally:
                    flipwritable(fn, old_mode)

        allfiles = [mm.locate(node.filename) for node in machfiles]
        return set(filter(None, allfiles))
<commit_msg>Avoid crash when a library is linked two under two names (real name and a symlink)<commit_after>import os

from macholib.MachOGraph import MachOGraph, MissingMachO
from macholib.util import iter_platform_files, in_system_path, mergecopy, \
    mergetree, flipwritable, has_filename_filter
from macholib.dyld import framework_info
from collections import deque

class ExcludedMachO(MissingMachO):
    pass

class FilteredMachOGraph(MachOGraph):
    def __init__(self, delegate, *args, **kwargs):
        super(FilteredMachOGraph, self).__init__(*args, **kwargs)
        self.delegate = delegate

    def createNode(self, cls, name):
        cls = self.delegate.getClass(name, cls)
        res = super(FilteredMachOGraph, self).createNode(cls, name)
        return res

    def locate(self, filename):
        newname = super(FilteredMachOGraph, self).locate(filename)
        if newname is None:
            return None
        return self.delegate.locate(newname)

class MachOStandalone(object):
    def __init__(self, base, dest=None, graph=None, env=None,
            executable_path=None):
        self.base = os.path.join(os.path.abspath(base), '')
        if dest is None:
            dest = os.path.join(self.base, 'Contents', 'Frameworks')
        self.dest = dest
        self.mm = FilteredMachOGraph(self, graph=graph, env=env,
            executable_path=executable_path)
        self.changemap = {}
        self.excludes = []
        self.pending = deque()

    def getClass(self, name, cls):
        if in_system_path(name):
            return ExcludedMachO
        for base in self.excludes:
            if name.startswith(base):
                return ExcludedMachO
        return cls

    def locate(self, filename):
        if in_system_path(filename):
            return filename
        if filename.startswith(self.base):
            return filename
        for base in self.excludes:
            if filename.startswith(base):
                return filename
        if filename in self.changemap:
            return self.changemap[filename]
        info = framework_info(filename)
        if info is None:
            res = self.copy_dylib(filename)
            self.changemap[filename] = res
            return res
        else:
            res = self.copy_framework(info)
            self.changemap[filename] = res
            return res

    def copy_dylib(self, filename):
        # When the filename is a symlink use the basename of the target of the link
        # as the name in standalone bundle. This avoids problems when two libraries
        # link to the same dylib but using different symlinks.
        if os.path.islink(filename):
            dest = os.path.join(self.dest, os.path.basename(os.path.realpath(filename)))
        else:
            dest = os.path.join(self.dest, os.path.basename(filename))

        if not os.path.exists(dest):
            self.mergecopy(filename, dest)
        return dest

    def mergecopy(self, src, dest):
        return mergecopy(src, dest)

    def mergetree(self, src, dest):
        return mergetree(src, dest)

    def copy_framework(self, info):
        dest = os.path.join(self.dest, info['shortname'] + '.framework')
        destfn = os.path.join(self.dest, info['name'])
        src = os.path.join(info['location'], info['shortname'] + '.framework')
        if not os.path.exists(dest):
            self.mergetree(src, dest)
            self.pending.append((destfn, iter_platform_files(dest)))
        return destfn

    def run(self, platfiles=None, contents=None):
        mm = self.mm
        if contents is None:
            contents = '@executable_path/..'
        if platfiles is None:
            platfiles = iter_platform_files(self.base)

        for fn in platfiles:
            mm.run_file(fn)

        while self.pending:
            fmwk, files = self.pending.popleft()
            ref = mm.findNode(fmwk)
            for fn in files:
                mm.run_file(fn, caller=ref)

        changemap = {}
        skipcontents = os.path.join(os.path.dirname(self.dest), '')
        machfiles = []

        for node in mm.flatten(has_filename_filter):
            machfiles.append(node)
            dest = os.path.join(contents, node.filename[len(skipcontents):])
            changemap[node.filename] = dest

        def changefunc(path):
            res = mm.locate(path)
            return changemap.get(res)

        for node in machfiles:
            fn = mm.locate(node.filename)
            if fn is None:
                continue
            rewroteAny = False
            for header in node.headers:
                if node.rewriteLoadCommands(changefunc):
                    rewroteAny = True
            if rewroteAny:
                old_mode = flipwritable(fn)
                try:
                    with open(fn, 'rb+') as f:
                        for header in node.headers:
                            f.seek(0)
                            node.write(f)
                        f.seek(0, 2)
                        f.flush()
                finally:
                    flipwritable(fn, old_mode)

        allfiles = [mm.locate(node.filename) for node in machfiles]
        return set(filter(None, allfiles))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import (
    division, absolute_import, print_function, unicode_literals,
)

# commented out since __all__ require string in Python 2.
# from builtins import *                  # noqa
# from future.builtins.disabled import *  # noqa

from magic_parameter.decorator import (
    function_parameter,
    method_parameter,
    method_init_parameter,
    class_init_parameter,
)
from magic_parameter.type_declaration import (
    list_t,
    tuple_t,
    set_t,
    dict_t,
    or_t,
)

__all__ = [
    'function_parameter',
    'method_parameter',
    'method_init_parameter',
    'class_init_parameter',

    'list_t',
    'tuple_t',
    'set_t',
    'dict_t',
    'or_t',
]
<commit_msg>Fixed Python2 __all__ issue.<commit_after># -*- coding: utf-8 -*-
from __future__ import (
    division, absolute_import, print_function,
    # commented out since __all__ require str in Python 2.
    # unicode_literals,
)

from builtins import *                  # noqa
from future.builtins.disabled import *  # noqa

from magic_parameter.decorator import (
    function_parameter,
    method_parameter,
    method_init_parameter,
    class_init_parameter,
)
from magic_parameter.type_declaration import (
    list_t,
    tuple_t,
    set_t,
    dict_t,
    or_t,
)

__all__ = [
    'function_parameter',
    'method_parameter',
    'method_init_parameter',
    'class_init_parameter',

    'list_t',
    'tuple_t',
    'set_t',
    'dict_t',
    'or_t',
]
<|endoftext|>


<commit_before>import datetime
import hashlib
import re
import urllib
import json
from gettext import gettext as _
from jinja2 import nodes
from jinja2.ext import Extension
from jinja2.utils import contextfunction, import_string, Markup

from mailpile.commands import Action
from mailpile.util import *
from mailpile.plugins import get_activities, get_selection_actions
from mailpile.plugins import get_display_actions, get_display_modes
from mailpile.plugins import get_assets, get_body_blocks

class MailpileCommand(Extension):
    """Run Mailpile Commands, """
    tags = set(['mpcmd'])

    def __init__(self, environment):
        Extension.__init__(self, environment)
        self.env = environment
        environment.globals['mailpile'] = self._command
        environment.globals['regex_replace'] = self._regex_replace
        environment.filters['regex_replace'] = self._regex_replace
        environment.globals['friendly_bytes'] = self._friendly_bytes
        environment.filters['friendly_bytes'] = self._friendly_bytes
        environment.globals['friendly_number'] = self._friendly_number
        environment.filters['friendly_number'] = self._friendly_number
        environment.globals['show_avatar'] = self._show_avatar
        environment.filters['show_avatar'] = self._show_avatar
        environment.globals['navigation_on'] = self._navigation_on
        environment.filters['navigation_on'] = self._navigation_on
        environment.globals['show_tags'] = self._show_tags
        environment.filters['show_tags'] = self._show_tags
        environment.globals['message_encryption_status'
                            ] = self._message_encryption_status
        environment.filters['message_encryption_status'
                            ] = self._message_encryption_status
        environment.globals['contact_url'] = self._contact_url
        environment.filters['contact_url'] = self._contact_url
        environment.globals['contact_name'] = self._contact_name
        environment.filters['contact_name'] = self._contact_name

        # See utils.py for these functions:
        environment.globals['elapsed_datetime'] = elapsed_datetime
        environment.filters['elapsed_datetime'] = elapsed_datetime
        environment.globals['friendly_datetime'] = friendly_datetime
        environment.filters['friendly_datetime'] = friendly_datetime
        environment.globals['friendly_time'] = friendly_time
        environment.filters['friendly_time'] = friendly_time

        # See plugins/__init__.py for these functions:
        environment.globals['get_activities'] = get_activities
        environment.globals['get_selection_actions'] = get_selection_actions
        environment.globals['get_display_actions'] = get_display_actions
        environment.globals['get_display_modes'] = get_display_modes
        environment.globals['get_assets'] = get_assets
        environment.globals['get_body_blocks'] = get_body_blocks

    def _command(self, command, *args, **kwargs):
        return Action(self.env.session, command, args, data=kwargs).as_dict()

    def _regex_replace(self, s, find, replace):
        """A non-optimal implementation of a regex filter"""
        return re.sub(find, replace, s)

    def _friendly_number(self, number, decimals=0):
        # See mailpile/util.py:friendly_number if this needs fixing
        return friendly_number(number, decimals=decimals, base=1000)

    def _friendly_bytes(self, number, decimals=0):
        # See mailpile/util.py:friendly_number if this needs fixing
        return friendly_number(number,
                               decimals=decimals, base=1024, suffix='B')

    def _show_avatar(self, contact):
        if "photo" in contact:
            photo = contact['photo']
        else:
            photo = '/static/img/avatar-default.png'
        return photo

    def _navigation_on(self, search_tag_ids, on_tid):
        if search_tag_ids:
            for tid in search_tag_ids:
                if tid == on_tid:
                    return "navigation-on"
                else:
                    return ""

    def _show_tags(self, search_terms, tags):
        return ""

    def _message_encryption_status(self, message_parts):
        state = "unencrypted"
        count = 0
        for part in message_parts:
            if 'encryption_info' in part:
                if part["encryption_info"]['status'] == "decrypted":
                    count += 1
                    state = "encrypted"
                elif part["encryption_info"]['status'] == "missingkey":
                    count += 1
                    state = "missingkey"
                elif part["encryption_info"]['status'] == "error":
                    state = "error"
                else:
                    count += 0
        return state

    def _contact_url(self, person):
        if 'contact' in person['flags']:
            url = "/contact/" + person['address'] + "/"
        else:
            url = "/contact/add/" + person['address'] + "/"
        return url

    def _contact_name(self, profiles, person):
        name = person['fn']
        for profile in profiles:
            if profile['email'] == person['address']:
                name = _('You')
                break
        return name
<commit_msg>Added & Modified message_encryption_classes and message_signature_classes<commit_after>import datetime
import hashlib
import re
import urllib
import json
from gettext import gettext as _
from jinja2 import nodes
from jinja2.ext import Extension
from jinja2.utils import contextfunction, import_string, Markup

from mailpile.commands import Action
from mailpile.util import *
from mailpile.plugins import get_activities, get_selection_actions
from mailpile.plugins import get_display_actions, get_display_modes
from mailpile.plugins import get_assets, get_body_blocks

class MailpileCommand(Extension):
    """Run Mailpile Commands, """
    tags = set(['mpcmd'])

    def __init__(self, environment):
        Extension.__init__(self, environment)
        self.env = environment
        environment.globals['mailpile'] = self._command
        environment.globals['regex_replace'] = self._regex_replace
        environment.filters['regex_replace'] = self._regex_replace
        environment.globals['friendly_bytes'] = self._friendly_bytes
        environment.filters['friendly_bytes'] = self._friendly_bytes
        environment.globals['friendly_number'] = self._friendly_number
        environment.filters['friendly_number'] = self._friendly_number
        environment.globals['show_avatar'] = self._show_avatar
        environment.filters['show_avatar'] = self._show_avatar
        environment.globals['navigation_on'] = self._navigation_on
        environment.filters['navigation_on'] = self._navigation_on
        environment.globals['show_tags'] = self._show_tags
        environment.filters['show_tags'] = self._show_tags
        environment.globals['message_signature_classes'
                            ] = self._message_signature_classes
        environment.filters['message_signature_classes'
                            ] = self._message_signature_classes
        environment.globals['message_encryption_classes'
                            ] = self._message_encryption_classes
        environment.filters['message_encryption_classes'
                            ] = self._message_encryption_classes
        environment.globals['contact_url'] = self._contact_url
        environment.filters['contact_url'] = self._contact_url
        environment.globals['contact_name'] = self._contact_name
        environment.filters['contact_name'] = self._contact_name

        # See utils.py for these functions:
        environment.globals['elapsed_datetime'] = elapsed_datetime
        environment.filters['elapsed_datetime'] = elapsed_datetime
        environment.globals['friendly_datetime'] = friendly_datetime
        environment.filters['friendly_datetime'] = friendly_datetime
        environment.globals['friendly_time'] = friendly_time
        environment.filters['friendly_time'] = friendly_time

        # See plugins/__init__.py for these functions:
        environment.globals['get_activities'] = get_activities
        environment.globals['get_selection_actions'] = get_selection_actions
        environment.globals['get_display_actions'] = get_display_actions
        environment.globals['get_display_modes'] = get_display_modes
        environment.globals['get_assets'] = get_assets
        environment.globals['get_body_blocks'] = get_body_blocks

    def _command(self, command, *args, **kwargs):
        return Action(self.env.session, command, args, data=kwargs).as_dict()

    def _regex_replace(self, s, find, replace):
        """A non-optimal implementation of a regex filter"""
        return re.sub(find, replace, s)

    def _friendly_number(self, number, decimals=0):
        # See mailpile/util.py:friendly_number if this needs fixing
        return friendly_number(number, decimals=decimals, base=1000)

    def _friendly_bytes(self, number, decimals=0):
        # See mailpile/util.py:friendly_number if this needs fixing
        return friendly_number(number,
                               decimals=decimals, base=1024, suffix='B')

    def _show_avatar(self, contact):
        if "photo" in contact:
            photo = contact['photo']
        else:
            photo = '/static/img/avatar-default.png'
        return photo

    def _navigation_on(self, search_tag_ids, on_tid):
        if search_tag_ids:
            for tid in search_tag_ids:
                if tid == on_tid:
                    return "navigation-on"
                else:
                    return ""

    def _show_tags(self, search_terms, tags):
        return ""

    def _message_signature_classes(self, status):
        if status == "none":
            state = "icon-signature-none"
        elif status in ("error", "invalid", "revoked"):
            status = "crypto-color-red icon-signature-" + status
        elif status in ("expired", "unknown"):
            state = "crypto-color-orange icon-signature-" + status
        elif status == "unverified":
            state = "crypto-color-blue icon-signature-unverified"
        elif status == "verified":
            state = "crypto-color-green icon-signature-verified"
        elif status.startswith("mixed-"):
            state = "crypto-color-blue icon-signature-unknown"
        else:
            state = "icon-signature-none"
        return state

    def _message_encryption_classes(self, status):
        if status == "none":
            state = "icon-lock-open"
        elif status == "decrypted":
            state = "crypto-color-green icon-lock-closed"
        elif status == "missingkey":
            state = "crypto-color-red icon-lock-closed"
        elif status == "error":
            state = "crypto-color-red icon-lock-error"
        elif status == "partial-decrypted":
            state = "crypto-color-orange icon-lock-open"
        elif status.startswith("mixed-"):
            state = "crypto-color-orange icon-lock-open"
        else:
            state = "icon-lock-open"
        return state

    def _contact_url(self, person):
        if 'contact' in person['flags']:
            url = "/contact/" + person['address'] + "/"
        else:
            url = "/contact/add/" + person['address'] + "/"
        return url

    def _contact_name(self, profiles, person):
        name = person['fn']
        for profile in profiles:
            if profile['email'] == person['address']:
                name = _('You')
                break
        return name
<|endoftext|>


<commit_before>'''Internal utilities.

.. warning::

    For users of lima there should be no need to use anything within
    :mod:`lima.util` directly. Name and contents of this module may change at
    any time without deprecation notice or upgrade path.

'''
from collections import abc
from contextlib import contextmanager

@contextmanager
def complain_about(name):
    '''A Context manager that makes every Exception about name'''
    try:
        yield
    except Exception as e:
        # adapted from http://stackoverflow.com/a/17677938
        msg = '{}: {}'.format(name, e) if e.args else name
        e.args = (msg, ) + e.args[1:]
        raise

class suppress:
    '''Context manager to suppress specified exceptions

    This context manager is taken directly from the Python 3.4 standard library
    to get support for Python 3.3.

    See https://docs.python.org/3.4/library/contextlib.html#contextlib.suppress

    '''
    def __init__(self, *exceptions):
        self._exceptions = exceptions

    def __enter__(self):
        pass

    def __exit__(self, exctype, excinst, exctb):
        return exctype is not None and issubclass(exctype, self._exceptions)

def vector_context(obj):
    '''Return obj if obj is a vector, or [obj] in case obj is a scalar.

    For this function, a *vector* is an iterable that's no string. Everything
    else counts as a *scalar*.

    Inspired by Perl's list context (this has nothing to do with Python context
    managers). Useful to provide scalar values to operations that expect
    vectors (so there's no need to put brackets around single elements).

    Args:
        obj: Any object

    Returns:
        ``obj`` if obj is a vector, otherwise ``[obj]``.

    '''
    if isinstance(obj, abc.Iterable) and not isinstance(obj, str):
        return obj
    return [obj]

def ensure_iterable(obj):
    '''Raise TypeError if obj is not iterable.'''
    if not isinstance(obj, abc.Iterable):
        raise TypeError('{!r} is not iterable.'.format(obj))

def ensure_mapping(obj):
    '''Raise TypeError if obj is no mapping.'''
    if not isinstance(obj, abc.Mapping):
        raise TypeError('{!r} is not a mapping.'.format(obj))

def ensure_only_one_of(collection, elements):
    '''Raise ValueError if collection contains more than one of elements.

    Only distinct elements are considered. For mappings, the keys are
    considered.

    Args:
        collection: An iterable container.

        elements: A set of elements that must not appear together.

    Raises:
        ValueError: If *collection* contains more than one (distinct) element
            of *elements*.

    '''
    found = set(collection) & set(elements)
    if len(found) > 1:
        raise ValueError('Mutually exclusive: {!r}'.format(found))

def ensure_subset_of(collection, superset):
    '''Raise ValueError if collection is no subset of superset

    Only distinct elements are considered. For mappings, only the keys are
    considered.

    Args:
        collection: An iterable container.

        superset: A set of allowed elements.

    Raises:
        ValueError: If *collection* contains more than one (distinct) element
            of *elements*.

    '''
    excess = set(collection) - set(superset)
    if excess:
        raise ValueError('Excess element(s): {!r}'.format(excess))

def ensure_only_instances_of(collection, cls):
    '''Raise TypeError, if collection contains elements not of type cls.'''
    found = [obj for obj in collection if not isinstance(obj, cls)]
    if found:
        raise TypeError('No instances of {}: {!r}'.format(cls, found))

def make_function(name, code, globals_=None):
    '''Return a function created by executing a code string in a new namespace.

    This is not much more than a wrapper around :func:`exec`.

    Args:
        name: The name of the function to create. Must match the function name
            in ``code``.

        code: A String containing the function definition code. The name of the
            function must match ``name``.

        globals_: A dict of globals to mix into the new function's namespace.
            ``__builtins__`` must be provided explicitly if required.

    .. warning:

        All pitfalls of using :func:`exec` apply to this function as well.

    '''
    namespace = dict(__builtins__={})
    if globals_:
        namespace.update(globals_)
    exec(code, namespace)
    return namespace[name]
<commit_msg>utils: refactor docstrings & comments for suppress.<commit_after>'''Internal utilities.

.. warning::

    For users of lima there should be no need to use anything within
    :mod:`lima.util` directly. Name and contents of this module may change at
    any time without deprecation notice or upgrade path.

'''
from collections import abc
from contextlib import contextmanager

@contextmanager
def complain_about(name):
    '''A Context manager that makes every Exception about name'''
    try:
        yield
    except Exception as e:
        # adapted from http://stackoverflow.com/a/17677938
        msg = '{}: {}'.format(name, e) if e.args else name
        e.args = (msg, ) + e.args[1:]
        raise

# The code for this class is taken directly from the Python 3.4 standard
# library (to support Python 3.3), licensed under the PSF License (see
# https://docs.python.org/3/license.html)
class suppress:
    '''Context manager to suppress specified exceptions

    After the exception is suppressed, execution proceeds with the next
    statement following the with statement.

         with suppress(FileNotFoundError):
             os.remove(somefile)
         # Execution still resumes here if the file was already removed

    Backported for Python 3.3 from Python 3.4 (see source for license info).

    '''
    def __init__(self, *exceptions):
        self._exceptions = exceptions

    def __enter__(self):
        pass

    def __exit__(self, exctype, excinst, exctb):
        # Unlike isinstance and issubclass, CPython exception handling
        # currently only looks at the concrete type hierarchy (ignoring
        # the instance and subclass checking hooks). While Guido considers
        # that a bug rather than a feature, it's a fairly hard one to fix
        # due to various internal implementation details. suppress provides
        # the simpler issubclass based semantics, rather than trying to
        # exactly reproduce the limitations of the CPython interpreter.
        #
        # See http://bugs.python.org/issue12029 for more details
        return exctype is not None and issubclass(exctype, self._exceptions)

def vector_context(obj):
    '''Return obj if obj is a vector, or [obj] in case obj is a scalar.

    For this function, a *vector* is an iterable that's no string. Everything
    else counts as a *scalar*.

    Inspired by Perl's list context (this has nothing to do with Python context
    managers). Useful to provide scalar values to operations that expect
    vectors (so there's no need to put brackets around single elements).

    Args:
        obj: Any object

    Returns:
        ``obj`` if obj is a vector, otherwise ``[obj]``.

    '''
    if isinstance(obj, abc.Iterable) and not isinstance(obj, str):
        return obj
    return [obj]

def ensure_iterable(obj):
    '''Raise TypeError if obj is not iterable.'''
    if not isinstance(obj, abc.Iterable):
        raise TypeError('{!r} is not iterable.'.format(obj))

def ensure_mapping(obj):
    '''Raise TypeError if obj is no mapping.'''
    if not isinstance(obj, abc.Mapping):
        raise TypeError('{!r} is not a mapping.'.format(obj))

def ensure_only_one_of(collection, elements):
    '''Raise ValueError if collection contains more than one of elements.

    Only distinct elements are considered. For mappings, the keys are
    considered.

    Args:
        collection: An iterable container.

        elements: A set of elements that must not appear together.

    Raises:
        ValueError: If *collection* contains more than one (distinct) element
            of *elements*.

    '''
    found = set(collection) & set(elements)
    if len(found) > 1:
        raise ValueError('Mutually exclusive: {!r}'.format(found))

def ensure_subset_of(collection, superset):
    '''Raise ValueError if collection is no subset of superset

    Only distinct elements are considered. For mappings, only the keys are
    considered.

    Args:
        collection: An iterable container.

        superset: A set of allowed elements.

    Raises:
        ValueError: If *collection* contains more than one (distinct) element
            of *elements*.

    '''
    excess = set(collection) - set(superset)
    if excess:
        raise ValueError('Excess element(s): {!r}'.format(excess))

def ensure_only_instances_of(collection, cls):
    '''Raise TypeError, if collection contains elements not of type cls.'''
    found = [obj for obj in collection if not isinstance(obj, cls)]
    if found:
        raise TypeError('No instances of {}: {!r}'.format(cls, found))

def make_function(name, code, globals_=None):
    '''Return a function created by executing a code string in a new namespace.

    This is not much more than a wrapper around :func:`exec`.

    Args:
        name: The name of the function to create. Must match the function name
            in ``code``.

        code: A String containing the function definition code. The name of the
            function must match ``name``.

        globals_: A dict of globals to mix into the new function's namespace.
            ``__builtins__`` must be provided explicitly if required.

    .. warning:

        All pitfalls of using :func:`exec` apply to this function as well.

    '''
    namespace = dict(__builtins__={})
    if globals_:
        namespace.update(globals_)
    exec(code, namespace)
    return namespace[name]
<|endoftext|>


<commit_before>"""Module containing the data for various IGM metal lines.
Species covered are:
H, He,C, N, O, Ne, Mg, Si, Fe
Reads the data from VPFIT's atom.dat file
Transitions are various.

"""

import re
import os.path

class LineData(object):
    """Class to aggregate a number of lines from VPFIT tables. Reads from atom.dat,
    to get the results call get_line(species, ion), where ion is the transition number.
    At the moment only lowest level transition in each sequence is read."""
    def __init__(self, vpdat = None):
        self.species = ('H', 'He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'Fe')
        #Put in the masses by hand for simplicity
        self.masses = {'H': 1.00794,'He': 4.002602,'C': 12.011,'N': 14.00674,'O': 15.9994,'Ne': 20.18,'Mg': 24.3050,'Si': 28.0855,'Fe': 55.847 }
        #9 empty lists, one list per species
        self.lines = {}
        if vpdat is None:
            vpdata = os.path.join(os.path.dirname(os.path.realpath(__file__)), "atom.dat")
        self.read_vpfit(vpdata)

    def read_vpfit(self,vpdat):
        """Read VPFIT's atom.dat file for the 9 species"""
        vp = open(vpdat)
        inline = "--"
        while inline != "":
            inline = vp.readline()
            try:
                (specie, ion) = find_species(inline)
            except ValueError:
                continue
            #is this a line we care about?
            if self.species.count(specie) == 0 or ion < 0:
                continue
            (lambda_X, fosc_X, gamma_X) = parse_line_contents(inline)
            line = Line(lambda_X, fosc_X, gamma_X)
            #Only add the first transition for lines
            if (specie,ion) not in self.lines:
                self.lines[(specie,ion)]={int(lambda_X):line}
            else:
                self.lines[(specie,ion)][int(lambda_X)] = line
#                 print "Read line: ",specie,ion

    def __getitem__(self,specion):
        """Get data for a particular line.
        specion: a tuple of (specie, ion)
        specie: number of species
        ion: transition number (from 1)"""
        lines = self.lines[specion]
        return lines

    def __len__(self):
        return len(self.lines)

    def get_mass(self, specie):
        """Get the mass in amu for a species"""
        return self.masses[specie]

class Line(object):
    """Class to store the parameters of a single line.
    Data is:
        lambda_X - line width in angstroms (10^-10 m)
        fosc_X - oscillator strength (dimensionless)
        gamma_X - line's gamma in Hz"""
    def __init__(self, lambda_X, fosc_X, gamma_X):
        self.lambda_X = lambda_X
        self.fosc_X = fosc_X
        self.gamma_X = gamma_X

def find_species(line):
    """Parse a line to find which species it is: species is given by first two characters, unless the second character is a capital.
    There may be whitespace.
    Ionisation number is then a capital letter followed by three characters."""
    #Match a capital and possibly a lower case, followed by a capital followed by any three characters.
    mat = re.match(r"([A-Z]\s*[a-z]?)([A-Z].{3})",line)
    if mat is None:
        raise ValueError
    species = re.sub(r"\s","",mat.groups()[0])
    ion = mat.groups()[1]
    ion = re.sub(r"\s","",mat.groups()[1])
    ion_r = roman_to_int(ion)
    return (species, ion_r)

def parse_line_contents(line):
    """Extract the sigma, gamma, lambda and m from the line:
       just read until we get something looking like a float and then store the first three of them"""
    #Split by spaces
    spl = line.split()
    res = []
    for ss in spl:
        try:
            float(ss)
            res.append(float(ss))
        except ValueError:
            pass
        if len(res) == 3:
            break
    return res[0], res[1], res[2]

def roman_to_int(roman):
    """
    Convert a roman numeral to an integer.

    >>> r = range(1, 4000)
    >>> nums = [int_to_roman(i) for i in r]
    >>> ints = [roman_to_int(n) for n in nums]
    >>> print r == ints
    1

    >>> roman_to_int('VVVIV')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: VVVIV
    >>> roman_to_int(1)
    Traceback (most recent call last):
     ...
    TypeError: expected string, got <type 'int'>
    >>> roman_to_int('a')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: A
    >>> roman_to_int('IL')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: IL

    Thanks Paul Winkler, on the internet.
    """
    if not isinstance(roman, str):
        raise TypeError("expected string, got %s" % type(roman))
    #roman = roman.upper()
    nums = ['M', 'D', 'C', 'L', 'X', 'V', 'I']
    ints = [1000, 500, 100, 50,  10,  5,   1]
    places = []
    for c in roman:
        if not c in nums:
            raise ValueError("roman is not a valid roman numeral: %s" % roman)
    for i in range(len(roman)):
        c = roman[i]
        value = ints[nums.index(c)]
        # If the next place holds a larger number, this value is negative.
        try:
            nextvalue = ints[nums.index(roman[i +1])]
            if nextvalue > value:
                value *= -1
        except IndexError:
            # there is no next place.
            pass
        places.append(value)
    tot = 0
    for n in places:
        tot += n
    return tot
<commit_msg>Add line data cache so we don't have to reparse atom.dat all the time<commit_after>"""Module containing the data for various IGM metal lines.
Species covered are:
H, He,C, N, O, Ne, Mg, Si, Fe
Reads the data from VPFIT's atom.dat file
Transitions are various.

"""

import re
import os.path

_lines_cache_ = None

class LineData(object):
    """Class to aggregate a number of lines from VPFIT tables. Reads from atom.dat,
    to get the results call get_line(species, ion), where ion is the transition number.
    At the moment only lowest level transition in each sequence is read."""
    def __init__(self, vpdat = None):
        self.species = ('H', 'He', 'C', 'N', 'O', 'Ne', 'Mg', 'Si', 'Fe')
        #Put in the masses by hand for simplicity
        self.masses = {'H': 1.00794,'He': 4.002602,'C': 12.011,'N': 14.00674,'O': 15.9994,'Ne': 20.18,'Mg': 24.3050,'Si': 28.0855,'Fe': 55.847 }
        #9 empty lists, one list per species
        self.lines = {}
        if vpdat is None:
            vpdata = os.path.join(os.path.dirname(os.path.realpath(__file__)), "atom.dat")
            if _lines_cache_ is None:
                _lines_cache_ = read_vpfit(vpdata,self.species)
            self.lines = _lines_cache_
        else:
            self.lines = read_vpfit(vpdata,self.species)

    def __getitem__(self,specion):
        """Get data for a particular line.
        specion: a tuple of (specie, ion)
        specie: number of species
        ion: transition number (from 1)"""
        lines = self.lines[specion]
        return lines

    def __len__(self):
        return len(self.lines)

    def get_mass(self, specie):
        """Get the mass in amu for a species"""
        return self.masses[specie]

class Line(object):
    """Class to store the parameters of a single line.
    Data is:
        lambda_X - line width in angstroms (10^-10 m)
        fosc_X - oscillator strength (dimensionless)
        gamma_X - line's gamma in Hz"""
    def __init__(self, lambda_X, fosc_X, gamma_X):
        self.lambda_X = lambda_X
        self.fosc_X = fosc_X
        self.gamma_X = gamma_X

def read_vpfit(vpdat, species):
    """Read VPFIT's atom.dat file for the 9 species"""
    vp = open(vpdat)
    lines = {}
    inline = "--"
    while inline != "":
        inline = vp.readline()
        try:
            (specie, ion) = find_species(inline)
        except ValueError:
            continue
        #is this a line we care about?
        if species.count(specie) == 0 or ion < 0:
            continue
        (lambda_X, fosc_X, gamma_X) = parse_line_contents(inline)
        line = Line(lambda_X, fosc_X, gamma_X)
        #Only add the first transition for lines
        if (specie,ion) not in lines:
            lines[(specie,ion)]={int(lambda_X):line}
        else:
            lines[(specie,ion)][int(lambda_X)] = line
#                 print "Read line: ",specie,ion
    return lines

def find_species(line):
    """Parse a line to find which species it is: species is given by first two characters, unless the second character is a capital.
    There may be whitespace.
    Ionisation number is then a capital letter followed by three characters."""
    #Match a capital and possibly a lower case, followed by a capital followed by any three characters.
    mat = re.match(r"([A-Z]\s*[a-z]?)([A-Z].{3})",line)
    if mat is None:
        raise ValueError
    species = re.sub(r"\s","",mat.groups()[0])
    ion = mat.groups()[1]
    ion = re.sub(r"\s","",mat.groups()[1])
    ion_r = roman_to_int(ion)
    return (species, ion_r)

def parse_line_contents(line):
    """Extract the sigma, gamma, lambda and m from the line:
       just read until we get something looking like a float and then store the first three of them"""
    #Split by spaces
    spl = line.split()
    res = []
    for ss in spl:
        try:
            float(ss)
            res.append(float(ss))
        except ValueError:
            pass
        if len(res) == 3:
            break
    return res[0], res[1], res[2]

def roman_to_int(roman):
    """
    Convert a roman numeral to an integer.

    >>> r = range(1, 4000)
    >>> nums = [int_to_roman(i) for i in r]
    >>> ints = [roman_to_int(n) for n in nums]
    >>> print r == ints
    1

    >>> roman_to_int('VVVIV')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: VVVIV
    >>> roman_to_int(1)
    Traceback (most recent call last):
     ...
    TypeError: expected string, got <type 'int'>
    >>> roman_to_int('a')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: A
    >>> roman_to_int('IL')
    Traceback (most recent call last):
     ...
    ValueError: roman is not a valid roman numeral: IL

    Thanks Paul Winkler, on the internet.
    """
    if not isinstance(roman, str):
        raise TypeError("expected string, got %s" % type(roman))
    #roman = roman.upper()
    nums = ['M', 'D', 'C', 'L', 'X', 'V', 'I']
    ints = [1000, 500, 100, 50,  10,  5,   1]
    places = []
    for c in roman:
        if not c in nums:
            raise ValueError("roman is not a valid roman numeral: %s" % roman)
    for i in range(len(roman)):
        c = roman[i]
        value = ints[nums.index(c)]
        # If the next place holds a larger number, this value is negative.
        try:
            nextvalue = ints[nums.index(roman[i +1])]
            if nextvalue > value:
                value *= -1
        except IndexError:
            # there is no next place.
            pass
        places.append(value)
    tot = 0
    for n in places:
        tot += n
    return tot
<|endoftext|>


<commit_before>import os
import shutil
import tempfile
import subprocess

def merge_user_settings(settings):
	default = settings.get('default') or {}
	user = settings.get('user') or {}
	if user:
		plugins = default.pop('plugins', {})
		user_plugins = user.get('plugins', {})
		for name, data in user_plugins.items():
			if name in plugins:
				plugins[name].update(data)
			else:
				plugins[name] = data
		default['plugins'] = plugins

		user.pop('plugins', None)
		default.update(user)

	return default

def memoize(f):
	rets = {}

	def wrap(*args):
		if not args in rets:
			rets[args] = f(*args)

		return rets[args]

	wrap.__name__ = f.__name__
	return wrap

def climb(top):
    right = True
    while right:
        top, right = os.path.split(top)
        yield top

@memoize
def find(top, name, parent=False):
    for d in climb(top):
        target = os.path.join(d, name)
        if os.path.exists(target):
            if parent:
                return d

            return target

def extract_path(cmd, delim=':'):
	path = popen(cmd, os.environ).communicate()[0].decode()
	path = path.split('__SUBL__', 1)[1].strip('\r\n')
	return ':'.join(path.split(delim))

def find_path(env):
	# find PATH using shell --login
	if 'SHELL' in env:
		shell_path = env['SHELL']
		shell = os.path.basename(shell_path)

		if shell in ('bash', 'zsh'):
			return extract_path(
				(shell_path, '--login', '-c', 'echo __SUBL__$PATH')
			)
		elif shell == 'fish':
			return extract_path(
				(shell_path, '--login', '-c', 'echo __SUBL__; for p in $PATH; echo $p; end'),
				'\n'
			)

	# guess PATH if we haven't returned yet
	split = env['PATH'].split(':')
	p = env['PATH']
	for path in (
		'/usr/bin', '/usr/local/bin',
		'/usr/local/php/bin', '/usr/local/php5/bin'
				):
		if not path in split:
			p += (':' + path)

	return p

@memoize
def create_environment():
	if os.name == 'posix':
		os.environ['PATH'] = find_path(os.environ)

	return os.environ

def can_exec(fpath):
	return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

def which(cmd):
	env = create_environment()
	for base in env.get('PATH', '').split(os.pathsep):
		path = os.path.join(base, cmd)
		if can_exec(path):
			return path

	return None

# popen methods
def communicate(cmd, code):
	code = code.encode('utf8')
	out = popen(cmd)
	if out is not None:
		out = out.communicate(code)
		return (out[0].decode() or '') + (out[1].decode() or '')
	else:
		return ''

def tmpfile(cmd, code, suffix=''):
	if isinstance(cmd, str):
		cmd = cmd,

	f = tempfile.NamedTemporaryFile(suffix=suffix)
	f.write(code.encode('utf8'))
	f.flush()

	cmd = tuple(cmd) + (f.name,)
	out = popen(cmd)
	if out:
		out = out.communicate()
		return (out[0].decode('utf8', 'replace') or '') + (out[1].decode('utf8', 'replace') or '')
	else:
		return ''

def tmpdir(cmd, files, filename, code):
	filename = os.path.split(filename)[1]
	d = tempfile.mkdtemp()

	for f in files:
		try: os.makedirs(os.path.join(d, os.path.split(f)[0]))
		except: pass

		target = os.path.join(d, f)
		if os.path.split(target)[1] == filename:
			# source file hasn't been saved since change, so update it from our live buffer
			f = open(target, 'wb')
			f.write(code)
			f.close()
		else:
			shutil.copyfile(f, target)

	os.chdir(d)
	out = popen(cmd)
	if out:
		out = out.communicate()
		out = (out[0].decode('utf8') or '') + '\n' + (out[1].decode('utf8') or '')

		# filter results from build to just this filename
		# no guarantee all languages are as nice about this as Go
		# may need to improve later or just defer to communicate()
		out = '\n'.join([
			line for line in out.split('\n') if filename in line.split(':', 1)[0]
		])
	else:
		out = ''

	shutil.rmtree(d, True)
	return out

def popen(cmd, env=None):
	import lint.persist as persist
	if isinstance(cmd, str):
		cmd = cmd,

	info = None
	if os.name == 'nt':
		info = subprocess.STARTUPINFO()
		info.dwFlags |= subprocess.STARTF_USESHOWWINDOW
		info.wShowWindow = subprocess.SW_HIDE

	if env is None:
		env = create_environment()

	try:
		return subprocess.Popen(cmd, stdin=subprocess.PIPE,
			stdout=subprocess.PIPE, stderr=subprocess.PIPE,
			startupinfo=info, env=env)
	except OSError as err:
		persist.debug('SublimeLint: Error launching', repr(cmd))
		persist.debug('Error was:', err.strerror)
		persist.debug('Environment:', env)
<commit_msg>consolidate cli output combining code<commit_after>import os
import shutil
import tempfile
import subprocess

def merge_user_settings(settings):
	default = settings.get('default') or {}
	user = settings.get('user') or {}
	if user:
		plugins = default.pop('plugins', {})
		user_plugins = user.get('plugins', {})
		for name, data in user_plugins.items():
			if name in plugins:
				plugins[name].update(data)
			else:
				plugins[name] = data
		default['plugins'] = plugins

		user.pop('plugins', None)
		default.update(user)

	return default

def memoize(f):
	rets = {}

	def wrap(*args):
		if not args in rets:
			rets[args] = f(*args)

		return rets[args]

	wrap.__name__ = f.__name__
	return wrap

def climb(top):
    right = True
    while right:
        top, right = os.path.split(top)
        yield top

@memoize
def find(top, name, parent=False):
    for d in climb(top):
        target = os.path.join(d, name)
        if os.path.exists(target):
            if parent:
                return d

            return target

def extract_path(cmd, delim=':'):
	path = popen(cmd, os.environ).communicate()[0].decode()
	path = path.split('__SUBL__', 1)[1].strip('\r\n')
	return ':'.join(path.split(delim))

def find_path(env):
	# find PATH using shell --login
	if 'SHELL' in env:
		shell_path = env['SHELL']
		shell = os.path.basename(shell_path)

		if shell in ('bash', 'zsh'):
			return extract_path(
				(shell_path, '--login', '-c', 'echo __SUBL__$PATH')
			)
		elif shell == 'fish':
			return extract_path(
				(shell_path, '--login', '-c', 'echo __SUBL__; for p in $PATH; echo $p; end'),
				'\n'
			)

	# guess PATH if we haven't returned yet
	split = env['PATH'].split(':')
	p = env['PATH']
	for path in (
		'/usr/bin', '/usr/local/bin',
		'/usr/local/php/bin', '/usr/local/php5/bin'
				):
		if not path in split:
			p += (':' + path)

	return p

@memoize
def create_environment():
	if os.name == 'posix':
		os.environ['PATH'] = find_path(os.environ)

	return os.environ

def can_exec(fpath):
	return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

def which(cmd):
	env = create_environment()
	for base in env.get('PATH', '').split(os.pathsep):
		path = os.path.join(base, cmd)
		if can_exec(path):
			return path

	return None

# popen methods
def combine_output(out, sep=''):
	return sep.join((
		(out[0].decode('utf8') or ''),
		(out[1].decode('utf8') or ''),
	))

def communicate(cmd, code):
	code = code.encode('utf8')
	out = popen(cmd)
	if out is not None:
		out = out.communicate(code)
		return combine_output(out)
	else:
		return ''

def tmpfile(cmd, code, suffix=''):
	if isinstance(cmd, str):
		cmd = cmd,

	f = tempfile.NamedTemporaryFile(suffix=suffix)
	f.write(code.encode('utf8'))
	f.flush()

	cmd = tuple(cmd) + (f.name,)
	out = popen(cmd)
	if out:
		out = out.communicate()
		return combine_output(out)
	else:
		return ''

def tmpdir(cmd, files, filename, code):
	filename = os.path.split(filename)[1]
	d = tempfile.mkdtemp()

	for f in files:
		try: os.makedirs(os.path.join(d, os.path.split(f)[0]))
		except: pass

		target = os.path.join(d, f)
		if os.path.split(target)[1] == filename:
			# source file hasn't been saved since change, so update it from our live buffer
			f = open(target, 'wb')
			f.write(code)
			f.close()
		else:
			shutil.copyfile(f, target)

	os.chdir(d)
	out = popen(cmd)
	if out:
		out = out.communicate()
		out = combine_output(out, '\n')

		# filter results from build to just this filename
		# no guarantee all languages are as nice about this as Go
		# may need to improve later or just defer to communicate()
		out = '\n'.join([
			line for line in out.split('\n') if filename in line.split(':', 1)[0]
		])
	else:
		out = ''

	shutil.rmtree(d, True)
	return out

def popen(cmd, env=None):
	import lint.persist as persist
	if isinstance(cmd, str):
		cmd = cmd,

	info = None
	if os.name == 'nt':
		info = subprocess.STARTUPINFO()
		info.dwFlags |= subprocess.STARTF_USESHOWWINDOW
		info.wShowWindow = subprocess.SW_HIDE

	if env is None:
		env = create_environment()

	try:
		return subprocess.Popen(cmd, stdin=subprocess.PIPE,
			stdout=subprocess.PIPE, stderr=subprocess.PIPE,
			startupinfo=info, env=env)
	except OSError as err:
		persist.debug('SublimeLint: Error launching', repr(cmd))
		persist.debug('Error was:', err.strerror)
		persist.debug('Environment:', env)
<|endoftext|>


<commit_before>import sys

from .loader import Loader
from .parser import Parser, Context, Argument
from .exceptions import Failure, CollectionNotFound
from ._version import __version__

def parse(argv):
    # Initial/core parsing (core options can affect the rest of the parsing)
    initial_context = Context(args=(
        # TODO: make collection a list-building arg, not a string
        Argument(names=('collection', 'c')),
        Argument(names=('root', 'r')),
        Argument(names=('version', 'V'), kind=bool, default=False),
    ))
    # 'core' will result an .unparsed attribute with what was left over.
    parser = Parser(initial=initial_context, ignore_unknown=True)
    core = parser.parse_argv(argv)
    args = core[0].args

    # Print version & exit if necessary
    if args.version.value:
        print "Invoke %s" % __version__
        sys.exit(0)

    # Load collection (default or specified) and parse leftovers
    loader = Loader(root=args.root.value)
    collection = loader.load_collection(args.collection.value)
    tasks = Parser(contexts=collection.to_contexts()).parse_argv(core.unparsed)

    return args, collection, tasks

def main():
    # Parse command line
    args, collection, tasks = parse(sys.argv[1:])

    # Take action based on 'core' options and the 'tasks' found
    for context in tasks:
        kwargs = {}
        for name, arg in context.args.iteritems():
            kwargs[name] = arg.value
        try:
            collection[context.name].body(**kwargs)
        except Failure, f:
            sys.exit(f.result.exited)
<commit_msg>Don't traceback during parse errors<commit_after>import sys

from .loader import Loader
from .parser import Parser, Context, Argument
from .exceptions import Failure, CollectionNotFound, ParseError
from ._version import __version__

def parse_gracefully(parser, argv):
    """
    Run ``parser.parse_argv(argv)`` & gracefully handle ``ParseError``s.

    'Gracefully' meaning to print a useful human-facing error message instead
    of a traceback; the program will still exit if an error is raised.

    If no error is raised, returns the result of the ``parse_argv`` call.
    """
    try:
        return parser.parse_argv(argv)
    except ParseError, e:
        print str(e)
        sys.exit(1)

def parse(argv):
    # Initial/core parsing (core options can affect the rest of the parsing)
    initial_context = Context(args=(
        # TODO: make collection a list-building arg, not a string
        Argument(names=('collection', 'c')),
        Argument(names=('root', 'r')),
        Argument(names=('version', 'V'), kind=bool, default=False),
    ))
    # 'core' will result an .unparsed attribute with what was left over.
    parser = Parser(initial=initial_context, ignore_unknown=True)
    core = parse_gracefully(parser, argv)
    args = core[0].args

    # Print version & exit if necessary
    if args.version.value:
        print "Invoke %s" % __version__
        sys.exit(0)

    # Load collection (default or specified) and parse leftovers
    loader = Loader(root=args.root.value)
    collection = loader.load_collection(args.collection.value)
    parser = Parser(contexts=collection.to_contexts())
    tasks = parse_gracefully(parser, core.unparsed)

    return args, collection, tasks

def main():
    # Parse command line
    args, collection, tasks = parse(sys.argv[1:])

    # Take action based on 'core' options and the 'tasks' found
    for context in tasks:
        kwargs = {}
        for name, arg in context.args.iteritems():
            kwargs[name] = arg.value
        try:
            collection[context.name].body(**kwargs)
        except Failure, f:
            sys.exit(f.result.exited)
<|endoftext|>


<commit_before>import multiprocessing
from multiprocessing import Pool
from margaritashotgun import remote_host
from margaritashotgun import logger
import logging

class Workers():
    """
    """
    cpu_count = None
    worker_count = None
    progress_bar = True
    hosts = None

    def __init__(self, conf, workers, name, library=True):
        """
        """
        self.name = name
        self.library = library
        self.progressbar = True
        self.cpu_count = multiprocessing.cpu_count()
        host_count = len(conf)
        self.worker_count = self.count(workers, self.cpu_count, host_count)
        self.queue = multiprocessing.Queue(-1)

        try:
            log_dir = conf[0]['logging']['log_dir']
        except KeyError:
            log_dir = ""
        try:
            log_prefix = conf[0]['logging']['log_prefix']
        except KeyError:
            log_prefix = ""

        self.log_file = "{}{}memory_capture.log".format(
                            log_dir, log_prefix)

        if self.worker_count > 1 or self.library is True:
            self.progressbar = False
        self.conf = []
        for c in conf:
            c['host']['progressbar'] = self.progressbar
            self.conf.append(c)

    def count(self, workers, cpu_count, host_count):
        """
        """
        if workers == 'auto':
            if cpu_count > host_count:
                worker_count = host_count
            else:
                worker_count = cpu_count
        elif workers > host_count:
            worker_count = host_count
        else:
            worker_count = int(workers)
        return worker_count

    def spawn(self, desc, timeout=1800):
        """
        """
        self.pool = Pool(self.worker_count, initializer=remote_host._init,
                        initargs=(self.queue,))

        self.listener = logger.Logger(target=logger.listener,
                                      args=(self.queue,
                                            self.name,
                                            self.log_file,
                                            desc))
        self.listener.start()
        res = self.pool.map_async(remote_host.process, self.conf)
        results = res.get(timeout)

        self.cleanup()
        return results

    def cleanup(self):
        self.pool.close()
        self.pool.join()
        self.queue.put_nowait(None)
        self.queue.close()
        self.listener.join()
<commit_msg>correct log prefix<commit_after>import multiprocessing
from multiprocessing import Pool
from margaritashotgun import remote_host
from margaritashotgun import logger
import logging

class Workers():
    """
    """
    cpu_count = None
    worker_count = None
    progress_bar = True
    hosts = None

    def __init__(self, conf, workers, name, library=True):
        """
        """
        self.name = name
        self.library = library
        self.progressbar = True
        self.cpu_count = multiprocessing.cpu_count()
        host_count = len(conf)
        self.worker_count = self.count(workers, self.cpu_count, host_count)
        self.queue = multiprocessing.Queue(-1)

        try:
            log_dir = conf[0]['logging']['log_dir']
        except KeyError:
            log_dir = ""
        try:
            log_prefix = conf[0]['logging']['prefix']
        except KeyError:
            log_prefix = ""

        self.log_file = "{}{}memory_capture.log".format(
                            log_dir, log_prefix)

        if self.worker_count > 1 or self.library is True:
            self.progressbar = False
        self.conf = []
        for c in conf:
            c['host']['progressbar'] = self.progressbar
            self.conf.append(c)

    def count(self, workers, cpu_count, host_count):
        """
        """
        if workers == 'auto':
            if cpu_count > host_count:
                worker_count = host_count
            else:
                worker_count = cpu_count
        elif workers > host_count:
            worker_count = host_count
        else:
            worker_count = int(workers)
        return worker_count

    def spawn(self, desc, timeout=1800):
        """
        """
        self.pool = Pool(self.worker_count, initializer=remote_host._init,
                        initargs=(self.queue,))

        self.listener = logger.Logger(target=logger.listener,
                                      args=(self.queue,
                                            self.name,
                                            self.log_file,
                                            desc))
        self.listener.start()
        res = self.pool.map_async(remote_host.process, self.conf)
        results = res.get(timeout)

        self.cleanup()
        return results

    def cleanup(self):
        self.pool.close()
        self.pool.join()
        self.queue.put_nowait(None)
        self.queue.close()
        self.listener.join()
<|endoftext|>