<commit_before>import os, traceback, sys, json, shutil
import subprocess32 as subprocess
from threading import Timer
from contextlib import contextmanager

WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
LIBS_DIR = os.path.join(WORKING_DIR, "libs")
CORPUS_DIR = os.path.join(WORKING_DIR, "corpus")
CORPUS_INFO = None
TOOLS_DIR = os.path.join(WORKING_DIR, "tools")
SIMPROG_DIR = os.path.join(WORKING_DIR, "simprog")

CLUSTER_FILE = os.path.join(WORKING_DIR, "clusters.json")
CLASS2FIELDS_FILE = os.path.join(WORKING_DIR, "c2f.json")

DLJC_BINARY = os.path.join(TOOLS_DIR, "do-like-javac", "dljc")
DLJC_OUTPUT_DIR = "dljc-out"

DYNTRACE_ADDONS_DIR = os.path.join(WORKING_DIR, "dyntrace")

LIMITED_PROJECT_LIST = ["dyn4j", "jreactphysics3d", "jbox2d", "react", "jmonkeyengine"]

def run_cmd(cmd, print_output=False, timeout=None):
  def kill_proc(proc, stats):
    stats['timed_out'] = True
    proc.kill()

  stats = {'timed_out': False,
           'output': ''}
  timer = None

  if print_output:
    print ("Running %s" % ' '.join(cmd))
  try:
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

    if timeout:
      timer = Timer(timeout, kill_proc, [process, stats])
      timer.start()

    for line in iter(process.stdout.readline, b''):
      stats['output'] = stats['output'] + line
      if print_output:
        sys.stdout.write(line)
        sys.stdout.flush()
    process.stdout.close()
    process.wait()
    stats['return_code'] = process.returncode
    if timer:
      timer.cancel()

  except:
    print ('calling {cmd} failed\n{trace}'.format(cmd=' '.join(cmd),trace=traceback.format_exc()))
  return stats

@contextmanager
def cd(newdir):
  prevdir = os.getcwd()
  os.chdir(os.path.expanduser(newdir))
  try:
    yield
  finally:
    os.chdir(prevdir)

def mkdir(newdir):
  if not os.path.isdir(newdir):
    os.makedirs(newdir)

def get_method_from_daikon_out(daikon_out):
  arr1 = daikon_out.split('.')
  arr2 = arr1[1].split(':::')
  method = arr2[0]
  return method

def get_jar(jar_name):
  path = os.path.join(LIBS_DIR, jar_name)
  if os.path.isfile(path):
    return path
  else:
    return None

def get_corpus_info():
  global CORPUS_INFO
  if not CORPUS_INFO:
    with open(os.path.join(WORKING_DIR, 'corpus.json')) as f:
      CORPUS_INFO = json.loads(f.read())

  return CORPUS_INFO

def get_project_dir(project_name):
  project = project_info(project_name)

  if 'build-dir' in project:
    return os.path.join(CORPUS_DIR, project['name'], project['build-dir'])
  else:
    return os.path.join(CORPUS_DIR, project['name'])

def get_project_list():
  return get_corpus_info()['projects'].keys()

def project_info(project_name):
  return get_corpus_info()['projects'][project_name]

def get_simprog(py_file):
  return os.path.join(SIMPROG_DIR, py_file)

def get_dljc_dir_for_project(project_name):
  dtrace_path = os.path.join(get_project_dir(project_name), DLJC_OUTPUT_DIR)
  if os.path.exists(dtrace_path):
    return dtrace_path
  else:
    return None

def clean_project(project_name):
  print "Cleaning {}".format(project_name)
  project_dir = get_project_dir(project_name)
  with cd(project_dir):
    clean_command = project_info(project_name)['clean'].strip().split()
    run_cmd(clean_command)
    run_cmd(['rm', '-r', DLJC_OUTPUT_DIR])

def get_class_dirs(project_name):
  classdirs = []

  dljc_output = os.path.join(get_project_dir(project_name),
                             DLJC_OUTPUT_DIR,
                             'javac.json')

  if not os.path.exists(dljc_output):
    print 'Tried to get classdirs from project where DLJC has not been run.'
    return None

  with open(dljc_output, 'r') as f:
    javac_commands = json.loads(f.read())
    for command in javac_commands:
      classdir = command['javac_switches'].get('d')
      if classdir:
        classdirs.append(classdir)

  return classdirs

def copy_dyntrace_files(project_name):
  project_dir = get_project_dir(project_name)
  out_dir = os.path.join(project_dir, DLJC_OUTPUT_DIR)
  mkdir(out_dir)
  addons = [f for f in os.listdir(DYNTRACE_ADDONS_DIR) \
            if f.startswith(project_info(project_name)['name']+".")]

  for addon in addons:
    addon_type = addon.rsplit('.', 1)[-1]
    shutil.copyfile(os.path.join(DYNTRACE_ADDONS_DIR, addon),
                    os.path.join(out_dir, addon_type))

def run_dljc(project_name, tools, options=[], timelimit=1800.0):
  copy_dyntrace_files(project_name)
  project_dir = get_project_dir(project_name)
  with cd(project_dir):
    build_command = project_info(project_name)['build'].strip().split()
    dljc_command = [DLJC_BINARY,
                    '-l', LIBS_DIR,
                    '-o', DLJC_OUTPUT_DIR,
                    '-t', ','.join(tools)]
    dljc_command.extend(options)
    dljc_command.append('--')
    dljc_command.extend(build_command)
    run_cmd(dljc_command, print_output=True, timeout=timelimit)

CHECKER_ENV_SETUP = False
def setup_checker_framework_env():
  global CHECKER_ENV_SETUP
  if CHECKER_ENV_SETUP:
    return

  jsr308 = TOOLS_DIR
  os.environ['JSR308'] = jsr308

  afu = os.path.join(jsr308, 'annotation-tools', 'annotation-file-utilities')
  os.environ['AFU'] = afu
  os.environ['PATH'] += ':' + os.path.join(afu, 'scripts')
  CHECKER_ENV_SETUP = True
<commit_msg>Add DAIKONDIR to DLJC environment to enable DynComp<commit_after>import os, traceback, sys, json, shutil
import subprocess32 as subprocess
from threading import Timer
from contextlib import contextmanager

WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
LIBS_DIR = os.path.join(WORKING_DIR, "libs")
CORPUS_DIR = os.path.join(WORKING_DIR, "corpus")
CORPUS_INFO = None
TOOLS_DIR = os.path.join(WORKING_DIR, "tools")
SIMPROG_DIR = os.path.join(WORKING_DIR, "simprog")

CLUSTER_FILE = os.path.join(WORKING_DIR, "clusters.json")
CLASS2FIELDS_FILE = os.path.join(WORKING_DIR, "c2f.json")

DLJC_BINARY = os.path.join(TOOLS_DIR, "do-like-javac", "dljc")
DLJC_OUTPUT_DIR = "dljc-out"

DYNTRACE_ADDONS_DIR = os.path.join(WORKING_DIR, "dyntrace")

LIMITED_PROJECT_LIST = ["dyn4j", "jreactphysics3d", "jbox2d", "react", "jmonkeyengine"]

def run_cmd(cmd, print_output=False, timeout=None):
  def kill_proc(proc, stats):
    stats['timed_out'] = True
    proc.kill()

  stats = {'timed_out': False,
           'output': ''}
  timer = None

  if print_output:
    print ("Running %s" % ' '.join(cmd))
  try:
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)

    if timeout:
      timer = Timer(timeout, kill_proc, [process, stats])
      timer.start()

    for line in iter(process.stdout.readline, b''):
      stats['output'] = stats['output'] + line
      if print_output:
        sys.stdout.write(line)
        sys.stdout.flush()
    process.stdout.close()
    process.wait()
    stats['return_code'] = process.returncode
    if timer:
      timer.cancel()

  except:
    print ('calling {cmd} failed\n{trace}'.format(cmd=' '.join(cmd),trace=traceback.format_exc()))
  return stats

@contextmanager
def cd(newdir):
  prevdir = os.getcwd()
  os.chdir(os.path.expanduser(newdir))
  try:
    yield
  finally:
    os.chdir(prevdir)

def mkdir(newdir):
  if not os.path.isdir(newdir):
    os.makedirs(newdir)

def get_method_from_daikon_out(daikon_out):
  arr1 = daikon_out.split('.')
  arr2 = arr1[1].split(':::')
  method = arr2[0]
  return method

def get_jar(jar_name):
  path = os.path.join(LIBS_DIR, jar_name)
  if os.path.isfile(path):
    return path
  else:
    return None

def get_corpus_info():
  global CORPUS_INFO
  if not CORPUS_INFO:
    with open(os.path.join(WORKING_DIR, 'corpus.json')) as f:
      CORPUS_INFO = json.loads(f.read())

  return CORPUS_INFO

def get_project_dir(project_name):
  project = project_info(project_name)

  if 'build-dir' in project:
    return os.path.join(CORPUS_DIR, project['name'], project['build-dir'])
  else:
    return os.path.join(CORPUS_DIR, project['name'])

def get_project_list():
  return get_corpus_info()['projects'].keys()

def project_info(project_name):
  return get_corpus_info()['projects'][project_name]

def get_simprog(py_file):
  return os.path.join(SIMPROG_DIR, py_file)

def get_dljc_dir_for_project(project_name):
  dtrace_path = os.path.join(get_project_dir(project_name), DLJC_OUTPUT_DIR)
  if os.path.exists(dtrace_path):
    return dtrace_path
  else:
    return None

def clean_project(project_name):
  print "Cleaning {}".format(project_name)
  project_dir = get_project_dir(project_name)
  with cd(project_dir):
    clean_command = project_info(project_name)['clean'].strip().split()
    run_cmd(clean_command)
    run_cmd(['rm', '-r', DLJC_OUTPUT_DIR])

def get_class_dirs(project_name):
  classdirs = []

  dljc_output = os.path.join(get_project_dir(project_name),
                             DLJC_OUTPUT_DIR,
                             'javac.json')

  if not os.path.exists(dljc_output):
    print 'Tried to get classdirs from project where DLJC has not been run.'
    return None

  with open(dljc_output, 'r') as f:
    javac_commands = json.loads(f.read())
    for command in javac_commands:
      classdir = command['javac_switches'].get('d')
      if classdir:
        classdirs.append(classdir)

  return classdirs

def copy_dyntrace_files(project_name):
  project_dir = get_project_dir(project_name)
  out_dir = os.path.join(project_dir, DLJC_OUTPUT_DIR)
  mkdir(out_dir)
  addons = [f for f in os.listdir(DYNTRACE_ADDONS_DIR) \
            if f.startswith(project_info(project_name)['name']+".")]

  for addon in addons:
    addon_type = addon.rsplit('.', 1)[-1]
    shutil.copyfile(os.path.join(DYNTRACE_ADDONS_DIR, addon),
                    os.path.join(out_dir, addon_type))

def run_dljc(project_name, tools, options=[], timelimit=1800.0):
  copy_dyntrace_files(project_name)
  os.environ['DAIKONDIR'] = os.path.join(TOOLS_DIR, 'daikon-src')
  project_dir = get_project_dir(project_name)
  with cd(project_dir):
    build_command = project_info(project_name)['build'].strip().split()
    dljc_command = [DLJC_BINARY,
                    '-l', LIBS_DIR,
                    '-o', DLJC_OUTPUT_DIR,
                    '-t', ','.join(tools)]
    dljc_command.extend(options)
    dljc_command.append('--')
    dljc_command.extend(build_command)
    run_cmd(dljc_command, print_output=True, timeout=timelimit)

CHECKER_ENV_SETUP = False
def setup_checker_framework_env():
  global CHECKER_ENV_SETUP
  if CHECKER_ENV_SETUP:
    return

  jsr308 = TOOLS_DIR
  os.environ['JSR308'] = jsr308

  afu = os.path.join(jsr308, 'annotation-tools', 'annotation-file-utilities')
  os.environ['AFU'] = afu
  os.environ['PATH'] += ':' + os.path.join(afu, 'scripts')
  CHECKER_ENV_SETUP = True
<|endoftext|>


<commit_before>"""
sentry.tsdb.redis
~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2014 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import

import logging
import six

from binascii import crc32
from collections import defaultdict
from datetime import timedelta
from django.conf import settings
from django.utils import timezone
from hashlib import md5

from sentry.tsdb.base import BaseTSDB
from sentry.utils.dates import to_timestamp
from sentry.utils.redis import (
    check_cluster_versions,
    make_rb_cluster,
)
from sentry.utils.versioning import Version

logger = logging.getLogger(__name__)

class RedisTSDB(BaseTSDB):
    """
    A time series storage backend for Redis.

    The time series API supports two data types:

        * simple counters
        * distinct counters (number of unique elements seen)

    The backend also supports virtual nodes (``vnodes``) which controls shard
    distribution. This value should be set to the anticipated maximum number of
    physical hosts and not modified after data has been written.

    Simple counters are stored in hashes. The key of the hash is composed of
    the model, epoch (which defines the start of the rollup period), and a
    shard identifier. This allows TTLs to be applied to the entire bucket,
    instead of having to be stored for every individual element in the rollup
    period. This results in a data layout that looks something like this::

        {
            "<model>:<epoch>:<shard id>": {
                "<key>": value,
                ...
            },
            ...
        }

    Distinct counters are stored using HyperLogLog, which provides a
    cardinality estimate with a standard error of 0.8%. The data layout looks
    something like this::

        {
            "<model>:<epoch>:<key>": value,
            ...
        }

    """
    def __init__(self, hosts=None, prefix='ts:', vnodes=64, **kwargs):
        # inherit default options from REDIS_OPTIONS
        defaults = settings.SENTRY_REDIS_OPTIONS

        if hosts is None:
            hosts = defaults.get('hosts', {0: {}})

        self.cluster = make_rb_cluster(hosts)
        self.prefix = prefix
        self.vnodes = vnodes
        super(RedisTSDB, self).__init__(**kwargs)

    def validate(self):
        logger.info('Validating Redis version...')
        check_cluster_versions(
            self.cluster,
            Version((2, 8, 9)),
            label='TSDB',
        )

    def make_key(self, model, epoch, model_key):
        if isinstance(model_key, six.integer_types):
            vnode = model_key % self.vnodes
        else:
            vnode = crc32(model_key) % self.vnodes

        return '{0}{1}:{2}:{3}'.format(self.prefix, model.value, epoch, vnode)

    def get_model_key(self, key):
        # We specialize integers so that a pure int-map can be optimized by
        # Redis, whereas long strings (say tag values) will store in a more
        # efficient hashed format.
        if not isinstance(key, six.integer_types):
            # enforce utf-8 encoding
            if isinstance(key, unicode):
                key = key.encode('utf-8')
            return md5(repr(key)).hexdigest()
        return key

    def incr(self, model, key, timestamp=None, count=1):
        self.incr_multi([(model, key)], timestamp, count)

    def incr_multi(self, items, timestamp=None, count=1):
        """
        Increment project ID=1 and group ID=5:

        >>> incr_multi([(TimeSeriesModel.project, 1), (TimeSeriesModel.group, 5)])
        """
        make_key = self.make_key
        normalize_to_rollup = self.normalize_to_rollup
        if timestamp is None:
            timestamp = timezone.now()

        with self.cluster.map() as client:
            for rollup, max_values in self.rollups:
                norm_rollup = normalize_to_rollup(timestamp, rollup)
                for model, key in items:
                    model_key = self.get_model_key(key)
                    hash_key = make_key(model, norm_rollup, model_key)
                    client.hincrby(hash_key, model_key, count)
                    client.expireat(
                        hash_key,
                        self.calculate_expiry(rollup, max_values, timestamp),
                    )

    def get_range(self, model, keys, start, end, rollup=None):
        """
        To get a range of data for group ID=[1, 2, 3]:

        Start and end are both inclusive.

        >>> now = timezone.now()
        >>> get_keys(TimeSeriesModel.group, [1, 2, 3],
        >>>          start=now - timedelta(days=1),
        >>>          end=now)
        """
        normalize_to_epoch = self.normalize_to_epoch
        normalize_to_rollup = self.normalize_to_rollup
        make_key = self.make_key

        if rollup is None:
            rollup = self.get_optimal_rollup(start, end)

        results = []
        timestamp = end
        with self.cluster.map() as client:
            while timestamp >= start:
                real_epoch = normalize_to_epoch(timestamp, rollup)
                norm_epoch = normalize_to_rollup(timestamp, rollup)

                for key in keys:
                    model_key = self.get_model_key(key)
                    hash_key = make_key(model, norm_epoch, model_key)
                    results.append((real_epoch, key,
                                    client.hget(hash_key, model_key)))

                timestamp = timestamp - timedelta(seconds=rollup)

        results_by_key = defaultdict(dict)
        for epoch, key, count in results:
            results_by_key[key][epoch] = int(count.value or 0)

        for key, points in results_by_key.iteritems():
            results_by_key[key] = sorted(points.items())
        return dict(results_by_key)

    def make_distinct_counter_key(self, model, rollup, timestamp, key):
        return '{prefix}{model}:{epoch}:{key}'.format(
            prefix=self.prefix,
            model=model.value,
            epoch=self.normalize_ts_to_rollup(timestamp, rollup),
            key=self.get_model_key(key),
        )

    def record(self, model, key, values, timestamp=None):
        self.record_multi(((model, key, values),), timestamp)

    def record_multi(self, items, timestamp=None):
        """
        Record an occurence of an item in a distinct counter.
        """
        if timestamp is None:
            timestamp = timezone.now()

        ts = int(to_timestamp(timestamp))  # ``timestamp`` is not actually a timestamp :(

        with self.cluster.fanout() as client:
            for model, key, values in items:
                c = client.target_key(key)
                for rollup, max_values in self.rollups:
                    k = self.make_distinct_counter_key(
                        model,
                        rollup,
                        ts,
                        key,
                    )
                    c.pfadd(k, *values)
                    c.expireat(
                        k,
                        self.calculate_expiry(
                            rollup,
                            max_values,
                            timestamp,
                        ),
                    )

    def get_distinct_counts_series(self, model, keys, start, end=None, rollup=None):
        """
        Fetch counts of distinct items for each rollup interval within the range.
        """
        rollup, series = self.get_optimal_rollup_series(start, end, rollup)

        responses = {}
        with self.cluster.fanout() as client:
            for key in keys:
                c = client.target_key(key)
                r = responses[key] = []
                for timestamp in series:
                    r.append((
                        timestamp,
                        c.pfcount(
                            self.make_distinct_counter_key(
                                model,
                                rollup,
                                timestamp,
                                key,
                            ),
                        ),
                    ))

        return {key: [(timestamp, promise.value) for timestamp, promise in value] for key, value in responses.iteritems()}

    def get_distinct_counts_totals(self, model, keys, start, end=None, rollup=None):
        """
        Count distinct items during a time range.
        """
        rollup, series = self.get_optimal_rollup_series(start, end, rollup)

        responses = {}
        with self.cluster.fanout() as client:
            for key in keys:
                # XXX: The current versions of the Redis driver don't implement
                # ``PFCOUNT`` correctly (although this is fixed in the Git
                # master, so should be available in the next release) and only
                # supports a single key argument -- not the variadic signature
                # supported by the protocol -- so we have to call the commnand
                # directly here instead.
                ks = []
                for timestamp in series:
                    ks.append(self.make_distinct_counter_key(model, rollup, timestamp, key))

                responses[key] = client.target_key(key).execute_command('PFCOUNT', *ks)

        return {key: value.value for key, value in responses.iteritems()}
<commit_msg>Clean up TSDB Redis backend instantiation.<commit_after>"""
sentry.tsdb.redis
~~~~~~~~~~~~~~~~~

:copyright: (c) 2010-2014 by the Sentry Team, see AUTHORS for more details.
:license: BSD, see LICENSE for more details.
"""
from __future__ import absolute_import

import logging
import six

from binascii import crc32
from collections import defaultdict
from datetime import timedelta
from django.conf import settings
from django.utils import timezone
from hashlib import md5

from sentry.tsdb.base import BaseTSDB
from sentry.utils.dates import to_timestamp
from sentry.utils.redis import (
    check_cluster_versions,
    make_rb_cluster,
)
from sentry.utils.versioning import Version

logger = logging.getLogger(__name__)

class RedisTSDB(BaseTSDB):
    """
    A time series storage backend for Redis.

    The time series API supports two data types:

        * simple counters
        * distinct counters (number of unique elements seen)

    The backend also supports virtual nodes (``vnodes``) which controls shard
    distribution. This value should be set to the anticipated maximum number of
    physical hosts and not modified after data has been written.

    Simple counters are stored in hashes. The key of the hash is composed of
    the model, epoch (which defines the start of the rollup period), and a
    shard identifier. This allows TTLs to be applied to the entire bucket,
    instead of having to be stored for every individual element in the rollup
    period. This results in a data layout that looks something like this::

        {
            "<model>:<epoch>:<shard id>": {
                "<key>": value,
                ...
            },
            ...
        }

    Distinct counters are stored using HyperLogLog, which provides a
    cardinality estimate with a standard error of 0.8%. The data layout looks
    something like this::

        {
            "<model>:<epoch>:<key>": value,
            ...
        }

    """
    def __init__(self, hosts=None, prefix='ts:', vnodes=64, **kwargs):
        # inherit default options from REDIS_OPTIONS
        defaults = settings.SENTRY_REDIS_OPTIONS

        if hosts is None:
            hosts = defaults.get('hosts', {0: {}})

        self.cluster = make_rb_cluster(hosts)
        self.prefix = prefix
        self.vnodes = vnodes
        super(RedisTSDB, self).__init__(**kwargs)

    def validate(self):
        logger.debug('Validating Redis version...')
        check_cluster_versions(
            self.cluster,
            Version((2, 8, 9)),
            label='TSDB',
        )

    def make_key(self, model, epoch, model_key):
        if isinstance(model_key, six.integer_types):
            vnode = model_key % self.vnodes
        else:
            vnode = crc32(model_key) % self.vnodes

        return '{0}{1}:{2}:{3}'.format(self.prefix, model.value, epoch, vnode)

    def get_model_key(self, key):
        # We specialize integers so that a pure int-map can be optimized by
        # Redis, whereas long strings (say tag values) will store in a more
        # efficient hashed format.
        if not isinstance(key, six.integer_types):
            # enforce utf-8 encoding
            if isinstance(key, unicode):
                key = key.encode('utf-8')
            return md5(repr(key)).hexdigest()
        return key

    def incr(self, model, key, timestamp=None, count=1):
        self.incr_multi([(model, key)], timestamp, count)

    def incr_multi(self, items, timestamp=None, count=1):
        """
        Increment project ID=1 and group ID=5:

        >>> incr_multi([(TimeSeriesModel.project, 1), (TimeSeriesModel.group, 5)])
        """
        make_key = self.make_key
        normalize_to_rollup = self.normalize_to_rollup
        if timestamp is None:
            timestamp = timezone.now()

        with self.cluster.map() as client:
            for rollup, max_values in self.rollups:
                norm_rollup = normalize_to_rollup(timestamp, rollup)
                for model, key in items:
                    model_key = self.get_model_key(key)
                    hash_key = make_key(model, norm_rollup, model_key)
                    client.hincrby(hash_key, model_key, count)
                    client.expireat(
                        hash_key,
                        self.calculate_expiry(rollup, max_values, timestamp),
                    )

    def get_range(self, model, keys, start, end, rollup=None):
        """
        To get a range of data for group ID=[1, 2, 3]:

        Start and end are both inclusive.

        >>> now = timezone.now()
        >>> get_keys(TimeSeriesModel.group, [1, 2, 3],
        >>>          start=now - timedelta(days=1),
        >>>          end=now)
        """
        normalize_to_epoch = self.normalize_to_epoch
        normalize_to_rollup = self.normalize_to_rollup
        make_key = self.make_key

        if rollup is None:
            rollup = self.get_optimal_rollup(start, end)

        results = []
        timestamp = end
        with self.cluster.map() as client:
            while timestamp >= start:
                real_epoch = normalize_to_epoch(timestamp, rollup)
                norm_epoch = normalize_to_rollup(timestamp, rollup)

                for key in keys:
                    model_key = self.get_model_key(key)
                    hash_key = make_key(model, norm_epoch, model_key)
                    results.append((real_epoch, key,
                                    client.hget(hash_key, model_key)))

                timestamp = timestamp - timedelta(seconds=rollup)

        results_by_key = defaultdict(dict)
        for epoch, key, count in results:
            results_by_key[key][epoch] = int(count.value or 0)

        for key, points in results_by_key.iteritems():
            results_by_key[key] = sorted(points.items())
        return dict(results_by_key)

    def make_distinct_counter_key(self, model, rollup, timestamp, key):
        return '{prefix}{model}:{epoch}:{key}'.format(
            prefix=self.prefix,
            model=model.value,
            epoch=self.normalize_ts_to_rollup(timestamp, rollup),
            key=self.get_model_key(key),
        )

    def record(self, model, key, values, timestamp=None):
        self.record_multi(((model, key, values),), timestamp)

    def record_multi(self, items, timestamp=None):
        """
        Record an occurence of an item in a distinct counter.
        """
        if timestamp is None:
            timestamp = timezone.now()

        ts = int(to_timestamp(timestamp))  # ``timestamp`` is not actually a timestamp :(

        with self.cluster.fanout() as client:
            for model, key, values in items:
                c = client.target_key(key)
                for rollup, max_values in self.rollups:
                    k = self.make_distinct_counter_key(
                        model,
                        rollup,
                        ts,
                        key,
                    )
                    c.pfadd(k, *values)
                    c.expireat(
                        k,
                        self.calculate_expiry(
                            rollup,
                            max_values,
                            timestamp,
                        ),
                    )

    def get_distinct_counts_series(self, model, keys, start, end=None, rollup=None):
        """
        Fetch counts of distinct items for each rollup interval within the range.
        """
        rollup, series = self.get_optimal_rollup_series(start, end, rollup)

        responses = {}
        with self.cluster.fanout() as client:
            for key in keys:
                c = client.target_key(key)
                r = responses[key] = []
                for timestamp in series:
                    r.append((
                        timestamp,
                        c.pfcount(
                            self.make_distinct_counter_key(
                                model,
                                rollup,
                                timestamp,
                                key,
                            ),
                        ),
                    ))

        return {key: [(timestamp, promise.value) for timestamp, promise in value] for key, value in responses.iteritems()}

    def get_distinct_counts_totals(self, model, keys, start, end=None, rollup=None):
        """
        Count distinct items during a time range.
        """
        rollup, series = self.get_optimal_rollup_series(start, end, rollup)

        responses = {}
        with self.cluster.fanout() as client:
            for key in keys:
                # XXX: The current versions of the Redis driver don't implement
                # ``PFCOUNT`` correctly (although this is fixed in the Git
                # master, so should be available in the next release) and only
                # supports a single key argument -- not the variadic signature
                # supported by the protocol -- so we have to call the commnand
                # directly here instead.
                ks = []
                for timestamp in series:
                    ks.append(self.make_distinct_counter_key(model, rollup, timestamp, key))

                responses[key] = client.target_key(key).execute_command('PFCOUNT', *ks)

        return {key: value.value for key, value in responses.iteritems()}
<|endoftext|>


<commit_before>import os, subprocess, traceback, sys
from contextlib import contextmanager

WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
LIBS_DIR = os.path.join(WORKING_DIR, "libs")
CORPUS_DIR = os.path.join(WORKING_DIR, "corpus")
TOOLS_DIR = os.path.join(WORKING_DIR, "tools")

DLJC_BINARY = os.path.join(TOOLS_DIR, "do-like-javac", "dljc")
DLJC_OUTPUT_DIR = "dljc-out"

def jar_path(jar):
  path = os.path.join(LIBS_DIR, jar)
  if os.path.isfile(path):
    return path
  else:
    return None

def run_cmd(cmd, print_output=False):
  output = ""
  if print_output:
    print ("Running %s" % ' '.join(cmd))
  try:
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if print_output:
      for line in iter(process.stdout.readline, b''):
        output = output + line
        sys.stdout.write(line)
        sys.stdout.flush()
      process.stdout.close()
    process.wait()
  except:
    print ('calling {cmd} failed\n{trace}'.format(cmd=' '.join(cmd),trace=traceback.format_exc()))
  return output

@contextmanager
def cd(newdir):
  prevdir = os.getcwd()
  os.chdir(os.path.expanduser(newdir))
  try:
    yield
  finally:
    os.chdir(prevdir)

def get_project_dir(project_name):
  return os.path.join(CORPUS_DIR, project_name)

def get_project_list():
  return [project for project in  os.listdir(corpus_dir) if os.path.isdir(get_project_dir(project))]

def clean_project(project):
  project_dir = get_project_dir(project)
  with cd(project_dir):
    with open(os.path.join(project_dir, 'clean_command.txt'), 'r') as f:
      clean_command = f.readline().strip().split()
      run_cmd(clean_command)

def run_dljc(project, tools, options=[]):
  project_dir = get_project_dir(project)
  with cd(project_dir):
    with open(os.path.join(project_dir, 'build_command.txt'), 'r') as f:
      build_command = f.readline().strip().split()
      dljc_command = [DLJC_BINARY,
                      '-o', DLJC_OUTPUT_DIR,
                      '-t', ','.join(tools)]
      dljc_command.extend(options)
      dljc_command.append('--')
      dljc_command.extend(build_command)
      run_cmd(dljc_command, print_output=True)
<commit_msg>capitalize corpus_dir<commit_after>import os, subprocess, traceback, sys
from contextlib import contextmanager

WORKING_DIR = os.path.dirname(os.path.realpath(__file__))
LIBS_DIR = os.path.join(WORKING_DIR, "libs")
CORPUS_DIR = os.path.join(WORKING_DIR, "corpus")
TOOLS_DIR = os.path.join(WORKING_DIR, "tools")

DLJC_BINARY = os.path.join(TOOLS_DIR, "do-like-javac", "dljc")
DLJC_OUTPUT_DIR = "dljc-out"

def jar_path(jar):
  path = os.path.join(LIBS_DIR, jar)
  if os.path.isfile(path):
    return path
  else:
    return None

def run_cmd(cmd, print_output=False):
  output = ""
  if print_output:
    print ("Running %s" % ' '.join(cmd))
  try:
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if print_output:
      for line in iter(process.stdout.readline, b''):
        output = output + line
        sys.stdout.write(line)
        sys.stdout.flush()
      process.stdout.close()
    process.wait()
  except:
    print ('calling {cmd} failed\n{trace}'.format(cmd=' '.join(cmd),trace=traceback.format_exc()))
  return output

@contextmanager
def cd(newdir):
  prevdir = os.getcwd()
  os.chdir(os.path.expanduser(newdir))
  try:
    yield
  finally:
    os.chdir(prevdir)

def get_project_dir(project_name):
  return os.path.join(CORPUS_DIR, project_name)

def get_project_list():
  return [project for project in  os.listdir(CORPUS_DIR) if os.path.isdir(get_project_dir(project))]

def clean_project(project):
  project_dir = get_project_dir(project)
  with cd(project_dir):
    with open(os.path.join(project_dir, 'clean_command.txt'), 'r') as f:
      clean_command = f.readline().strip().split()
      run_cmd(clean_command)

def run_dljc(project, tools, options=[]):
  project_dir = get_project_dir(project)
  with cd(project_dir):
    with open(os.path.join(project_dir, 'build_command.txt'), 'r') as f:
      build_command = f.readline().strip().split()
      dljc_command = [DLJC_BINARY,
                      '-o', DLJC_OUTPUT_DIR,
                      '-t', ','.join(tools)]
      dljc_command.extend(options)
      dljc_command.append('--')
      dljc_command.extend(build_command)
      run_cmd(dljc_command, print_output=True)
<|endoftext|>


<commit_before>'''Business days calculations and utilities for a given calendar specification.
The calendar specification is a .cal file containing the weekdays to be
considered as non-business days and a iso-formated list of dates representing
holidays. It follows an example::

    Saturday
    Sunday
    2001-01-01
    2002-01-01
    2013-01-01

Classes:
    Calendar
    Date
    DateIndex
'''

import os
from datetime import datetime, date, timedelta
import re

class DateIndex(object):
    def __init__(self, holidays, startdate, enddate, weekdays):
        self._index = {}
        self.startdate = Date(startdate).date
        self.enddate = Date(enddate).date
        self.weekdays = weekdays
        self.holidays = [d.date for d in holidays]
        d1 = timedelta(1)
        dt = self.startdate
        w = c = 1
        while dt <= self.enddate:
            is_hol = dt in self.holidays or dt.weekday() in weekdays
            self._index[dt] = (w, c, is_hol)
            c += 1
            if not is_hol:
                w += 1
            dt = dt + d1

    def get(self, dt):
        dt = Date(dt)
        return self._index[dt.date]

    def __getitem__(self, dt):
        return self.get(dt)

class Date(object):
    def __init__(self, d=None, format='%Y-%m-%d'):
        d = d if d else date.today()
        if type(d) in (str, unicode):
            d = datetime.strptime(d, format).date()
        elif type(d) is datetime:
            d = d.date()
        elif type(d) is Date:
            d = d.date
        elif type(d) is date:
            pass
        else:
            raise ValueError()
        self.date = d

    def format(self, fmts='%Y-%m-%d'):
        return datetime.strftime(self.date, fmts)

    def __cmp__(self, other):
        return cmp(self.date, other.date)

    def __eq__(self, other):
        return self.date == other.date

    def __repr__(self):
        return self.format()

    __str__ = __repr__

class Calendar(object):
    ''' Calendar class to compute business days accordingly a list of holidays.
    '''
    _weekdays = ('Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday')
    def __init__(self, holidays=[], startdate='1970-01-01', enddate='2071-01-01', weekdays=[], name=None,
                       adjust_from='next', adjust_to='previous'):
        self.name = name
        self._holidays = [Date(d) for d in holidays]
        self._nonwork_weekdays = [self._weekdays.index(wd) for wd in weekdays]
        if len(self._holidays):
            self._startdate = min(self._holidays)
            self._enddate = max(self._holidays)
        else:
            self._startdate = Date(startdate)
            self._enddate = Date(enddate)
        self._index = DateIndex(self._holidays, self._startdate, self._enddate, self._nonwork_weekdays)

    def __get_weekdays(self):
        return tuple( self._weekdays[nwd] for nwd in self._nonwork_weekdays )
    weekdays = property(__get_weekdays)

    def __get_startdate(self):
        return self._startdate.date
    startdate = property(__get_startdate)

    def __get_enddate(self):
        return self._enddate.date
    enddate = property(__get_enddate)

    def __get_holidays(self):
        return [d.date for d in self._holidays]
    holidays = property(__get_holidays)

    def __get_index(self):
        return self._index
    index = property(__get_index)

    def bizdays(self, date_from, date_to):
        '''Returns the number of business days between 2 dates.
        dates is a tuple with 2 strings ISO-formated dates (%%Y-%%m-%%d).'''
        d1 = self.__adjust_next(date_from)
        d2 = self.__adjust_previous(date_to)
        if d1 > d2:
            raise ValueError("The first date must be before the second.")
        return self._index[d2][0] - self._index[d1][0]

    def isbizday(self, dt):
        '''Returns True whether dt is a business day or False otherwise.
        dt is a string ISO-formated date (%%Y-%%m-%%d).'''
        return not self._index[dt][2]

    def __adjust_next(self, dt):
        d1 = timedelta(1)
        dt = Date(dt).date
        while self._index[dt][2]:
            dt = dt + d1
        return Date(dt)

    def adjust_next(self, dt):
        """Returns the next business day whether the passed date isn't a
        business day or returns the given date.
        dt is a string ISO-formated date (%%Y-%%m-%%d).
        """
        return self.__adjust_next(dt).date

    def __adjust_previous(self, dt):
        d1 = timedelta(1)
        dt = Date(dt).date
        while self._index[dt][2]:
            dt = dt - d1
        return Date(dt)

    def adjust_previous(self, dt):
        """Returns the first business day before the passed date whether
        the given date isn't a business day or returns the given date.
        dt is a string ISO-formated date (%%Y-%%m-%%d).
        """
        return self.__adjust_previous(dt).date

    def seq(self, date_from, date_to):
        '''Returns a sequence generator which generates business days between
        the 2 given dates.
        dates is a tuple with 2 strings ISO-formated dates (%%Y-%%m-%%d).
        '''
        _from, _to = Date(date_from), Date(date_to)
        d1 = timedelta(1)
        if _from > _to:
            raise ValueError("The first date must be before the second.")
        _from = self.__adjust_next(_from)
        while _from <= _to:
            yield _from.date
            _from = self.__adjust_next(_from.date + d1)

    def offset(self, dt, n):
        """
        Offsets the given date by n days
        """
        dt = Date(dt)
        if n >= 0:
            d1 = timedelta(1)
            adjust = lambda d: self.__adjust_next(d)
            dt = adjust(dt)
        else:
            d1 = timedelta(-1)
            adjust = lambda d: self.__adjust_previous(d)
            dt = adjust(dt)
            n = abs(n)
        i = 0
        while i < n:
            dt.date += d1
            dt = adjust(dt)
            i += 1
        return dt.date

    @classmethod
    def load(cls, fname):
        if not os.path.exists(fname):
            raise Exception('Invalid calendar specification: \
            file not found (%s)' % fname)
        w = '|'.join(cls._weekdays)
        _holidays = []
        _nonwork_weekdays = []
        with open(fname) as fcal:
            for cal_reg in fcal:
                cal_reg = cal_reg.strip()
                if cal_reg is '': continue
                if re.match('^%s$' % w, cal_reg):
                    _nonwork_weekdays.append(cal_reg)
                else:
                    _holidays.append(Date(cal_reg))
        return Calendar(_holidays, weekdays=_nonwork_weekdays)

    @staticmethod
    def load_holidays(fname):
        if not os.path.exists(fname):
            raise Exception('Invalid calendar specification: \
            file not found (%s)' % fname)
        _holidays = []
        with open(fname) as fcal:
            for cal_reg in fcal:
                cal_reg = cal_reg.strip()
                if cal_reg is '': continue
                _holidays.append(Date(cal_reg).date)
        return _holidays

    def __eq__(self, other):
        return self.startdate == other.startdate and \
            self.enddate == other.enddate and \
            self._cal_spec == other._cal_spec

<commit_msg>Changed arguments order<commit_after>'''Business days calculations and utilities for a given calendar specification.
The calendar specification is a .cal file containing the weekdays to be
considered as non-business days and a iso-formated list of dates representing
holidays. It follows an example::

    Saturday
    Sunday
    2001-01-01
    2002-01-01
    2013-01-01

Classes:
    Calendar
    Date
    DateIndex
'''

import os
from datetime import datetime, date, timedelta
import re

class DateIndex(object):
    def __init__(self, holidays, startdate, enddate, weekdays):
        self._index = {}
        self.startdate = Date(startdate).date
        self.enddate = Date(enddate).date
        self.weekdays = weekdays
        self.holidays = [d.date for d in holidays]
        d1 = timedelta(1)
        dt = self.startdate
        w = c = 1
        while dt <= self.enddate:
            is_hol = dt in self.holidays or dt.weekday() in weekdays
            self._index[dt] = (w, c, is_hol)
            c += 1
            if not is_hol:
                w += 1
            dt = dt + d1

    def get(self, dt):
        dt = Date(dt)
        return self._index[dt.date]

    def __getitem__(self, dt):
        return self.get(dt)

class Date(object):
    def __init__(self, d=None, format='%Y-%m-%d'):
        d = d if d else date.today()
        if type(d) in (str, unicode):
            d = datetime.strptime(d, format).date()
        elif type(d) is datetime:
            d = d.date()
        elif type(d) is Date:
            d = d.date
        elif type(d) is date:
            pass
        else:
            raise ValueError()
        self.date = d

    def format(self, fmts='%Y-%m-%d'):
        return datetime.strftime(self.date, fmts)

    def __cmp__(self, other):
        return cmp(self.date, other.date)

    def __eq__(self, other):
        return self.date == other.date

    def __repr__(self):
        return self.format()

    __str__ = __repr__

class Calendar(object):
    ''' Calendar class to compute business days accordingly a list of holidays.
    '''
    _weekdays = ('Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday')
    def __init__(self, holidays=[], weekdays=[], startdate='1970-01-01', enddate='2071-01-01', name=None,
                       adjust_from='next', adjust_to='previous'):
        self.name = name
        self._holidays = [Date(d) for d in holidays]
        self._nonwork_weekdays = [self._weekdays.index(wd) for wd in weekdays]
        if len(self._holidays):
            self._startdate = min(self._holidays)
            self._enddate = max(self._holidays)
        else:
            self._startdate = Date(startdate)
            self._enddate = Date(enddate)
        self._index = DateIndex(self._holidays, self._startdate, self._enddate, self._nonwork_weekdays)

    def __get_weekdays(self):
        return tuple( self._weekdays[nwd] for nwd in self._nonwork_weekdays )
    weekdays = property(__get_weekdays)

    def __get_startdate(self):
        return self._startdate.date
    startdate = property(__get_startdate)

    def __get_enddate(self):
        return self._enddate.date
    enddate = property(__get_enddate)

    def __get_holidays(self):
        return [d.date for d in self._holidays]
    holidays = property(__get_holidays)

    def __get_index(self):
        return self._index
    index = property(__get_index)

    def bizdays(self, date_from, date_to):
        '''Returns the number of business days between 2 dates.
        dates is a tuple with 2 strings ISO-formated dates (%%Y-%%m-%%d).'''
        d1 = self.__adjust_next(date_from)
        d2 = self.__adjust_previous(date_to)
        if d1 > d2:
            raise ValueError("The first date must be before the second.")
        return self._index[d2][0] - self._index[d1][0]

    def isbizday(self, dt):
        '''Returns True whether dt is a business day or False otherwise.
        dt is a string ISO-formated date (%%Y-%%m-%%d).'''
        return not self._index[dt][2]

    def __adjust_next(self, dt):
        d1 = timedelta(1)
        dt = Date(dt).date
        while self._index[dt][2]:
            dt = dt + d1
        return Date(dt)

    def adjust_next(self, dt):
        """Returns the next business day whether the passed date isn't a
        business day or returns the given date.
        dt is a string ISO-formated date (%%Y-%%m-%%d).
        """
        return self.__adjust_next(dt).date

    def __adjust_previous(self, dt):
        d1 = timedelta(1)
        dt = Date(dt).date
        while self._index[dt][2]:
            dt = dt - d1
        return Date(dt)

    def adjust_previous(self, dt):
        """Returns the first business day before the passed date whether
        the given date isn't a business day or returns the given date.
        dt is a string ISO-formated date (%%Y-%%m-%%d).
        """
        return self.__adjust_previous(dt).date

    def seq(self, date_from, date_to):
        '''Returns a sequence generator which generates business days between
        the 2 given dates.
        dates is a tuple with 2 strings ISO-formated dates (%%Y-%%m-%%d).
        '''
        _from, _to = Date(date_from), Date(date_to)
        d1 = timedelta(1)
        if _from > _to:
            raise ValueError("The first date must be before the second.")
        _from = self.__adjust_next(_from)
        while _from <= _to:
            yield _from.date
            _from = self.__adjust_next(_from.date + d1)

    def offset(self, dt, n):
        """
        Offsets the given date by n days
        """
        dt = Date(dt)
        if n >= 0:
            d1 = timedelta(1)
            adjust = lambda d: self.__adjust_next(d)
            dt = adjust(dt)
        else:
            d1 = timedelta(-1)
            adjust = lambda d: self.__adjust_previous(d)
            dt = adjust(dt)
            n = abs(n)
        i = 0
        while i < n:
            dt.date += d1
            dt = adjust(dt)
            i += 1
        return dt.date

    @classmethod
    def load(cls, fname):
        if not os.path.exists(fname):
            raise Exception('Invalid calendar specification: \
            file not found (%s)' % fname)
        w = '|'.join(cls._weekdays)
        _holidays = []
        _nonwork_weekdays = []
        with open(fname) as fcal:
            for cal_reg in fcal:
                cal_reg = cal_reg.strip()
                if cal_reg is '': continue
                if re.match('^%s$' % w, cal_reg):
                    _nonwork_weekdays.append(cal_reg)
                else:
                    _holidays.append(Date(cal_reg))
        return Calendar(_holidays, weekdays=_nonwork_weekdays)

    @staticmethod
    def load_holidays(fname):
        if not os.path.exists(fname):
            raise Exception('Invalid calendar specification: \
            file not found (%s)' % fname)
        _holidays = []
        with open(fname) as fcal:
            for cal_reg in fcal:
                cal_reg = cal_reg.strip()
                if cal_reg is '': continue
                _holidays.append(Date(cal_reg).date)
        return _holidays

    def __eq__(self, other):
        return self.startdate == other.startdate and \
            self.enddate == other.enddate and \
            self._cal_spec == other._cal_spec

<|endoftext|>


<commit_before># Copyright 2014 The Alive authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from z3 import *

gbl_unique_id = 0
def mk_unique_id():
  global gbl_unique_id
  id = str(gbl_unique_id)
  gbl_unique_id += 1
  return id

def freshBV(prefix, size):
  return BitVec('%s_%s' % (prefix, mk_unique_id()), size)

def mk_and(l):
  l = [e for e in l if not is_true(e)]
  if len(l) == 0:
    return BoolVal(True)
  if len(l) == 1:
    return l[0]
  return And(l)

def mk_or(l):
  l = [e for e in l if not is_false(e)]
  if len(l) == 0:
    return BoolVal(False)
  if len(l) == 1:
    return l[0]
  return Or(l)

def mk_not(e):
  if is_false(e):
    return BoolVal(True)
  if is_true(e):
    return BoolVal(False)
  return Not(e)

def mk_distinct(l):
  if len(l) < 2:
    return BoolVal(True)
  return Distinct(l)

def mk_implies(a, b):
  if is_true(a):
    return b
  if is_false(a) or is_true(b):
    return BoolVal(True)
  if is_false(b):
    return Not(a)
  return Implies(a, b)

def mk_forall(l, f):
  if l == []:
    return f
  return ForAll(l, f)

def mk_exists(l, f):
  if l == []:
    return f
  return Exists(l, f)

def toBV(b):
  return If(b, BitVecVal(1, 1), BitVecVal(0, 1))

def truncateOrZExt(src, tgt):
  srcb = src.size()
  if isinstance(tgt, int):
    tgtb = tgt
  else:
    tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(tgtb - 1, 0, src)
  return ZeroExt(tgtb - srcb, src)

def truncateOrSExt(src, tgt):
  srcb = src.size()
  tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(tgtb - 1, 0, src)
  return SignExt(tgtb - srcb, src)

def truncateOrPad(src, tgt):
  srcb = src.size()
  tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(srcb - 1, srcb - tgtb, src)
  return Concat(src, BitVecVal(0, tgtb - srcb))

"""
def no_overflow_smul(a, b):
  size = a.size()
  assert b.size() == size
  m = SignExt(size, a) * SignExt(size, b)
  min = BitVecVal(-(1 << (size-1)), 2*size)
  max = BitVecVal((1 << (size-1)) -1, 2*size)
  return And(m >= min, m <= max)
"""
def no_overflow_smul(a, b):
  size = a.size()
  assert b.size() == size
  m = SignExt(size, a) * SignExt(size, b)
  return m == SignExt(size, a * b)

def no_overflow_umul(a, b):
  size = a.size()
  assert b.size() == size
  m = Extract(2*size-1, size, ZeroExt(size, a) * ZeroExt(size, b))
  return m == BitVecVal(0, size)

def bv_log2(v, bitwidth):
  size = v.size()
  def rec(h, l):
    if h <= l:
      return BitVecVal(l, bitwidth)
    mid = l+int((h-l)/2)
    return If(Extract(h,mid+1,v) != 0, rec(h, mid+1), rec(mid, l))
  return rec(size-1, 0)

"""
linear version of log2
def bv_log2(v):
  size = v.size()
  def rec(i):
    if i == 0:
      return BitVecVal(0, size)
    return If(Extract(i,i,v) == BitVecVal(1, 1), BitVecVal(i, size), rec(i-1))
  return rec(size-1)
"""

def ctlz(v, output_width):
  size = v.size()
  def rec(i):
    if i < 0:
      return BitVecVal(size, output_width)
    return If(Extract(i,i,v) == BitVecVal(1, 1),
              BitVecVal(size-1-i, output_width),
              rec(i-1))
  return rec(size-1)

def cttz(v, output_width):
  size = v.size()
  def rec(i):
    if i == size:
      return BitVecVal(size, output_width)
    return If(Extract(i,i,v) == BitVecVal(1, 1),
              BitVecVal(i, output_width),
              rec(i+1))
  return rec(0)

##########################
# Type inference utilities

def register_pick_one_type(v):
  global gbl_one_type_only
  gbl_one_type_only.add(str(v))

def unregister_pick_one_type(vs):
  global gbl_one_type_only
  for v in vs.iterkeys():
    gbl_one_type_only.discard(v)

def reset_pick_one_type():
  global gbl_one_type_only
  gbl_one_type_only = set([])

def get_pick_one_type():
  return gbl_one_type_only

##########################

# number of users of an instruction
def get_users_var(name):
  return BitVec('u_' + name, 8)

def get_flag_var(flag, inst):
  dst = 'src' if gbl_is_source else 'tgt'
  return BitVec('f_%s_%s_%s' % (flag, inst, dst), 1)

def set_smt_is_source(s):
  global gbl_is_source
  gbl_is_source = s

gbl_infer_flags = False
def set_infer_flags(f):
  global gbl_infer_flags
  gbl_infer_flags = f

def do_infer_flags():
  return gbl_infer_flags

##########################
# Error handling

class AliveError(Exception):
  pass

class ParseError():
  def __init__(self, msgs, token = None):
    if isinstance(msgs, list):
      self.msgs = msgs
    else:
      self.msgs = [msgs]
    self.token = token

  def __repr__(self):
    lineno = get_lineno()
    line = get_line(lineno)
    col = get_column(line, self.token)
    return exception2str("\n".join(self.msgs), line, lineno, col)

gbl_line_offset = 0
def exception2str(msg, line, lineno, col, line_offset = None):
  if line_offset is None:
    line_offset = gbl_line_offset
  s  = "ERROR: %s (line: %d)\n" % (msg, line_offset + lineno)
  s += line + '\n'
  s += ' ' * col + '^'
  return s

def get_lineno():
  return gbl_parse_str.count('\n', 0, gbl_parse_loc) + 1

def get_line(lineno):
  return gbl_parse_str.split('\n')[lineno-1]

def get_column(s, tok):
  col = gbl_parse_loc - (gbl_parse_str.rfind('\n', 0, gbl_parse_loc)+1)
  if not tok:
    return col
  token_col = s.find(tok, col)
  return token_col if token_col >= 0 else col

def save_parse_str(s, line):
  global gbl_parse_str, gbl_line_offset
  gbl_parse_str = s
  gbl_line_offset = line-1

def save_loc(loc):
  global gbl_parse_loc
  gbl_parse_loc = loc
<commit_msg>fix linear version of log2 as well<commit_after># Copyright 2014 The Alive authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from z3 import *

gbl_unique_id = 0
def mk_unique_id():
  global gbl_unique_id
  id = str(gbl_unique_id)
  gbl_unique_id += 1
  return id

def freshBV(prefix, size):
  return BitVec('%s_%s' % (prefix, mk_unique_id()), size)

def mk_and(l):
  l = [e for e in l if not is_true(e)]
  if len(l) == 0:
    return BoolVal(True)
  if len(l) == 1:
    return l[0]
  return And(l)

def mk_or(l):
  l = [e for e in l if not is_false(e)]
  if len(l) == 0:
    return BoolVal(False)
  if len(l) == 1:
    return l[0]
  return Or(l)

def mk_not(e):
  if is_false(e):
    return BoolVal(True)
  if is_true(e):
    return BoolVal(False)
  return Not(e)

def mk_distinct(l):
  if len(l) < 2:
    return BoolVal(True)
  return Distinct(l)

def mk_implies(a, b):
  if is_true(a):
    return b
  if is_false(a) or is_true(b):
    return BoolVal(True)
  if is_false(b):
    return Not(a)
  return Implies(a, b)

def mk_forall(l, f):
  if l == []:
    return f
  return ForAll(l, f)

def mk_exists(l, f):
  if l == []:
    return f
  return Exists(l, f)

def toBV(b):
  return If(b, BitVecVal(1, 1), BitVecVal(0, 1))

def truncateOrZExt(src, tgt):
  srcb = src.size()
  if isinstance(tgt, int):
    tgtb = tgt
  else:
    tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(tgtb - 1, 0, src)
  return ZeroExt(tgtb - srcb, src)

def truncateOrSExt(src, tgt):
  srcb = src.size()
  tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(tgtb - 1, 0, src)
  return SignExt(tgtb - srcb, src)

def truncateOrPad(src, tgt):
  srcb = src.size()
  tgtb = tgt.size()
  if srcb == tgtb:
    return src
  if srcb > tgtb:
    return Extract(srcb - 1, srcb - tgtb, src)
  return Concat(src, BitVecVal(0, tgtb - srcb))

"""
def no_overflow_smul(a, b):
  size = a.size()
  assert b.size() == size
  m = SignExt(size, a) * SignExt(size, b)
  min = BitVecVal(-(1 << (size-1)), 2*size)
  max = BitVecVal((1 << (size-1)) -1, 2*size)
  return And(m >= min, m <= max)
"""
def no_overflow_smul(a, b):
  size = a.size()
  assert b.size() == size
  m = SignExt(size, a) * SignExt(size, b)
  return m == SignExt(size, a * b)

def no_overflow_umul(a, b):
  size = a.size()
  assert b.size() == size
  m = Extract(2*size-1, size, ZeroExt(size, a) * ZeroExt(size, b))
  return m == BitVecVal(0, size)

def bv_log2(v, bitwidth):
  def rec(h, l):
    if h <= l:
      return BitVecVal(l, bitwidth)
    mid = l+int((h-l)/2)
    return If(Extract(h,mid+1,v) != 0, rec(h, mid+1), rec(mid, l))
  return rec(v.size()-1, 0)

"""
linear version of log2
def bv_log2(v, bitwidth):
  def rec(i):
    if i == 0:
      return BitVecVal(0, bitwidth)
    return If(Extract(i,i,v) == BitVecVal(1,1), BitVecVal(i,bitwidth), rec(i-1))
  return rec(v.size()-1)
"""

def ctlz(v, output_width):
  size = v.size()
  def rec(i):
    if i < 0:
      return BitVecVal(size, output_width)
    return If(Extract(i,i,v) == BitVecVal(1, 1),
              BitVecVal(size-1-i, output_width),
              rec(i-1))
  return rec(size-1)

def cttz(v, output_width):
  size = v.size()
  def rec(i):
    if i == size:
      return BitVecVal(size, output_width)
    return If(Extract(i,i,v) == BitVecVal(1, 1),
              BitVecVal(i, output_width),
              rec(i+1))
  return rec(0)

##########################
# Type inference utilities

def register_pick_one_type(v):
  global gbl_one_type_only
  gbl_one_type_only.add(str(v))

def unregister_pick_one_type(vs):
  global gbl_one_type_only
  for v in vs.iterkeys():
    gbl_one_type_only.discard(v)

def reset_pick_one_type():
  global gbl_one_type_only
  gbl_one_type_only = set([])

def get_pick_one_type():
  return gbl_one_type_only

##########################

# number of users of an instruction
def get_users_var(name):
  return BitVec('u_' + name, 8)

def get_flag_var(flag, inst):
  dst = 'src' if gbl_is_source else 'tgt'
  return BitVec('f_%s_%s_%s' % (flag, inst, dst), 1)

def set_smt_is_source(s):
  global gbl_is_source
  gbl_is_source = s

gbl_infer_flags = False
def set_infer_flags(f):
  global gbl_infer_flags
  gbl_infer_flags = f

def do_infer_flags():
  return gbl_infer_flags

##########################
# Error handling

class AliveError(Exception):
  pass

class ParseError():
  def __init__(self, msgs, token = None):
    if isinstance(msgs, list):
      self.msgs = msgs
    else:
      self.msgs = [msgs]
    self.token = token

  def __repr__(self):
    lineno = get_lineno()
    line = get_line(lineno)
    col = get_column(line, self.token)
    return exception2str("\n".join(self.msgs), line, lineno, col)

gbl_line_offset = 0
def exception2str(msg, line, lineno, col, line_offset = None):
  if line_offset is None:
    line_offset = gbl_line_offset
  s  = "ERROR: %s (line: %d)\n" % (msg, line_offset + lineno)
  s += line + '\n'
  s += ' ' * col + '^'
  return s

def get_lineno():
  return gbl_parse_str.count('\n', 0, gbl_parse_loc) + 1

def get_line(lineno):
  return gbl_parse_str.split('\n')[lineno-1]

def get_column(s, tok):
  col = gbl_parse_loc - (gbl_parse_str.rfind('\n', 0, gbl_parse_loc)+1)
  if not tok:
    return col
  token_col = s.find(tok, col)
  return token_col if token_col >= 0 else col

def save_parse_str(s, line):
  global gbl_parse_str, gbl_line_offset
  gbl_parse_str = s
  gbl_line_offset = line-1

def save_loc(loc):
  global gbl_parse_loc
  gbl_parse_loc = loc
<|endoftext|>


<commit_before>"""Serve nix store objects over HTTP."""
import argparse
import logging
import os
from os.path import exists, isdir, join, basename, dirname
import re
from subprocess import check_output, Popen, PIPE

from flask import Flask, make_response, send_file

from servenix.utils import decode_str
from servenix.exceptions import NoSuchObject, NoNarGenerated, CouldNotUpdateHash

_HASH_REGEX=re.compile(r"[a-z0-9]{32}")
_PATH_REGEX=re.compile(r"([a-z0-9]{32})-.*")

class NixServer(Flask):
    """Serves nix packages."""
    def __init__(self, nix_store_path, nix_bin_path, nix_state_path,
                 compression_type):
        self._nix_store_path = nix_store_path
        self._nix_state_path = nix_state_path
        self._nix_bin_path = nix_bin_path
        self._compression_type = compression_type
        # Cache mapping object hashes to store paths.
        self._hashes_to_paths = {}
        # Cache mapping store paths to object info.
        self._paths_to_info = {}
        # A static string telling a nix client what this store serves.
        self._cache_info = "\n".join([
            "StoreDir: {}".format(self._nix_store_path),
            "WantMassQuery: 1",
            "Priority: 30"
        ])
        if self._compression_type == "bzip2":
            self._content_type = "application/x-bzip2"
            self._nar_extension = ".nar.bz2"
        else:
            self._content_type = "application/x-xz"
            self._nar_extension = ".nar.xz"

    def store_path_from_hash(self, store_object_hash):
        """Look up a store path using its hash.

        The name of every object in the nix store starts with a
        32-character hash. We can find the full path by finding an
        object that starts with this hash.

        :param store_object_hash: The 32-character hash prefix of the object.
        :type store_object_hash: ``str``

        :return: The full store path to the object.
        :rtype: ``str``

        :raises: :py:class:`NoSuchObject` if the object isn't in the store.
        """
        if store_object_hash in self._hashes_to_paths:
            return self._hashes_to_paths[store_object_hash]
        store_objects = map(decode_str, os.listdir(self._nix_store_path))
        for obj in store_objects:
            match = _PATH_REGEX.match(obj)
            if match is None:
                continue
            full_path = join(self._nix_store_path, obj)
            _hash = match.group(1)
            self._hashes_to_paths[_hash] = full_path
            if _hash == store_object_hash:
                return full_path
        raise NoSuchObject("No object with hash {} was found."
                           .format(store_object_hash))

    def get_object_info(self, store_path):
        """Given a store path, get some information about the path.

        :param store_path: Path to the object in the store.
        :type store_path: ``str``

        :return: A dictionary of store object information.
        :rtype: ``dict``
        """
        if store_path in self._paths_to_info:
            return self._paths_to_info[store_path]
        # Invoke nix-store with various queries to get package info.
        nix_store_q = lambda option: decode_str(check_output([
            "{}/nix-store".format(self._nix_bin_path),
            "--query", option, store_path
        ])).strip()
        # Build the compressed version. Compute its hash and size.
        nar_path = self.build_nar(store_path)
        du = check_output("du -sb {}".format(nar_path), shell=True)
        file_size = int(du.split()[0])
        file_hash = decode_str(check_output(
            "nix-hash --type sha256 --base32 --flat {}".format(nar_path),
            shell=True)).strip()
        # Some paths have corrupt hashes stored in the sqlite
        # database. I'm not sure why this happens, but we check the
        # actual hash using nix-hash and if it doesn't match what
        # `nix-store -q --hash` says, we update the sqlite database to
        # fix the error.
        computed_store_obj_hash = decode_str(check_output(
            "nix-hash --type sha256 {}".format(store_path),
            shell=True)).strip()
        registered_store_obj_hash = nix_store_q("--hash")
        if computed_store_obj_hash != registered_store_obj_hash:
            db = join(self._nix_state_path, "nix", "db", "db.sqlite")
            logging.warn("Incorrect hash {} stored for path {}. Updating."
                         .format(registered_store_obj_hash, store_path))
            query = ("UPDATE ValidPaths SET hash = '{}' where path = '{}';"
                     .format(computed_store_obj_hash, store_path))
            proc = Popen(['sqlite3', db], stdin=PIPE, stderr=PIPE)
            err = decode_str(proc.communicate(input=bytes(query, "utf-8"))[1])
            if proc.wait() != 0:
                raise CouldNotUpdateHash(path, registered_store_obj_hash,
                                         computed_store_obj_hash, err)
        info = {
            "StorePath": store_path,
            "NarHash": computed_store_obj_hash,
            "NarSize": nix_store_q("--size"),
            "FileSize": str(file_size),
            "FileHash": "sha256:{}".format(file_hash)
        }
        references = nix_store_q("--references").split()
        if references != []:
            info["References"] = " ".join(map(basename, references))
        deriver = nix_store_q("--deriver")
        if deriver != "unknown-deriver":
            info["Deriver"] = basename(deriver)
        self._paths_to_info[store_path] = info
        return info

    def build_nar(self, store_path):
        """Build a nix archive (nar) and return the resulting path."""
        # Construct a nix expression which will produce a nar.
        nar_expr = "".join([
            "(import <nix/nar.nix> {",
            'storePath = "{}";'.format(store_path),
            'hashAlgo = "sha256";',
            'compressionType = "{}";'.format(self._compression_type),
            "})"])
        # Nix-build this expression, resulting in a store object.
        compressed_path = decode_str(check_output([
            join(self._nix_bin_path, "nix-build"),
            "--expr", nar_expr, "--no-out-link"
        ])).strip()

        # This path will contain a compressed file; return its path.
        contents = map(decode_str, os.listdir(compressed_path))
        for filename in contents:
            if filename.endswith(self._nar_extension):
                return join(compressed_path, filename)
        raise NoNarGenerated(compressed_path, self._nar_extension)

    def make_app(self):
        """Create a flask app and set up routes on it.

        :return: A flask app.
        :rtype: :py:class:`Flask`
        """
        app = Flask(__name__)

        @app.route("/nix-cache-info")
        def nix_cache_info():
            """Return information about the binary cache."""
            return self._cache_info

        @app.route("/<obj_hash>.narinfo")
        def get_narinfo(obj_hash):
            """Given an object's 32-character hash, return information on it.

            The information includes the object's size (uncompressed), sha256
            hash, store path, and reference graph.

            If the object isn't found, return a 404.

            :param obj_hash: First 32 characters of the object's store path.
            :type obj_hash: ``str``
            """
            if _HASH_REGEX.match(obj_hash) is None:
                 return ("Hash {} must match {}"
                         .format(obj_hash, _HASH_REGEX.pattern), 400)
            try:
                store_path = self.store_path_from_hash(obj_hash)
                store_info = self.get_object_info(store_path)
            except NoSuchObject as err:
                return (err.message, 404)
            # Add a few more keys to the store object, specific to the
            # compression type we're serving.
            store_info["URL"] = "nar/{}{}".format(obj_hash,
                                                  self._nar_extension)
            store_info["Compression"] = self._compression_type
            info_string = "\n".join("{}: {}".format(k, v)
                             for k, v in store_info.items()) + "\n"
            return make_response((info_string, 200,
                                 {"Content-Type": "text/x-nix-narinfo"}))

        @app.route("/nar/<obj_hash>{}".format(self._nar_extension))
        def serve_nar(obj_hash):
            """Return the compressed binary from the nix store.

            If the object isn't found, return a 404.

            :param obj_hash: First 32 characters of the object's store path.
            :type obj_hash: ``str``
            """
            try:
                store_path = self.store_path_from_hash(obj_hash)
            except NoSuchObject as err:
                return (err.message, 404)
            nar_path = self.build_nar(store_path)
            return send_file(nar_path, mimetype=self._content_type)

        return app

def _get_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(prog="servenix")
    parser.add_argument("--port", type=int, default=5000,
                        help="Port to listen on.")
    parser.add_argument("--compression-type", default="xz",
                        choices=("xz", "bzip2"),
                        help="How served objects should be compressed.")
    return parser.parse_args()

def main():
    """Main entry point."""
    try:
        nix_bin_path = os.environ["NIX_BIN_PATH"]
        assert exists(join(nix_bin_path, "nix-store"))
        # The store path can be given explicitly, or else it will be
        # inferred to be 2 levels up from the bin path. E.g., if the
        # bin path is /foo/bar/123-nix/bin, the store directory will
        # be /foo/bar.
        nix_store_path = os.environ.get("NIX_STORE_PATH",
                                        dirname(dirname(nix_bin_path)))
        assert isdir(nix_store_path), \
            "Nix store directory {} doesn't exist".format(nix_store_path)
        # The state path can be given explicitly, or else it will be
        # inferred to be sibling to the store directory.
        nix_state_path = os.environ.get("NIX_STATE_PATH",
                                        join(dirname(nix_store_path), "var"))
        assert isdir(nix_state_path), \
            "Nix state directory {} doesn't exist".format(nix_state_path)
    except KeyError as err:
        exit("Invalid environment: variable {} must be set.".format(err))
    db = join(nix_state_path, "nix", "db", "db.sqlite")
    store_path = '/opt/ns/nix/store/002lkaccqahn2as9vgq56qgqrn36dlhb-python_testing'
    registered_store_obj_hash = '0hql4j2prkwx87f4d1isjh7kr9f2xfbywcjg14ldk56s2qx563y8'
    computed_store_obj_hash = registered_store_obj_hash
    logging.warn("Incorrect hash {} stored for path {}. Updating."
                 .format(registered_store_obj_hash, store_path))
    query = ("UPDATE ValidPaths SET hash = 'sha256:{}' where path = '{}';"
             .format(computed_store_obj_hash, store_path))
    proc = Popen(['sqlite3', db], stdin=PIPE, stderr=PIPE)
    err = decode_str(proc.communicate(input=bytes(query, "utf-8"))[1])
    if proc.wait() != 0:
        raise CouldNotUpdateHash(path, registered_store_obj_hash,
                                 computed_store_obj_hash, err)
    args = _get_args()
    nixserver = NixServer(nix_store_path=nix_store_path,
                          nix_state_path=nix_state_path,
                          nix_bin_path=nix_bin_path,
                          compression_type=args.compression_type)
    app = nixserver.make_app()
    app.run(port=args.port)

if __name__ == "__main__":
    main()
<commit_msg>remove debuggery<commit_after>"""Serve nix store objects over HTTP."""
import argparse
import logging
import os
from os.path import exists, isdir, join, basename, dirname
import re
from subprocess import check_output, Popen, PIPE

from flask import Flask, make_response, send_file

from servenix.utils import decode_str
from servenix.exceptions import NoSuchObject, NoNarGenerated, CouldNotUpdateHash

_HASH_REGEX=re.compile(r"[a-z0-9]{32}")
_PATH_REGEX=re.compile(r"([a-z0-9]{32})-.*")

class NixServer(Flask):
    """Serves nix packages."""
    def __init__(self, nix_store_path, nix_bin_path, nix_state_path,
                 compression_type):
        self._nix_store_path = nix_store_path
        self._nix_state_path = nix_state_path
        self._nix_bin_path = nix_bin_path
        self._compression_type = compression_type
        # Cache mapping object hashes to store paths.
        self._hashes_to_paths = {}
        # Cache mapping store paths to object info.
        self._paths_to_info = {}
        # A static string telling a nix client what this store serves.
        self._cache_info = "\n".join([
            "StoreDir: {}".format(self._nix_store_path),
            "WantMassQuery: 1",
            "Priority: 30"
        ])
        if self._compression_type == "bzip2":
            self._content_type = "application/x-bzip2"
            self._nar_extension = ".nar.bz2"
        else:
            self._content_type = "application/x-xz"
            self._nar_extension = ".nar.xz"

    def store_path_from_hash(self, store_object_hash):
        """Look up a store path using its hash.

        The name of every object in the nix store starts with a
        32-character hash. We can find the full path by finding an
        object that starts with this hash.

        :param store_object_hash: The 32-character hash prefix of the object.
        :type store_object_hash: ``str``

        :return: The full store path to the object.
        :rtype: ``str``

        :raises: :py:class:`NoSuchObject` if the object isn't in the store.
        """
        if store_object_hash in self._hashes_to_paths:
            return self._hashes_to_paths[store_object_hash]
        store_objects = map(decode_str, os.listdir(self._nix_store_path))
        for obj in store_objects:
            match = _PATH_REGEX.match(obj)
            if match is None:
                continue
            full_path = join(self._nix_store_path, obj)
            _hash = match.group(1)
            self._hashes_to_paths[_hash] = full_path
            if _hash == store_object_hash:
                return full_path
        raise NoSuchObject("No object with hash {} was found."
                           .format(store_object_hash))

    def get_object_info(self, store_path):
        """Given a store path, get some information about the path.

        :param store_path: Path to the object in the store.
        :type store_path: ``str``

        :return: A dictionary of store object information.
        :rtype: ``dict``
        """
        if store_path in self._paths_to_info:
            return self._paths_to_info[store_path]
        # Invoke nix-store with various queries to get package info.
        nix_store_q = lambda option: decode_str(check_output([
            "{}/nix-store".format(self._nix_bin_path),
            "--query", option, store_path
        ])).strip()
        # Build the compressed version. Compute its hash and size.
        nar_path = self.build_nar(store_path)
        du = check_output("du -sb {}".format(nar_path), shell=True)
        file_size = int(du.split()[0])
        file_hash = decode_str(check_output(
            "nix-hash --type sha256 --base32 --flat {}".format(nar_path),
            shell=True)).strip()
        # Some paths have corrupt hashes stored in the sqlite
        # database. I'm not sure why this happens, but we check the
        # actual hash using nix-hash and if it doesn't match what
        # `nix-store -q --hash` says, we update the sqlite database to
        # fix the error.
        computed_store_obj_hash = decode_str(check_output(
            "nix-hash --type sha256 {}".format(store_path),
            shell=True)).strip()
        registered_store_obj_hash = nix_store_q("--hash")
        if computed_store_obj_hash != registered_store_obj_hash:
            db = join(self._nix_state_path, "nix", "db", "db.sqlite")
            logging.warn("Incorrect hash {} stored for path {}. Updating."
                         .format(registered_store_obj_hash, store_path))
            query = ("UPDATE ValidPaths SET hash = '{}' where path = '{}';"
                     .format(computed_store_obj_hash, store_path))
            proc = Popen(['sqlite3', db], stdin=PIPE, stderr=PIPE)
            err = decode_str(proc.communicate(input=bytes(query, "utf-8"))[1])
            if proc.wait() != 0:
                raise CouldNotUpdateHash(path, registered_store_obj_hash,
                                         computed_store_obj_hash, err)
        info = {
            "StorePath": store_path,
            "NarHash": computed_store_obj_hash,
            "NarSize": nix_store_q("--size"),
            "FileSize": str(file_size),
            "FileHash": "sha256:{}".format(file_hash)
        }
        references = nix_store_q("--references").split()
        if references != []:
            info["References"] = " ".join(map(basename, references))
        deriver = nix_store_q("--deriver")
        if deriver != "unknown-deriver":
            info["Deriver"] = basename(deriver)
        self._paths_to_info[store_path] = info
        return info

    def build_nar(self, store_path):
        """Build a nix archive (nar) and return the resulting path."""
        # Construct a nix expression which will produce a nar.
        nar_expr = "".join([
            "(import <nix/nar.nix> {",
            'storePath = "{}";'.format(store_path),
            'hashAlgo = "sha256";',
            'compressionType = "{}";'.format(self._compression_type),
            "})"])
        # Nix-build this expression, resulting in a store object.
        compressed_path = decode_str(check_output([
            join(self._nix_bin_path, "nix-build"),
            "--expr", nar_expr, "--no-out-link"
        ])).strip()

        # This path will contain a compressed file; return its path.
        contents = map(decode_str, os.listdir(compressed_path))
        for filename in contents:
            if filename.endswith(self._nar_extension):
                return join(compressed_path, filename)
        raise NoNarGenerated(compressed_path, self._nar_extension)

    def make_app(self):
        """Create a flask app and set up routes on it.

        :return: A flask app.
        :rtype: :py:class:`Flask`
        """
        app = Flask(__name__)

        @app.route("/nix-cache-info")
        def nix_cache_info():
            """Return information about the binary cache."""
            return self._cache_info

        @app.route("/<obj_hash>.narinfo")
        def get_narinfo(obj_hash):
            """Given an object's 32-character hash, return information on it.

            The information includes the object's size (uncompressed), sha256
            hash, store path, and reference graph.

            If the object isn't found, return a 404.

            :param obj_hash: First 32 characters of the object's store path.
            :type obj_hash: ``str``
            """
            if _HASH_REGEX.match(obj_hash) is None:
                 return ("Hash {} must match {}"
                         .format(obj_hash, _HASH_REGEX.pattern), 400)
            try:
                store_path = self.store_path_from_hash(obj_hash)
                store_info = self.get_object_info(store_path)
            except NoSuchObject as err:
                return (err.message, 404)
            # Add a few more keys to the store object, specific to the
            # compression type we're serving.
            store_info["URL"] = "nar/{}{}".format(obj_hash,
                                                  self._nar_extension)
            store_info["Compression"] = self._compression_type
            info_string = "\n".join("{}: {}".format(k, v)
                             for k, v in store_info.items()) + "\n"
            return make_response((info_string, 200,
                                 {"Content-Type": "text/x-nix-narinfo"}))

        @app.route("/nar/<obj_hash>{}".format(self._nar_extension))
        def serve_nar(obj_hash):
            """Return the compressed binary from the nix store.

            If the object isn't found, return a 404.

            :param obj_hash: First 32 characters of the object's store path.
            :type obj_hash: ``str``
            """
            try:
                store_path = self.store_path_from_hash(obj_hash)
            except NoSuchObject as err:
                return (err.message, 404)
            nar_path = self.build_nar(store_path)
            return send_file(nar_path, mimetype=self._content_type)

        return app

def _get_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(prog="servenix")
    parser.add_argument("--port", type=int, default=5000,
                        help="Port to listen on.")
    parser.add_argument("--compression-type", default="xz",
                        choices=("xz", "bzip2"),
                        help="How served objects should be compressed.")
    return parser.parse_args()

def main():
    """Main entry point."""
    try:
        nix_bin_path = os.environ["NIX_BIN_PATH"]
        assert exists(join(nix_bin_path, "nix-store"))
        # The store path can be given explicitly, or else it will be
        # inferred to be 2 levels up from the bin path. E.g., if the
        # bin path is /foo/bar/123-nix/bin, the store directory will
        # be /foo/bar.
        nix_store_path = os.environ.get("NIX_STORE_PATH",
                                        dirname(dirname(nix_bin_path)))
        assert isdir(nix_store_path), \
            "Nix store directory {} doesn't exist".format(nix_store_path)
        # The state path can be given explicitly, or else it will be
        # inferred to be sibling to the store directory.
        nix_state_path = os.environ.get("NIX_STATE_PATH",
                                        join(dirname(nix_store_path), "var"))
        assert isdir(nix_state_path), \
            "Nix state directory {} doesn't exist".format(nix_state_path)
    except KeyError as err:
        exit("Invalid environment: variable {} must be set.".format(err))
    args = _get_args()
    nixserver = NixServer(nix_store_path=nix_store_path,
                          nix_state_path=nix_state_path,
                          nix_bin_path=nix_bin_path,
                          compression_type=args.compression_type)
    app = nixserver.make_app()
    app.run(port=args.port)

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