<commit_before>from __future__ import print_function

import atexit
import numbers
import os
import sys
import time

from six.moves import urllib

from mlflow.entities.experiment import Experiment
from mlflow.entities.param import Param
from mlflow.entities.metric import Metric
from mlflow.entities.run_status import RunStatus
from mlflow.entities.source_type import SourceType
from mlflow.store.file_store import FileStore
from mlflow.store.rest_store import RestStore
from mlflow.store.artifact_repo import ArtifactRepository
from mlflow.utils import env, rest_utils
from mlflow.utils.validation import _validate_metric_name, _validate_param_name, _validate_run_id

_RUN_ID_ENV_VAR = "MLFLOW_RUN_ID"
_TRACKING_URI_ENV_VAR = "MLFLOW_TRACKING_URI"
_EXPERIMENT_ID_ENV_VAR = "MLFLOW_EXPERIMENT_ID"
_DEFAULT_USER_ID = "unknown"
_LOCAL_FS_URI_PREFIX = "file:///"
_REMOTE_URI_PREFIX = "http://"

_active_run = None
_tracking_uri = None

def _get_user_id():
    """Get the ID of the user for the current run."""
    try:
        import pwd
        return pwd.getpwuid(os.getuid())[0]
    except ImportError:
        return _DEFAULT_USER_ID

def set_tracking_uri(uri):
    """
    Set the tracking server URI to the passed-in value. This does not affect the
    currently active run (if one exists), but takes effect for any successive runs.

    The provided URI can be one of three types:

    - An empty string, or a local file path, prefixed with ``file:/``. Data is stored
      locally at the provided file (or ``./mlruns`` if empty).
    - An HTTP URI like ``https://my-tracking-server:5000``.
    - A Databricks workspace, provided as just the string 'databricks' or, to use a specific
      Databricks profile (per the Databricks CLI), 'databricks://profileName'.
    """
    global _tracking_uri
    _tracking_uri = uri

def get_tracking_uri():
    """
    Return the current tracking URI. This may not correspond to the tracking URI of
    the currently active run, since the tracking URI can be updated via ``set_tracking_uri``.

    :return: the tracking URI.
    """
    global _tracking_uri
    if _tracking_uri is not None:
        return _tracking_uri
    elif env.get_env(_TRACKING_URI_ENV_VAR) is not None:
        return env.get_env(_TRACKING_URI_ENV_VAR)
    else:
        return os.path.abspath("./mlruns")

def is_local_uri(uri):
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == '' or scheme == 'file'

def _is_http_uri(uri):
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == 'http' or scheme == 'https'

def _is_databricks_uri(uri):
    """Databricks URIs look like 'databricks' (default profile) or 'databricks://profile'"""
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == 'databricks' or uri == 'databricks'

def _get_file_store(store_uri):
    path = urllib.parse.urlparse(store_uri).path
    return FileStore(path)

def _get_rest_store(store_uri):
    return RestStore({'hostname': store_uri})

def _get_databricks_rest_store(store_uri):
    parsed_uri = urllib.parse.urlparse(store_uri)

    profile = None
    if parsed_uri.scheme == 'databricks':
        profile = parsed_uri.hostname
    http_request_kwargs = rest_utils.get_databricks_http_request_kwargs_or_fail(profile)
    return RestStore(http_request_kwargs)

def _get_store():
    store_uri = get_tracking_uri()
    # Default: if URI hasn't been set, return a FileStore
    if store_uri is None:
        return FileStore()
    # Pattern-match on the URI
    if _is_databricks_uri(store_uri):
        return _get_databricks_rest_store(store_uri)
    if is_local_uri(store_uri):
        return _get_file_store(store_uri)
    if _is_http_uri(store_uri):
        return _get_rest_store(store_uri)

    raise Exception("Tracking URI must be a local filesystem URI of the form '%s...' or a "
                    "remote URI of the form '%s...'. Update the tracking URI via "
                    "mlflow.set_tracking_uri" % (_LOCAL_FS_URI_PREFIX, _REMOTE_URI_PREFIX))

class ActiveRun(object):
    """
    Class representing an active run. Has a reference to the store to which state for the run
    (e.g. run metadata, metrics, parameters, and artifacts) should be persisted.

    Contains methods for logging metrics, parameters, etc. under the current run.

    :param run_info: ``RunInfo`` describing the active run. A corresponding ``Run`` object is
                     assumed to already be persisted with state "running" in ``store``.
    :param store: Backend store to which the current run should persist state updates.
    """
    def __init__(self, run_info, store):
        self.store = store
        self.run_info = run_info
        if run_info.artifact_uri:
            self.artifact_repo = ArtifactRepository.from_artifact_uri(run_info.artifact_uri)
        else:
            self.artifact_repo = _get_legacy_artifact_repo(store, run_info)

    def set_terminated(self, status):
        self.run_info = self.store.update_run_info(
            self.run_info.run_uuid, run_status=RunStatus.from_string(status),
            end_time=_get_unix_timestamp())

    def log_metric(self, metric):
        _validate_metric_name(metric.key)
        self.store.log_metric(self.run_info.run_uuid, metric)

    def log_param(self, param):
        _validate_param_name(param.key)
        self.store.log_param(self.run_info.run_uuid, param)

    def log_artifact(self, local_path, artifact_path=None):
        self.artifact_repo.log_artifact(local_path, artifact_path)

    def log_artifacts(self, local_dir, artifact_path=None):
        self.artifact_repo.log_artifacts(local_dir, artifact_path)

    def get_artifact_uri(self):
        return self.artifact_repo.artifact_uri

    def get_run(self):
        return self.store.get_run(self.run_info.run_uuid)

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        status = "FINISHED" if exc_type is None else "FAILED"
        self.set_terminated(status)
        global _active_run
        _active_run = None
        return exc_type is None

def list_experiments():
    """
    Return a list of all experiments.
    """
    return _get_store().list_experiments()

def create_experiment(experiment_name):
    """
    Create an experiment with the specified name and return its ID.
    """
    if experiment_name is None or experiment_name == "":
        raise Exception("Invalid experiment name '%s'" % experiment_name)
    return _get_store().create_experiment(experiment_name)

def _get_main_file():
    if len(sys.argv) > 0:
        return sys.argv[0]
    return None

def _get_source_name():
    main_file = _get_main_file()
    if main_file is not None:
        return main_file
    return "<console>"

def _get_source_version():
    main_file = _get_main_file()
    if main_file is not None:
        return _get_git_commit(main_file)
    return None

def _get_source_type():
    return SourceType.LOCAL

def _get_experiment_id():
    return int(env.get_env(_EXPERIMENT_ID_ENV_VAR) or Experiment.DEFAULT_EXPERIMENT_ID)

def _get_unix_timestamp():
    return int(time.time() * 1000)

def _create_run(experiment_id, source_name, source_version, entry_point_name, source_type):
    store = _get_store()
    run = store.create_run(experiment_id=experiment_id, user_id=_get_user_id(), run_name=None,
                           source_type=source_type,
                           source_name=source_name,
                           entry_point_name=entry_point_name,
                           start_time=_get_unix_timestamp(),
                           source_version=source_version, tags=[])
    return ActiveRun(run.info, store)

def get_run(run_uuid):
    """Return the run with the specified run UUID from the current tracking server."""
    _validate_run_id(run_uuid)
    return _get_store().get_run(run_uuid)

def _get_existing_run(run_uuid):
    _validate_run_id(run_uuid)
    existing_run = get_run(run_uuid)
    if existing_run is None:
        raise Exception("Could not start run with UUID %s - no such run found." % run_uuid)
    store = _get_store()
    updated_info = store.update_run_info(
        run_uuid=run_uuid, run_status=RunStatus.RUNNING, end_time=None)
    return ActiveRun(updated_info, store)

def start_run(run_uuid=None, experiment_id=None, source_name=None, source_version=None,
              entry_point_name=None, source_type=None):
    """
    Start a new MLflow run, setting it as the active run under which metrics and params
    will be logged. The return value can be used as a context manager within a ``with`` block;
    otherwise, ``end_run()`` must be called to terminate the current run. If ``run_uuid``
    is passed or the ``MLFLOW_RUN_ID`` environment variable is set, ``start_run`` attempts to
    resume a run with the specified run ID (with ``run_uuid`` taking precedence over
    ``MLFLOW_RUN_ID``), and other parameters are ignored.

    :param run_uuid: If specified, get the run with the specified UUID and log metrics
                     and params under that run. The run's end time is unset and its status
                     is set to running, but the run's other attributes remain unchanged
                     (the run's ``source_version``, ``source_type``, etc. are not changed).
    :param experiment_id: Used only when ``run_uuid`` is unspecified. ID of the experiment under
                          which to create the current run. If unspecified, the run is created under
                          a new experiment with a randomly generated name.
    :param source_name: Name of the source file or URI of the project to be associated with the run.
                        Defaults to the current file if none provided.
    :param source_version: Optional Git commit hash to associate with the run.
    :param entry_point_name: Optional name of the entry point for to the current run.
    :param source_type: Integer enum value describing the type of the run
                        ("local", "project", etc.). Defaults to
                        ``mlflow.entities.source_type.SourceType.LOCAL``.
    :return: :py:class:`mlflow.tracking.ActiveRun` object that acts as a context manager wrapping
             the run's state.
    """
    global _active_run
    if _active_run:
        raise Exception("Run with UUID %s is already active, unable to start nested "
                        "run" % _active_run.run_info.run_uuid)
    existing_run_uuid = run_uuid or os.environ.get(_RUN_ID_ENV_VAR, None)
    if existing_run_uuid:
        _validate_run_id(existing_run_uuid)
        active_run_obj = _get_existing_run(existing_run_uuid)
    else:
        exp_id_for_run = experiment_id or _get_experiment_id()
        active_run_obj = _create_run(experiment_id=exp_id_for_run,
                                     source_name=source_name or _get_source_name(),
                                     source_version=source_version or _get_source_version(),
                                     entry_point_name=entry_point_name,
                                     source_type=source_type or _get_source_type())
    _active_run = active_run_obj
    return _active_run

def _get_or_start_run():
    if _active_run:
        return _active_run
    return start_run()

def end_run(status="FINISHED"):
    global _active_run
    if _active_run:
        _active_run.set_terminated(status)
        # Clear out the global existing run environment variable as well.
        env.unset_variable(_RUN_ID_ENV_VAR)
        _active_run = None

def active_run():
    """Return the currently active ``Run``, or None if no such run exists."""
    if _active_run:
        return _active_run.get_run()
    else:
        return None

def log_param(key, value):
    """
    Log the passed-in parameter under the current run, creating a run if necessary.

    :param key: Parameter name (string)
    :param value: Parameter value (string)
    """
    _get_or_start_run().log_param(Param(key, str(value)))

def log_metric(key, value):
    """
    Log the passed-in metric under the current run, creating a run if necessary.

    :param key: Metric name (string).
    :param value: Metric value (float).
    """
    if not isinstance(value, numbers.Number):
        print("WARNING: The metric {}={} was not logged because the value is not a number.".format(
            key, value), file=sys.stderr)
        return
    _get_or_start_run().log_metric(Metric(key, value, int(time.time())))

def log_artifact(local_path, artifact_path=None):
    """Log a local file or directory as an artifact of the currently active run."""
    _get_or_start_run().log_artifact(local_path, artifact_path)

def log_artifacts(local_dir, artifact_path=None):
    """Log all the contents of a local directory as artifacts of the run."""
    _get_or_start_run().log_artifacts(local_dir, artifact_path)

def get_artifact_uri():
    """
    Return the artifact URI of the currently active run. Calls to ``log_artifact`` and
    ``log_artifacts`` write artifact(s) to subdirectories of the returned URI.
    """
    return _get_or_start_run().get_artifact_uri()

atexit.register(end_run)

def _get_model_log_dir(model_name, run_id):
    if not run_id:
        raise Exception("Must specify a run_id to get logging directory for a model.")
    run = _get_store().get_run(run_id)
    if run.info.artifact_uri:
        artifact_repo = ArtifactRepository.from_artifact_uri(run.info.artifact_uri)
    else:
        artifact_repo = _get_legacy_artifact_repo(_get_store(), run.info)
    return artifact_repo.download_artifacts(model_name)

def _get_legacy_artifact_repo(file_store, run_info):
    # TODO(aaron) Remove this once everyone locally only has runs from after
    # the introduction of "artifact_uri".
    uri = os.path.join(file_store.root_directory, str(run_info.experiment_id),
                       run_info.run_uuid, "artifacts")
    return ArtifactRepository.from_artifact_uri(uri)

def _get_git_commit(path):
    try:
        from git import Repo, InvalidGitRepositoryError, GitCommandNotFound
    except ImportError as e:
        print("Notice: failed to import Git (the git executable is probably not on your PATH),"
              " so Git SHA is not available. Error: %s" % e, file=sys.stderr)
        return None
    try:
        if os.path.isfile(path):
            path = os.path.dirname(path)
        repo = Repo(path, search_parent_directories=True)
        commit = repo.head.commit.hexsha
        return commit
    except (InvalidGitRepositoryError, GitCommandNotFound, ValueError):
        return None
<commit_msg>Plug leak in Git commit finding logic (#243)<commit_after>from __future__ import print_function

import atexit
import numbers
import os
import sys
import time

from six.moves import urllib

from mlflow.entities.experiment import Experiment
from mlflow.entities.param import Param
from mlflow.entities.metric import Metric
from mlflow.entities.run_status import RunStatus
from mlflow.entities.source_type import SourceType
from mlflow.store.file_store import FileStore
from mlflow.store.rest_store import RestStore
from mlflow.store.artifact_repo import ArtifactRepository
from mlflow.utils import env, rest_utils
from mlflow.utils.validation import _validate_metric_name, _validate_param_name, _validate_run_id

_RUN_ID_ENV_VAR = "MLFLOW_RUN_ID"
_TRACKING_URI_ENV_VAR = "MLFLOW_TRACKING_URI"
_EXPERIMENT_ID_ENV_VAR = "MLFLOW_EXPERIMENT_ID"
_DEFAULT_USER_ID = "unknown"
_LOCAL_FS_URI_PREFIX = "file:///"
_REMOTE_URI_PREFIX = "http://"

_active_run = None
_tracking_uri = None

def _get_user_id():
    """Get the ID of the user for the current run."""
    try:
        import pwd
        return pwd.getpwuid(os.getuid())[0]
    except ImportError:
        return _DEFAULT_USER_ID

def set_tracking_uri(uri):
    """
    Set the tracking server URI to the passed-in value. This does not affect the
    currently active run (if one exists), but takes effect for any successive runs.

    The provided URI can be one of three types:

    - An empty string, or a local file path, prefixed with ``file:/``. Data is stored
      locally at the provided file (or ``./mlruns`` if empty).
    - An HTTP URI like ``https://my-tracking-server:5000``.
    - A Databricks workspace, provided as just the string 'databricks' or, to use a specific
      Databricks profile (per the Databricks CLI), 'databricks://profileName'.
    """
    global _tracking_uri
    _tracking_uri = uri

def get_tracking_uri():
    """
    Return the current tracking URI. This may not correspond to the tracking URI of
    the currently active run, since the tracking URI can be updated via ``set_tracking_uri``.

    :return: the tracking URI.
    """
    global _tracking_uri
    if _tracking_uri is not None:
        return _tracking_uri
    elif env.get_env(_TRACKING_URI_ENV_VAR) is not None:
        return env.get_env(_TRACKING_URI_ENV_VAR)
    else:
        return os.path.abspath("./mlruns")

def is_local_uri(uri):
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == '' or scheme == 'file'

def _is_http_uri(uri):
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == 'http' or scheme == 'https'

def _is_databricks_uri(uri):
    """Databricks URIs look like 'databricks' (default profile) or 'databricks://profile'"""
    scheme = urllib.parse.urlparse(uri).scheme
    return scheme == 'databricks' or uri == 'databricks'

def _get_file_store(store_uri):
    path = urllib.parse.urlparse(store_uri).path
    return FileStore(path)

def _get_rest_store(store_uri):
    return RestStore({'hostname': store_uri})

def _get_databricks_rest_store(store_uri):
    parsed_uri = urllib.parse.urlparse(store_uri)

    profile = None
    if parsed_uri.scheme == 'databricks':
        profile = parsed_uri.hostname
    http_request_kwargs = rest_utils.get_databricks_http_request_kwargs_or_fail(profile)
    return RestStore(http_request_kwargs)

def _get_store():
    store_uri = get_tracking_uri()
    # Default: if URI hasn't been set, return a FileStore
    if store_uri is None:
        return FileStore()
    # Pattern-match on the URI
    if _is_databricks_uri(store_uri):
        return _get_databricks_rest_store(store_uri)
    if is_local_uri(store_uri):
        return _get_file_store(store_uri)
    if _is_http_uri(store_uri):
        return _get_rest_store(store_uri)

    raise Exception("Tracking URI must be a local filesystem URI of the form '%s...' or a "
                    "remote URI of the form '%s...'. Update the tracking URI via "
                    "mlflow.set_tracking_uri" % (_LOCAL_FS_URI_PREFIX, _REMOTE_URI_PREFIX))

class ActiveRun(object):
    """
    Class representing an active run. Has a reference to the store to which state for the run
    (e.g. run metadata, metrics, parameters, and artifacts) should be persisted.

    Contains methods for logging metrics, parameters, etc. under the current run.

    :param run_info: ``RunInfo`` describing the active run. A corresponding ``Run`` object is
                     assumed to already be persisted with state "running" in ``store``.
    :param store: Backend store to which the current run should persist state updates.
    """
    def __init__(self, run_info, store):
        self.store = store
        self.run_info = run_info
        if run_info.artifact_uri:
            self.artifact_repo = ArtifactRepository.from_artifact_uri(run_info.artifact_uri)
        else:
            self.artifact_repo = _get_legacy_artifact_repo(store, run_info)

    def set_terminated(self, status):
        self.run_info = self.store.update_run_info(
            self.run_info.run_uuid, run_status=RunStatus.from_string(status),
            end_time=_get_unix_timestamp())

    def log_metric(self, metric):
        _validate_metric_name(metric.key)
        self.store.log_metric(self.run_info.run_uuid, metric)

    def log_param(self, param):
        _validate_param_name(param.key)
        self.store.log_param(self.run_info.run_uuid, param)

    def log_artifact(self, local_path, artifact_path=None):
        self.artifact_repo.log_artifact(local_path, artifact_path)

    def log_artifacts(self, local_dir, artifact_path=None):
        self.artifact_repo.log_artifacts(local_dir, artifact_path)

    def get_artifact_uri(self):
        return self.artifact_repo.artifact_uri

    def get_run(self):
        return self.store.get_run(self.run_info.run_uuid)

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        status = "FINISHED" if exc_type is None else "FAILED"
        self.set_terminated(status)
        global _active_run
        _active_run = None
        return exc_type is None

def list_experiments():
    """
    Return a list of all experiments.
    """
    return _get_store().list_experiments()

def create_experiment(experiment_name):
    """
    Create an experiment with the specified name and return its ID.
    """
    if experiment_name is None or experiment_name == "":
        raise Exception("Invalid experiment name '%s'" % experiment_name)
    return _get_store().create_experiment(experiment_name)

def _get_main_file():
    if len(sys.argv) > 0:
        return sys.argv[0]
    return None

def _get_source_name():
    main_file = _get_main_file()
    if main_file is not None:
        return main_file
    return "<console>"

def _get_source_version():
    main_file = _get_main_file()
    if main_file is not None:
        return _get_git_commit(main_file)
    return None

def _get_source_type():
    return SourceType.LOCAL

def _get_experiment_id():
    return int(env.get_env(_EXPERIMENT_ID_ENV_VAR) or Experiment.DEFAULT_EXPERIMENT_ID)

def _get_unix_timestamp():
    return int(time.time() * 1000)

def _create_run(experiment_id, source_name, source_version, entry_point_name, source_type):
    store = _get_store()
    run = store.create_run(experiment_id=experiment_id, user_id=_get_user_id(), run_name=None,
                           source_type=source_type,
                           source_name=source_name,
                           entry_point_name=entry_point_name,
                           start_time=_get_unix_timestamp(),
                           source_version=source_version, tags=[])
    return ActiveRun(run.info, store)

def get_run(run_uuid):
    """Return the run with the specified run UUID from the current tracking server."""
    _validate_run_id(run_uuid)
    return _get_store().get_run(run_uuid)

def _get_existing_run(run_uuid):
    _validate_run_id(run_uuid)
    existing_run = get_run(run_uuid)
    if existing_run is None:
        raise Exception("Could not start run with UUID %s - no such run found." % run_uuid)
    store = _get_store()
    updated_info = store.update_run_info(
        run_uuid=run_uuid, run_status=RunStatus.RUNNING, end_time=None)
    return ActiveRun(updated_info, store)

def start_run(run_uuid=None, experiment_id=None, source_name=None, source_version=None,
              entry_point_name=None, source_type=None):
    """
    Start a new MLflow run, setting it as the active run under which metrics and params
    will be logged. The return value can be used as a context manager within a ``with`` block;
    otherwise, ``end_run()`` must be called to terminate the current run. If ``run_uuid``
    is passed or the ``MLFLOW_RUN_ID`` environment variable is set, ``start_run`` attempts to
    resume a run with the specified run ID (with ``run_uuid`` taking precedence over
    ``MLFLOW_RUN_ID``), and other parameters are ignored.

    :param run_uuid: If specified, get the run with the specified UUID and log metrics
                     and params under that run. The run's end time is unset and its status
                     is set to running, but the run's other attributes remain unchanged
                     (the run's ``source_version``, ``source_type``, etc. are not changed).
    :param experiment_id: Used only when ``run_uuid`` is unspecified. ID of the experiment under
                          which to create the current run. If unspecified, the run is created under
                          a new experiment with a randomly generated name.
    :param source_name: Name of the source file or URI of the project to be associated with the run.
                        Defaults to the current file if none provided.
    :param source_version: Optional Git commit hash to associate with the run.
    :param entry_point_name: Optional name of the entry point for to the current run.
    :param source_type: Integer enum value describing the type of the run
                        ("local", "project", etc.). Defaults to
                        ``mlflow.entities.source_type.SourceType.LOCAL``.
    :return: :py:class:`mlflow.tracking.ActiveRun` object that acts as a context manager wrapping
             the run's state.
    """
    global _active_run
    if _active_run:
        raise Exception("Run with UUID %s is already active, unable to start nested "
                        "run" % _active_run.run_info.run_uuid)
    existing_run_uuid = run_uuid or os.environ.get(_RUN_ID_ENV_VAR, None)
    if existing_run_uuid:
        _validate_run_id(existing_run_uuid)
        active_run_obj = _get_existing_run(existing_run_uuid)
    else:
        exp_id_for_run = experiment_id or _get_experiment_id()
        active_run_obj = _create_run(experiment_id=exp_id_for_run,
                                     source_name=source_name or _get_source_name(),
                                     source_version=source_version or _get_source_version(),
                                     entry_point_name=entry_point_name,
                                     source_type=source_type or _get_source_type())
    _active_run = active_run_obj
    return _active_run

def _get_or_start_run():
    if _active_run:
        return _active_run
    return start_run()

def end_run(status="FINISHED"):
    global _active_run
    if _active_run:
        _active_run.set_terminated(status)
        # Clear out the global existing run environment variable as well.
        env.unset_variable(_RUN_ID_ENV_VAR)
        _active_run = None

def active_run():
    """Return the currently active ``Run``, or None if no such run exists."""
    if _active_run:
        return _active_run.get_run()
    else:
        return None

def log_param(key, value):
    """
    Log the passed-in parameter under the current run, creating a run if necessary.

    :param key: Parameter name (string)
    :param value: Parameter value (string)
    """
    _get_or_start_run().log_param(Param(key, str(value)))

def log_metric(key, value):
    """
    Log the passed-in metric under the current run, creating a run if necessary.

    :param key: Metric name (string).
    :param value: Metric value (float).
    """
    if not isinstance(value, numbers.Number):
        print("WARNING: The metric {}={} was not logged because the value is not a number.".format(
            key, value), file=sys.stderr)
        return
    _get_or_start_run().log_metric(Metric(key, value, int(time.time())))

def log_artifact(local_path, artifact_path=None):
    """Log a local file or directory as an artifact of the currently active run."""
    _get_or_start_run().log_artifact(local_path, artifact_path)

def log_artifacts(local_dir, artifact_path=None):
    """Log all the contents of a local directory as artifacts of the run."""
    _get_or_start_run().log_artifacts(local_dir, artifact_path)

def get_artifact_uri():
    """
    Return the artifact URI of the currently active run. Calls to ``log_artifact`` and
    ``log_artifacts`` write artifact(s) to subdirectories of the returned URI.
    """
    return _get_or_start_run().get_artifact_uri()

atexit.register(end_run)

def _get_model_log_dir(model_name, run_id):
    if not run_id:
        raise Exception("Must specify a run_id to get logging directory for a model.")
    run = _get_store().get_run(run_id)
    if run.info.artifact_uri:
        artifact_repo = ArtifactRepository.from_artifact_uri(run.info.artifact_uri)
    else:
        artifact_repo = _get_legacy_artifact_repo(_get_store(), run.info)
    return artifact_repo.download_artifacts(model_name)

def _get_legacy_artifact_repo(file_store, run_info):
    # TODO(aaron) Remove this once everyone locally only has runs from after
    # the introduction of "artifact_uri".
    uri = os.path.join(file_store.root_directory, str(run_info.experiment_id),
                       run_info.run_uuid, "artifacts")
    return ArtifactRepository.from_artifact_uri(uri)

def _get_git_commit(path):
    try:
        from git import Repo, InvalidGitRepositoryError, GitCommandNotFound, NoSuchPathError
    except ImportError as e:
        print("Notice: failed to import Git (the git executable is probably not on your PATH),"
              " so Git SHA is not available. Error: %s" % e, file=sys.stderr)
        return None
    try:
        if os.path.isfile(path):
            path = os.path.dirname(path)
        repo = Repo(path, search_parent_directories=True)
        commit = repo.head.commit.hexsha
        return commit
    except (InvalidGitRepositoryError, GitCommandNotFound, ValueError, NoSuchPathError):
        return None
<|endoftext|>


<commit_before>#esteban@attitude.cl
import sublime, sublime_plugin

#Base
purecss_classes = ["pure-img"]

#Grids
purecss_classes.extend(["pure-g", "pure-u-", "pure-u-sm-", "pure-u-md-", "pure-u-lg-", "pure-u-xl-"])

#Form
purecss_classes.extend(["pure-form", "pure-form-stacked", "pure-form-aligned", "pure-group", "pure-input-", "pure-input-rounded", "pure-checkbox", "pure-radio"])

#Buttons
purecss_classes.extend(["pure-button", "pure-button-disabled", "pure-button-active", "pure-button-primary"])

#Tables
purecss_classes.extend(["pure-table", "pure-table-bordered", "pure-table-horizontal", "pure-table-odd", "pure-table-striped"])

#Menus
purecss_classes.extend(["pure-menu", "pure-menu-horizontal", "pure-menu-selected", "pure-menu-heading", "pure-menu-open", "pure-menu-separator", "pure-menu-disabled", "pure-paginator"])

#Custom. Please Ref: https://github.com/yui/pure/issues/326
purecss_classes.extend(["pure-hidden-sd", "pure-hidden-md", "pure-hidden-lg", "pure-hidden-xl"])

#Hidden in PX
'''
/* pure-hidden-xs */
@media screen and (max-width:567px) {
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-xs{display:none}
}
/* pure-hidden-sm */
@media screen and (min-width:568px) and (max-width:767px) {
    .pure-visible-xs{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-sm{display:none}
}
/* pure-hidden-md */
@media screen and (min-width:768px) and (max-width:1023px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-md{display:none}
}
/* pure-hidden-lg */
@media screen and (min-width:1024px) and (max-width:1279px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-lg{display:none}
}
/* pure-hidden-xl */
@media screen and (min-width:1280px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-hidden-xl{display:none}
}
'''

#Hidden in EM
'''
/* pure-hidden-xs */
@media screen and (max-width:35.438em) {
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-xs{display:none}
}
/* pure-hidden-sm */
@media screen and (min-width:35.5em) and (max-width:47.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-sm{display:none}
}
/* pure-hidden-md */
@media screen and (min-width:48em) and (max-width:63.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-md{display:none}
}
/* pure-hidden-lg */
@media screen and (min-width:64em) and (max-width:79.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-lg{display:none}
}
/* pure-hidden-xl */
@media screen and (min-width:80em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-hidden-xl{display:none}
}
'''

#Attr
purecss_attr = ["hidden", "required", "disabled", "readonly"]

class PureCSSCompletions(sublime_plugin.EventListener):
    def __init__(self):
        self.class_completions = [("%s \tPureCSS class" % s, s) for s in purecss_classes]
        self.attr_completions = [("%s \tPureCSS attr" % s, s) for s in purecss_attr]

    def on_query_completions(self, view, prefix, locations):
        #Thanks to https://github.com/uikit/uikit-sublime for this snippet. All credits to YOOtheme GmbH (MIT licenced)
        if view.match_selector(locations[0], "text.html string.quoted"):
            # make sure it only triggers when: class="{here}"
            pt = locations[0] - len(prefix) - 1
            SEARCH_LIMIT = 250
            search_start = max(0, pt - SEARCH_LIMIT - len(prefix))
            line = view.substr(sublime.Region(search_start, pt))
            parts = line.split('=')
            if len(parts) > 2 and parts[-2].strip().endswith("class"):
                return self.class_completions
            else:
                return []
        elif view.match_selector(locations[0], "text.html meta.tag - text.html punctuation.definition.tag.begin"):
            # in tag: <div {here}>
            return self.attr_completions
        else:
            return []
<commit_msg>add pure-button-group as autocompletable class<commit_after>#esteban@attitude.cl
import sublime, sublime_plugin

#Base
purecss_classes = ["pure-img"]

#Grids
purecss_classes.extend(["pure-g", "pure-u-", "pure-u-sm-", "pure-u-md-", "pure-u-lg-", "pure-u-xl-"])

#Form
purecss_classes.extend(["pure-form", "pure-form-stacked", "pure-form-aligned", "pure-group", "pure-input-", "pure-input-rounded", "pure-checkbox", "pure-radio"])

#Buttons
purecss_classes.extend(["pure-button", "pure-button-disabled", "pure-button-active", "pure-button-primary", "pure-button-group"])

#Tables
purecss_classes.extend(["pure-table", "pure-table-bordered", "pure-table-horizontal", "pure-table-odd", "pure-table-striped"])

#Menus
purecss_classes.extend(["pure-menu", "pure-menu-horizontal", "pure-menu-selected", "pure-menu-heading", "pure-menu-open", "pure-menu-separator", "pure-menu-disabled", "pure-paginator"])

#Custom. Please Ref: https://github.com/yui/pure/issues/326
purecss_classes.extend(["pure-hidden-sd", "pure-hidden-md", "pure-hidden-lg", "pure-hidden-xl"])

#Hidden in PX
'''
/* pure-hidden-xs */
@media screen and (max-width:567px) {
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-xs{display:none}
}
/* pure-hidden-sm */
@media screen and (min-width:568px) and (max-width:767px) {
    .pure-visible-xs{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-sm{display:none}
}
/* pure-hidden-md */
@media screen and (min-width:768px) and (max-width:1023px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-md{display:none}
}
/* pure-hidden-lg */
@media screen and (min-width:1024px) and (max-width:1279px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-lg{display:none}
}
/* pure-hidden-xl */
@media screen and (min-width:1280px) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-hidden-xl{display:none}
}
'''

#Hidden in EM
'''
/* pure-hidden-xs */
@media screen and (max-width:35.438em) {
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-xs{display:none}
}
/* pure-hidden-sm */
@media screen and (min-width:35.5em) and (max-width:47.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-sm{display:none}
}
/* pure-hidden-md */
@media screen and (min-width:48em) and (max-width:63.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-lg{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-md{display:none}
}
/* pure-hidden-lg */
@media screen and (min-width:64em) and (max-width:79.938em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-xl{display:none}
    .pure-hidden-lg{display:none}
}
/* pure-hidden-xl */
@media screen and (min-width:80em) {
    .pure-visible-xs{display:none}
    .pure-visible-sm{display:none}
    .pure-visible-md{display:none}
    .pure-visible-lg{display:none}
    .pure-hidden-xl{display:none}
}
'''

#Attr
purecss_attr = ["hidden", "required", "disabled", "readonly"]

class PureCSSCompletions(sublime_plugin.EventListener):
    def __init__(self):
        self.class_completions = [("%s \tPureCSS class" % s, s) for s in purecss_classes]
        self.attr_completions = [("%s \tPureCSS attr" % s, s) for s in purecss_attr]

    def on_query_completions(self, view, prefix, locations):
        #Thanks to https://github.com/uikit/uikit-sublime for this snippet. All credits to YOOtheme GmbH (MIT licenced)
        if view.match_selector(locations[0], "text.html string.quoted"):
            # make sure it only triggers when: class="{here}"
            pt = locations[0] - len(prefix) - 1
            SEARCH_LIMIT = 250
            search_start = max(0, pt - SEARCH_LIMIT - len(prefix))
            line = view.substr(sublime.Region(search_start, pt))
            parts = line.split('=')
            if len(parts) > 2 and parts[-2].strip().endswith("class"):
                return self.class_completions
            else:
                return []
        elif view.match_selector(locations[0], "text.html meta.tag - text.html punctuation.definition.tag.begin"):
            # in tag: <div {here}>
            return self.attr_completions
        else:
            return []
<|endoftext|>


<commit_before>from io_utilities.base_importData import base_importData
from io_utilities.base_exportData import base_exportData
from .genome_diff_mutations import mutations
from calculate_utilities.base_calculate import base_calculate

class mutations_heatmap(mutations):
    def __init__(self,heatmap_I=[],dendrogram_col_I=[],dendrogram_row_I=[],mutations_I=[]):
        '''
        INPUT:
        mutations_I
        heatmap_I = heatmap data
        dendrogram_I = dendrogram data
        '''
        if mutations_I:
            self.mutations=mutations_I;
        else:
            self.mutations = [];
        if heatmap_I:
            self.heatmap=heatmap_I;
        else:
            self.heatmap = [];
        if dendrogram_col_I:
            self.dendrogram_col=dendrogram_col_I;
        else:
            self.dendrogram_col = [];
        if dendrogram_row_I:
            self.dendrogram_row=dendrogram_row_I;
        else:
            self.dendrogram_row = [];
    def make_heatmap(self, mutation_id_exclusion_list=[],max_position=4000000,
                row_pdist_metric_I='euclidean',row_linkage_method_I='complete',
                col_pdist_metric_I='euclidean',col_linkage_method_I='complete'):
        '''Execute hierarchical cluster on row and column data'''

        print('executing heatmap...');
        calculate = base_calculate();

        # partition into variables:
        mutation_data = self.mutations;
        mutation_data_O = [];
        mutation_ids_all = [];
        for end_cnt,mutation in enumerate(mutations):
            if mutation['mutation_position'] > max_position: #ignore positions great than 4000000
                continue;
            # mutation id
            mutation_genes_str = '';
            mutation_genes_str = self._make_mutationID(mutation['mutation_genes'],mutation['mutation_type'],mutation['mutation_position'])
            tmp = {};
            tmp.update(mutation);
            tmp.update({'mutation_id':mutation_id});
            mutation_data_O.append(tmp);
            mutation_ids_all.append(mutation_id);
        mutation_ids_all_unique = list(set(mutation_ids_all));
        mutation_ids = [x for x in mutation_ids_all_unique if not x in mutation_id_exclusion_list];
        # generate the frequency matrix data structure (mutation x intermediate)
        data_O = numpy.zeros((len(sample_names),len(mutation_ids)));
        samples=[];
        # order 2: groups each sample by mutation (intermediate x mutation)
        for sample_name_cnt,sample_name in enumerate(sample_names): #all samples for intermediate j / mutation i
            samples.append(sample_name); # corresponding label from hierarchical clustering
            for mutation_cnt,mutation in enumerate(mutation_ids): #all mutations i for intermediate j
                for row in mutation_data_O:
                    if row['mutation_id'] == mutation and row['sample_name'] == sample_name:
                        data_O[sample_name_cnt,mutation_cnt] = row['mutation_frequency'];
        # generate the clustering for the heatmap
        heatmap_O = [];
        dendrogram_col_O = {};
        dendrogram_row_O = {};
        heatmap_O,dendrogram_col_O,dendrogram_row_O = self.calculate.heatmap(data_O,samples,mutation_ids,
                row_pdist_metric_I=row_pdist_metric_I,row_linkage_method_I=row_linkage_method_I,
                col_pdist_metric_I=col_pdist_metric_I,col_linkage_method_I=col_linkage_method_I);
        # record the data
        self.heatmap = heatmap_O;
        self.dendrogram_col = dendrogram_col_O;
        self.dendrogram_row = dendrogram_row_O;

    def _make_mutationID(self,mutation_genes,mutation_type,mutation_position):
        '''return a unique mutation id string'''
        mutation_genes_str = '';
        for gene in mutation_genes:
            mutation_genes_str = mutation_genes_str + gene + '-/-'
        mutation_genes_str = mutation_genes_str[:-3];
        mutation_id = mutation_type + '_' + mutation_genes_str + '_' + str(mutation_position);
        return mutation_genes_str;

    def clear_data(self):
        del self.mutations[:];
        del self.heatmap[:];
        del self.dendrogram[:];<commit_msg>updates to heatmap<commit_after>from io_utilities.base_importData import base_importData
from io_utilities.base_exportData import base_exportData
from .genome_diff_mutations import mutations
from calculate_utilities.base_calculate import base_calculate
import numpy

class mutations_heatmap(mutations):
    def __init__(self,heatmap_I=[],dendrogram_col_I=[],dendrogram_row_I=[],mutations_I=[],sample_names_I=[]):
        '''
        INPUT:
        mutations_I
        sample_names_I
        heatmap_I = heatmap data
        dendrogram_I = dendrogram data
        '''
        if mutations_I:
            self.mutations=mutations_I;
        else:
            self.mutations = [];
        if sample_names_I:
            self.sample_names=sample_names_I;
        else:
            self.sample_names = [];
        if heatmap_I:
            self.heatmap=heatmap_I;
        else:
            self.heatmap = [];
        if dendrogram_col_I:
            self.dendrogram_col=dendrogram_col_I;
        else:
            self.dendrogram_col = [];
        if dendrogram_row_I:
            self.dendrogram_row=dendrogram_row_I;
        else:
            self.dendrogram_row = [];
    def make_heatmap(self, mutation_id_exclusion_list=[],max_position=4000000,
                row_pdist_metric_I='euclidean',row_linkage_method_I='complete',
                col_pdist_metric_I='euclidean',col_linkage_method_I='complete'):
        '''Execute hierarchical cluster on row and column data'''

        print('executing heatmap...');
        calculate = base_calculate();

        # partition into variables:
        mutation_data = self.mutations;
        sample_names = self.sample_names;
        mutation_data_O = [];
        mutation_ids_all = [];
        for end_cnt,mutation in enumerate(mutation_data):
            if mutation['mutation_position'] > max_position: #ignore positions great than 4000000
                continue;
            # mutation id
            mutation_id = '';
            mutation_id = self._make_mutationID(mutation['mutation_genes'],mutation['mutation_type'],mutation['mutation_position'])
            tmp = {};
            tmp.update(mutation);
            tmp.update({'mutation_id':mutation_id});
            mutation_data_O.append(tmp);
            mutation_ids_all.append(mutation_id);
        mutation_ids_all_unique = list(set(mutation_ids_all));
        mutation_ids = [x for x in mutation_ids_all_unique if not x in mutation_id_exclusion_list];
        # generate the frequency matrix data structure (mutation x intermediate)
        data_O = numpy.zeros((len(sample_names),len(mutation_ids)));
        samples=[];
        # order 2: groups each sample by mutation (intermediate x mutation)
        for sample_name_cnt,sample_name in enumerate(sample_names): #all samples for intermediate j / mutation i
            samples.append(sample_name); # corresponding label from hierarchical clustering
            for mutation_cnt,mutation in enumerate(mutation_ids): #all mutations i for intermediate j
                for row in mutation_data_O:
                    if row['mutation_id'] == mutation and row['sample_name'] == sample_name:
                        data_O[sample_name_cnt,mutation_cnt] = row['mutation_frequency'];
        # generate the clustering for the heatmap
        heatmap_O = [];
        dendrogram_col_O = {};
        dendrogram_row_O = {};
        heatmap_O,dendrogram_col_O,dendrogram_row_O = calculate.heatmap(data_O,samples,mutation_ids,
                row_pdist_metric_I=row_pdist_metric_I,row_linkage_method_I=row_linkage_method_I,
                col_pdist_metric_I=col_pdist_metric_I,col_linkage_method_I=col_linkage_method_I);
        # record the data
        self.heatmap = heatmap_O;
        self.dendrogram_col = dendrogram_col_O;
        self.dendrogram_row = dendrogram_row_O;

    def _make_mutationID(self,mutation_genes,mutation_type,mutation_position):
        '''return a unique mutation id string'''
        mutation_genes_str = '';
        for gene in mutation_genes:
            mutation_genes_str = mutation_genes_str + gene + '-/-'
        mutation_genes_str = mutation_genes_str[:-3];
        mutation_id = mutation_type + '_' + mutation_genes_str + '_' + str(mutation_position);
        return mutation_id;

    def clear_data(self):
        del self.mutations[:];
        del self.heatmap[:];
        del self.dendrogram[:];<|endoftext|>


<commit_before>"""
python -m lifelines.tests.test_suit

"""
import unittest
import numpy as np
import numpy.testing as npt
from collections import Counter
import matplotlib.pyplot as plt
import pdb

import pandas as pd

from ..estimation import KaplanMeierFitter, NelsonAalenFitter, AalenAdditiveFitter
from ..statistics import logrank_test
from ..generate_datasets import *
from ..plotting import plot_lifetimes
from ..utils import datetimes_to_durations

LIFETIMES = np.array([2,4,4,4,5,7,10,11,11,12])
CENSORSHIP = np.array([1,1,0,1,0,1,1,1,1,0])

waltonT2 = np.array([6., 13., 13., 13., 19., 19., 19., 26., 26., 26., 26., 26., 33., 33., 47., 62., 62., 9., 9., 9., 15., 15., 22., 22., 22., 22., 29., 29., 29., 29., 29., 36., 36., 43.])
waltonT1 = np.array([33., 54., 54., 61., 61., 61., 61., 61., 61., 61., 61., 61., 61., 61., 69., 69., 69., 69., 69., 69., 69., 69., 69., 69., 69., 32., 53., 53., 60., 60., 60., 60., 60., 68., 68., 68., 68., 68., 68., 68., 68., 68., 68., 75., 17., 51., 51., 51., 58., 58., 58., 58., 66., 66., 7., 7., 41., 41., 41., 41., 41., 41., 41., 48., 48., 48., 48., 48., 48., 48., 48., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 63., 63., 63., 63., 63., 63., 63., 63., 63., 69., 69., 38., 38., 45., 45., 45., 45., 45., 45., 45., 45., 45., 45., 53., 53., 53., 53., 53., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 66.])

N = len(LIFETIMES)

class MiscTests(unittest.TestCase):

    def test_quadratic_integration_identity(self):
      #integrate x between 0 and 10:
      x = np.linspace(0,10,1000)[:,None]
      answer = 0.5*x**2
      approx_answer = cumulative_quadrature( x.T, x).T
      npt.assert_almost_equal(answer, approx_answer, decimal=5)

    def test_quadratic_integration_exp(self):
      #integrate exp(x) between 0 and 4:
      x = np.linspace(0,4,1000)[:,None]
      answer = np.exp(x) -1.
      approx_answer = cumulative_quadrature( np.exp(x).T, x).T
      #pdb.set_trace()
      npt.assert_almost_equal(answer, approx_answer, decimal=4)

    def test_aalen_additive_allows_numpy_or_df(self):
      t = np.random.random((10,1)).T
      dfX = pd.DataFrame(np.random.random((10,3)), columns = ["A", "B", "C"])
      c = dfX.columns
      npX = np.random.random((10,5))
      aaf = AalenAdditiveFitter()
      aaf.fit(t,npX)
      aaf.fit(t, dfX)
      npt.assert_array_equal( c, aaf.cumulative_hazards_.columns[:3] )

    def test_datetimes_to_durations_days(self):
        start_date = ['2013-10-10 0:00:00', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', '2013-10-10 0:00:00', '2013-10-15']
        T,C = datetimes_to_durations(start_date, end_date)
        npt.assert_almost_equal(T, np.array([3,1,5+365]) )
        npt.assert_almost_equal(C, np.array([1,1,1], dtype=bool) )
        return

    def test_datetimes_to_durations_years(self):
        start_date = ['2013-10-10', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', '2013-10-10', '2013-10-15']
        T,C = datetimes_to_durations(start_date, end_date, freq='Y')
        npt.assert_almost_equal(T, np.array([0,0,1]) )
        npt.assert_almost_equal(C, np.array([1,1,1], dtype=bool) )
        return

    def test_datetimes_to_durations_censor(self):
        start_date = ['2013-10-10', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', None, '']
        T,C = datetimes_to_durations(start_date, end_date, freq='Y')
        npt.assert_almost_equal(C, np.array([1,0,0], dtype=bool) )
        return

class StatisticalTests(unittest.TestCase):

  def setUp(self):
      self.lifetimes = Counter(LIFETIMES)
      self.km = self.kaplan_meier()
      self.kmc = self.kaplan_meier(censor=True)
      self.na = self.nelson_aalen()
      self.nac = self.nelson_aalen(censor=True)

  def test_kaplan_meier(self):
      kmf = KaplanMeierFitter()
      kmf.fit(LIFETIMES)
      npt.assert_almost_equal(kmf.survival_function_.values, self.km )

  def test_nelson_aalen(self):
      naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
      naf.fit(LIFETIMES)
      npt.assert_almost_equal(naf.cumulative_hazard_.values, self.na )

  def test_censor_nelson_aalen(self):
      naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
      naf.fit(LIFETIMES, censorship=CENSORSHIP)
      npt.assert_almost_equal(naf.cumulative_hazard_.values, self.nac )

  def test_censor_kaplan_meier(self):
      kmf = KaplanMeierFitter()
      kmf.fit(LIFETIMES, censorship = CENSORSHIP)
      npt.assert_almost_equal(kmf.survival_function_.values, self.kmc )

  def test_not_to_break(self):
      try:
        naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
        naf.fit(LIFETIMES)
        naf = NelsonAalenFitter(nelson_aalen_smoothing=True)
        naf.fit(LIFETIMES)
        self.assertTrue(1==1)
      except Exception as e:
        print e
        self.assertTrue(0==1)

  def test_equal_intensity(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(5, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result==None)

  def test_unequal_intensity(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result)

  def test_unequal_intensity_censorship(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      censorA = np.random.binomial(1,0.5, size=(200,1))
      censorB = np.random.binomial(1,0.5, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2, censorship_A = censorA, censorship_B = censorB)
      print summary
      self.assertTrue(result)

  def test_integer_times_logrank_test(self):
      data1 = np.random.exponential(5, size=(200,1)).astype(int)
      data2 = np.random.exponential(1, size=(200,1)).astype(int)
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result)

  def test_waltons_data(self):

      summary, p_value, result = logrank_test(waltonT1, waltonT2)
      print summary
      self.assertTrue(result)

  def kaplan_meier(self, censor=False):
      km = np.zeros((len(self.lifetimes.keys()),1))
      ordered_lifetimes = np.sort( self.lifetimes.keys())
      v = 1.
      n = N*1.0
      for i,t in enumerate(ordered_lifetimes):
        if censor:
           ix = LIFETIMES == t
           c = sum(1-CENSORSHIP[ix])
           n -=  c
           if n!=0:
              v *= ( 1-(self.lifetimes.get(t)-c)/n )
           n -= self.lifetimes.get(t) - c
        else:
           v *= ( 1-self.lifetimes.get(t)/n )
           n -= self.lifetimes.get(t)
        km[i] = v
      if km[0] < 1.:
        km = np.insert(km,0,1.)
      return km.reshape(len(km),1)

  def nelson_aalen(self, censor = False):
      na = np.zeros((len(self.lifetimes.keys()),1))
      ordered_lifetimes = np.sort( self.lifetimes.keys())
      v = 0.
      n = N*1.0
      for i,t in enumerate(ordered_lifetimes):
        if censor:
           ix = LIFETIMES == t
           c = sum(1-CENSORSHIP[ix])
           n -=  c
           if n!=0:
              v += ( (self.lifetimes.get(t)-c)/n )
           n -= self.lifetimes.get(t) - c
        else:
           v += ( self.lifetimes.get(t)/n )
           n -= self.lifetimes.get(t)
        na[i] = v
      if na[0] > 0:
        na = np.insert(na,0,0.)
      return na.reshape(len(na),1)

class PlottingTests(unittest.TestCase):

  def test_kmf_plotting(self):
      data1 = np.random.exponential(10, size=(200,1))
      data2 = np.random.exponential(2, size=(200,1))
      kmf = KaplanMeierFitter()
      kmf.fit(data1)
      ax = kmf.plot()
      kmf.fit(data2)
      kmf.plot(ax=ax, c="#A60628")
      return

  def test_naf_plotting(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      naf = NelsonAalenFitter()
      naf.fit(data1)
      ax = naf.plot()
      naf.fit(data2)
      naf.plot(ax=ax, c="#A60628")
      return

  def test_plot_lifetimes_calendar(self):
      plt.figure()
      t = np.linspace(0, 20, 1000)
      hz, coef, covrt =generate_hazard_rates(1,5, t)
      N = 20
      current = 10
      birthtimes = current*np.random.uniform(size=(N,))
      T, C= generate_random_lifetimes(hz, t, size=N, censor=current - birthtimes )
      plot_lifetimes(T, censorship=C, birthtimes=birthtimes)

  def test_plot_lifetimes_relative(self):
      plt.figure()
      t = np.linspace(0, 20, 1000)
      hz, coef, covrt =generate_hazard_rates(1,5, t)
      N = 20
      T, C= generate_random_lifetimes(hz, t, size=N, censor=True )
      plot_lifetimes(T, censorship=C)

if __name__ == '__main__':
    unittest.main()
<commit_msg>adding multivariate unittests<commit_after>"""
python -m lifelines.tests.test_suit

"""
import unittest
import numpy as np
import numpy.testing as npt
from collections import Counter
import matplotlib.pyplot as plt
import pdb

import pandas as pd

from ..estimation import KaplanMeierFitter, NelsonAalenFitter, AalenAdditiveFitter
from ..statistics import logrank_test,multi_logrank_test
from ..generate_datasets import *
from ..plotting import plot_lifetimes
from ..utils import datetimes_to_durations

LIFETIMES = np.array([2,4,4,4,5,7,10,11,11,12])
CENSORSHIP = np.array([1,1,0,1,0,1,1,1,1,0])

waltonT2 = np.array([6., 13., 13., 13., 19., 19., 19., 26., 26., 26., 26., 26., 33., 33., 47., 62., 62., 9., 9., 9., 15., 15., 22., 22., 22., 22., 29., 29., 29., 29., 29., 36., 36., 43.])
waltonT1 = np.array([33., 54., 54., 61., 61., 61., 61., 61., 61., 61., 61., 61., 61., 61., 69., 69., 69., 69., 69., 69., 69., 69., 69., 69., 69., 32., 53., 53., 60., 60., 60., 60., 60., 68., 68., 68., 68., 68., 68., 68., 68., 68., 68., 75., 17., 51., 51., 51., 58., 58., 58., 58., 66., 66., 7., 7., 41., 41., 41., 41., 41., 41., 41., 48., 48., 48., 48., 48., 48., 48., 48., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 56., 63., 63., 63., 63., 63., 63., 63., 63., 63., 69., 69., 38., 38., 45., 45., 45., 45., 45., 45., 45., 45., 45., 45., 53., 53., 53., 53., 53., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 60., 66.])

N = len(LIFETIMES)

class MiscTests(unittest.TestCase):

    def test_quadratic_integration_identity(self):
      #integrate x between 0 and 10:
      x = np.linspace(0,10,1000)[:,None]
      answer = 0.5*x**2
      approx_answer = cumulative_quadrature( x.T, x).T
      npt.assert_almost_equal(answer, approx_answer, decimal=5)

    def test_quadratic_integration_exp(self):
      #integrate exp(x) between 0 and 4:
      x = np.linspace(0,4,1000)[:,None]
      answer = np.exp(x) -1.
      approx_answer = cumulative_quadrature( np.exp(x).T, x).T
      #pdb.set_trace()
      npt.assert_almost_equal(answer, approx_answer, decimal=4)

    def test_aalen_additive_allows_numpy_or_df(self):
      t = np.random.random((10,1)).T
      dfX = pd.DataFrame(np.random.random((10,3)), columns = ["A", "B", "C"])
      c = dfX.columns
      npX = np.random.random((10,5))
      aaf = AalenAdditiveFitter()
      aaf.fit(t,npX)
      aaf.fit(t, dfX)
      npt.assert_array_equal( c, aaf.cumulative_hazards_.columns[:3] )

    def test_datetimes_to_durations_days(self):
        start_date = ['2013-10-10 0:00:00', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', '2013-10-10 0:00:00', '2013-10-15']
        T,C = datetimes_to_durations(start_date, end_date)
        npt.assert_almost_equal(T, np.array([3,1,5+365]) )
        npt.assert_almost_equal(C, np.array([1,1,1], dtype=bool) )
        return

    def test_datetimes_to_durations_years(self):
        start_date = ['2013-10-10', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', '2013-10-10', '2013-10-15']
        T,C = datetimes_to_durations(start_date, end_date, freq='Y')
        npt.assert_almost_equal(T, np.array([0,0,1]) )
        npt.assert_almost_equal(C, np.array([1,1,1], dtype=bool) )
        return

    def test_datetimes_to_durations_censor(self):
        start_date = ['2013-10-10', '2013-10-09', '2012-10-10']
        end_date = ['2013-10-13', None, '']
        T,C = datetimes_to_durations(start_date, end_date, freq='Y')
        npt.assert_almost_equal(C, np.array([1,0,0], dtype=bool) )
        return

class StatisticalTests(unittest.TestCase):

  def setUp(self):
      self.lifetimes = Counter(LIFETIMES)
      self.km = self.kaplan_meier()
      self.kmc = self.kaplan_meier(censor=True)
      self.na = self.nelson_aalen()
      self.nac = self.nelson_aalen(censor=True)

  def test_kaplan_meier(self):
      kmf = KaplanMeierFitter()
      kmf.fit(LIFETIMES)
      npt.assert_almost_equal(kmf.survival_function_.values, self.km )

  def test_nelson_aalen(self):
      naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
      naf.fit(LIFETIMES)
      npt.assert_almost_equal(naf.cumulative_hazard_.values, self.na )

  def test_censor_nelson_aalen(self):
      naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
      naf.fit(LIFETIMES, censorship=CENSORSHIP)
      npt.assert_almost_equal(naf.cumulative_hazard_.values, self.nac )

  def test_censor_kaplan_meier(self):
      kmf = KaplanMeierFitter()
      kmf.fit(LIFETIMES, censorship = CENSORSHIP)
      npt.assert_almost_equal(kmf.survival_function_.values, self.kmc )

  def test_not_to_break(self):
      try:
        naf = NelsonAalenFitter(nelson_aalen_smoothing=False)
        naf.fit(LIFETIMES)
        naf = NelsonAalenFitter(nelson_aalen_smoothing=True)
        naf.fit(LIFETIMES)
        self.assertTrue(1==1)
      except Exception as e:
        print e
        self.assertTrue(0==1)

  def test_equal_intensity(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(5, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result==None)

  def test_unequal_intensity(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result)

  def test_unequal_intensity_censorship(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      censorA = np.random.binomial(1,0.5, size=(200,1))
      censorB = np.random.binomial(1,0.5, size=(200,1))
      summary, p_value, result = logrank_test(data1, data2, censorship_A = censorA, censorship_B = censorB)
      print summary
      self.assertTrue(result)

  def test_integer_times_logrank_test(self):
      data1 = np.random.exponential(5, size=(200,1)).astype(int)
      data2 = np.random.exponential(1, size=(200,1)).astype(int)
      summary, p_value, result = logrank_test(data1, data2)
      print summary
      self.assertTrue(result)

  def test_waltons_data(self):

      summary, p_value, result = logrank_test(waltonT1, waltonT2)
      print summary
      self.assertTrue(result)

  def test_multivariate_unequal_intensities(self):
     T = np.random.exponential(10, size=300)
     g = np.random.binomial(2, 0.5, size=300)
     T[g==1] = np.random.exponential(6, size=(g==1).sum())
     s, _, result = multi_logrank_test(T, g)
     print s
     self.assertTrue(result==True)

  def test_multivariate_equal_intensities(self):
     T = np.random.exponential(10, size=300)
     g = np.random.binomial(2, 0.5, size=300)
     s, _, result = multi_logrank_test(T, g)
     print result
     self.assertTrue(result==None)

  def kaplan_meier(self, censor=False):
      km = np.zeros((len(self.lifetimes.keys()),1))
      ordered_lifetimes = np.sort( self.lifetimes.keys())
      v = 1.
      n = N*1.0
      for i,t in enumerate(ordered_lifetimes):
        if censor:
           ix = LIFETIMES == t
           c = sum(1-CENSORSHIP[ix])
           n -=  c
           if n!=0:
              v *= ( 1-(self.lifetimes.get(t)-c)/n )
           n -= self.lifetimes.get(t) - c
        else:
           v *= ( 1-self.lifetimes.get(t)/n )
           n -= self.lifetimes.get(t)
        km[i] = v
      if km[0] < 1.:
        km = np.insert(km,0,1.)
      return km.reshape(len(km),1)

  def nelson_aalen(self, censor = False):
      na = np.zeros((len(self.lifetimes.keys()),1))
      ordered_lifetimes = np.sort( self.lifetimes.keys())
      v = 0.
      n = N*1.0
      for i,t in enumerate(ordered_lifetimes):
        if censor:
           ix = LIFETIMES == t
           c = sum(1-CENSORSHIP[ix])
           n -=  c
           if n!=0:
              v += ( (self.lifetimes.get(t)-c)/n )
           n -= self.lifetimes.get(t) - c
        else:
           v += ( self.lifetimes.get(t)/n )
           n -= self.lifetimes.get(t)
        na[i] = v
      if na[0] > 0:
        na = np.insert(na,0,0.)
      return na.reshape(len(na),1)

class PlottingTests(unittest.TestCase):

  def test_kmf_plotting(self):
      data1 = np.random.exponential(10, size=(200,1))
      data2 = np.random.exponential(2, size=(200,1))
      kmf = KaplanMeierFitter()
      kmf.fit(data1)
      ax = kmf.plot()
      kmf.fit(data2)
      kmf.plot(ax=ax, c="#A60628")
      return

  def test_naf_plotting(self):
      data1 = np.random.exponential(5, size=(200,1))
      data2 = np.random.exponential(1, size=(200,1))
      naf = NelsonAalenFitter()
      naf.fit(data1)
      ax = naf.plot()
      naf.fit(data2)
      naf.plot(ax=ax, c="#A60628")
      return

  def test_plot_lifetimes_calendar(self):
      plt.figure()
      t = np.linspace(0, 20, 1000)
      hz, coef, covrt =generate_hazard_rates(1,5, t)
      N = 20
      current = 10
      birthtimes = current*np.random.uniform(size=(N,))
      T, C= generate_random_lifetimes(hz, t, size=N, censor=current - birthtimes )
      plot_lifetimes(T, censorship=C, birthtimes=birthtimes)

  def test_plot_lifetimes_relative(self):
      plt.figure()
      t = np.linspace(0, 20, 1000)
      hz, coef, covrt =generate_hazard_rates(1,5, t)
      N = 20
      T, C= generate_random_lifetimes(hz, t, size=N, censor=True )
      plot_lifetimes(T, censorship=C)

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


<commit_before>__title__ = 'MT940'
__package_name__ = 'mt-940'
__author__ = 'Rick van Hattem (wolph)'
__description__ = ' '.join('''
A library to parse MT940 files and returns smart Python collections for
statistics and manipulation.
'''.strip().split())
__email__ = 'wolph@wol.ph'
__version__ = '4.22.0'
__license__ = 'BSD'
__copyright__ = 'Copyright 2015 Rick van Hattem (wolph)'
__url__ = 'https://github.com/WoLpH/mt940'
<commit_msg>Incrementing version to v4.23.0<commit_after>__title__ = 'MT940'
__package_name__ = 'mt-940'
__author__ = 'Rick van Hattem (wolph)'
__description__ = ' '.join('''
A library to parse MT940 files and returns smart Python collections for
statistics and manipulation.
'''.strip().split())
__email__ = 'wolph@wol.ph'
__version__ = '4.23.0'
__license__ = 'BSD'
__copyright__ = 'Copyright 2015 Rick van Hattem (wolph)'
__url__ = 'https://github.com/WoLpH/mt940'
<|endoftext|>


<commit_before>'''
This module contains all fo the routines needed to set up a master server, this
involves preparing the three listeners and the workers needed by the master.
'''
# Import python modules
import os
import random
import multiprocessing
# Import zeromq
import zmq
# Import salt modules
import salt.utils

class Publisher(multiprocessing.Process):
    '''
    The publihing interface, a simple zeromq publisher that sends out the
    commands.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUB)

    def __bind(self):
        '''
        Bind to the interface specified in the configuration file
        '''
        binder = 'tcp://' + self.opts['interface'] + ':'\
               + self.opts['publish_port']
        self.socket.bind(binder)

    def command(self, cmd):
        '''
        Publish out a command to the minions, takes a cmd structure, one of
        those encrypted json jobs
        '''
        self.socket.send(cmd)

    def run(self):
        '''
        Start the publisher
        '''
        self.__bind()

class ReqServer(multiprocessing.Process):
    '''
    Starts up a multiproceesing request server, minions send results to this
    interface.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.num_threads = self.opts['worker_threads']
        self.context = zmq.Context(1)

        self.c_uri = 'tcp://' + self.opts['interface'] + ':'\
                   + self.opts['ret_port']
        self.clients = self.context.socket(zmq.XREP)

        self.workers = self.context.socket(zmq.XREQ)
        self.w_uri = 'tcp://localhost:' + self.opts['worker_port']

    def __worker(self):
        '''
        Starts up a worker multiprocess
        '''
        socket = self.context.socket(zmq.REP)
        socket.connect(self.w_uri)

        while True:
            message = socket.recv()
            salt.utils.pickle_message(message)
            socket.send(1)

    def __bind(self):
        '''
        Binds the reply server
        '''
        self.clients.bind(self.c_uri)

        self.workers.bind(self.w_uri)

        for ind in range(int(self.num_threads)):
            multiprocessing.Process(target=self.__worker).start()

        zmq.device(zmq.QUEUE, self.clients, self.workers)

    def run(self):
        '''
        Start up the ReqServer
        '''
        self.__bind()

class LocalServer(ReqServer):
    '''
    Create the localhost communication interface for root clients to connect
    to
    '''
    def __init__(self, opts):
        ReqServer.__init__(self, opts)
        self.num_threads = self.opts['local_threads']
        self.publisher = Publisher(opts)
        self.publisher.start()
        self.c_uri = 'tcp://localhost:' + self.opts['local_port']
        self.w_uri = 'tcp://localhost:' + self.opts['local_worker_port']
        self.key = self.__prep_key()

    def __prep_key(self):
        '''
        A key needs to be placed in the filesystem with permissions 0400 so
        clients are required to run as root.
        '''
        keyfile = os.path.join(self.opts['cachedir'], 'root_key')
        key = random.randint(100000000000000000000000, 999999999999999999999999)
        open(keyfile, 'w+').write(key)
        return key

    def __worker(self):
        '''
        A localserver worker needs to run some extra checks
        '''
        socket = self.context.socket(zmq.REP)
        socket.connect(self.w_uri)

        while True:
            message = socket.recv()
            if not message.has_key('key'):
                socket.send(0) # Reply false
                continue
            if not message.has_key('cmd'):
                socket.send(0) # Reply false
                continue
            if not message['key'] == self.key:
                socket.send(0) # Reply false
                continue
            cmd = salt.utils.prep_publish_cmd(message['cmd'])
            self.publisher.command(cmd)
<commit_msg>Salt master initialization hooks<commit_after>'''
This module contains all fo the routines needed to set up a master server, this
involves preparing the three listeners and the workers needed by the master.
'''
# Import python modules
import os
import random
import multiprocessing
# Import zeromq
import zmq
# Import salt modules
import salt.utils

class Master(object):
    '''
    The salt master server
    '''
    def __init__(self, opts):
        '''
        Create a salt master server instance
        '''
        self.opts = opts

    def start(self):
        '''
        Turn on the master server components
        '''
        publister = Publisher(opts)
        reqserv = ReqServer(opts)
        local = LocalServer(opts)
        publister.start()
        reqserv.start()
        local.start()

class Publisher(multiprocessing.Process):
    '''
    The publihing interface, a simple zeromq publisher that sends out the
    commands.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUB)

    def __bind(self):
        '''
        Bind to the interface specified in the configuration file
        '''
        binder = 'tcp://' + self.opts['interface'] + ':'\
               + self.opts['publish_port']
        self.socket.bind(binder)

    def command(self, cmd):
        '''
        Publish out a command to the minions, takes a cmd structure, one of
        those encrypted json jobs
        '''
        self.socket.send(cmd)

    def run(self):
        '''
        Start the publisher
        '''
        self.__bind()

class ReqServer(multiprocessing.Process):
    '''
    Starts up a multiproceesing request server, minions send results to this
    interface.
    '''
    def __init__(self, opts):
        multiprocessing.Process.__init__(self)
        self.opts = opts
        self.num_threads = self.opts['worker_threads']
        self.context = zmq.Context(1)

        self.c_uri = 'tcp://' + self.opts['interface'] + ':'\
                   + self.opts['ret_port']
        self.clients = self.context.socket(zmq.XREP)

        self.workers = self.context.socket(zmq.XREQ)
        self.w_uri = 'tcp://localhost:' + self.opts['worker_port']

    def __worker(self):
        '''
        Starts up a worker multiprocess
        '''
        socket = self.context.socket(zmq.REP)
        socket.connect(self.w_uri)

        while True:
            message = socket.recv()
            salt.utils.pickle_message(message)
            socket.send(1)

    def __bind(self):
        '''
        Binds the reply server
        '''
        self.clients.bind(self.c_uri)

        self.workers.bind(self.w_uri)

        for ind in range(int(self.num_threads)):
            multiprocessing.Process(target=self.__worker).start()

        zmq.device(zmq.QUEUE, self.clients, self.workers)

    def run(self):
        '''
        Start up the ReqServer
        '''
        self.__bind()

class LocalServer(ReqServer):
    '''
    Create the localhost communication interface for root clients to connect
    to
    '''
    def __init__(self, opts):
        ReqServer.__init__(self, opts)
        self.num_threads = self.opts['local_threads']
        self.publisher = Publisher(opts)
        self.publisher.start()
        self.c_uri = 'tcp://localhost:' + self.opts['local_port']
        self.w_uri = 'tcp://localhost:' + self.opts['local_worker_port']
        self.key = self.__prep_key()

    def __prep_key(self):
        '''
        A key needs to be placed in the filesystem with permissions 0400 so
        clients are required to run as root.
        '''
        keyfile = os.path.join(self.opts['cachedir'], 'root_key')
        key = random.randint(100000000000000000000000, 999999999999999999999999)
        open(keyfile, 'w+').write(key)
        return key

    def __worker(self):
        '''
        A localserver worker needs to run some extra checks
        '''
        socket = self.context.socket(zmq.REP)
        socket.connect(self.w_uri)

        while True:
            message = socket.recv()
            if not message.has_key('key'):
                socket.send(0) # Reply false
                continue
            if not message.has_key('cmd'):
                socket.send(0) # Reply false
                continue
            if not message['key'] == self.key:
                socket.send(0) # Reply false
                continue
            cmd = salt.utils.prep_publish_cmd(message['cmd'])
            self.publisher.command(cmd)
<|endoftext|>


<commit_before># -*- test-case-name: vumi.transports.smpp.test_smpp -*-

from datetime import datetime

import redis
from twisted.python import log
from twisted.internet import reactor

# from vumi.service import Worker
from vumi.utils import get_operator_number, get_deploy_int
from vumi.transports.base import Transport
from vumi.transports.smpp.client import EsmeTransceiverFactory
from vumi.message import Message

class SmppTransport(Transport):
    """
    An SMPP transport.

    The SMPP transport has many configuration parameters. These are
    divided up into sections below.

    SMPP sequence number configuration options:

    :type smpp_increment: int
    :param smpp_increment:
        Increment for SMPP sequence number (must be >= number of
        SMPP workers on a single SMPP account).
    :type smpp_offset: int
    :param smpp_offset:
        Offset for this worker's SMPP sequence numbers (no duplicates
        on a single SMPP account and must be <= increment)

    SMPP server account configuration options:

    :type system_id: str
    :param system_id:
        User id used to connect to the SMPP server.
    :type password: str
    :param password:
        Password for the system id.
    :type system_type: str, optional
    :param system_type:
        Additional system metadata that is passed through to the SMPP server
        on connect.
    :type host: str
    :param host:
        Hostname of the SMPP server.
    :type port: int
    :param port:
        Port the SMPP server is listening on.
    :type initial_reconnect_delay: int, optional
    :param initial_reconnect_delay:
        Number of seconds to delay before reconnecting to the server after
        being disconnected. Default is 5s. Some WASPs, e.g. Clickatell,
        require a 30s delay before reconnecting. In these cases a 45s
        initial_reconnect_delay is recommended.

    SMPP protocol configuration options:

    :type interface_version: str, optional
    :param interface_version:
        SMPP protocol version. Default is '34' (i.e. version 3.4).
    :type dest_addr_ton:
    :param dest_addr_ton:
        Destination TON (type of number). Default .
    :type dest_addr_npi:
    :param dest_addr_npi:
        Destination NPI (number plan identifier). Default 1 (ISDN/E.164/E.163).
    :type registered_delivery:
    :param registered_delivery:
        Whether to ask for delivery reports. Default 1 (request delivery
        reports).

    The list of SMPP protocol configuration options given above is not
    exhaustive. Any other options specified are passed through to the
    python-smpp library PDU (protocol data unit) builder.

    Cellphone number routing options:

    :type COUNTRY_CODE: str, optional
    :param COUNTRY_CODE:
        Used to translate a leading zero in a destination MSISDN into a
        country code. Default '',
    :type OPERATOR_PREFIX: str, optional
    :param OPERATOR_PREFIX:
        Nested dictionary of prefix to network name mappings. Default {} (set
        network to 'UNKNOWN'). E.g. { '27': { '27761': 'NETWORK1' }}.
    :type OPERATOR_NUMBER:
    :param OPERATOR_NUMBER:
        Dictionary of source MSISDN to use for each network listed in
        OPERATOR_PREFIX. If a network is not listed, the source MSISDN
        specified by the message sender is used. Default {} (always used the
        from address specified by the message sender). E.g. { 'NETWORK1':
        '27761234567'}.
    """

    def setup_transport(self):
        log.msg("Starting the SmppTransport with %s" % self.config)

        # TODO: move this to a config file
        dbindex = get_deploy_int(self._amqp_client.vhost)
        self.smpp_offset = int(self.config['smpp_offset'])

        # Connect to Redis
        if not hasattr(self, 'r_server'):
            # Only set up redis if we don't have a test stub already
            self.r_server = redis.Redis("localhost", db=dbindex)
        self.r_prefix = "%(system_id)s@%(host)s:%(port)s" % self.config
        log.msg("Connected to Redis, prefix: %s" % self.r_prefix)
        last_sequence_number = int(self.r_get_last_sequence()
                                   or self.smpp_offset)
        log.msg("Last sequence_number: %s" % last_sequence_number)

        if not hasattr(self, 'esme_client'):
            # start the Smpp transport (if we don't have one)
            factory = EsmeTransceiverFactory(self.config,
                                             self._amqp_client.vumi_options)
            factory.loadDefaults(self.config)
            factory.setLastSequenceNumber(last_sequence_number)
            factory.setConnectCallback(self.esme_connected)
            factory.setDisconnectCallback(self.esme_disconnected)
            factory.setSubmitSMRespCallback(self.submit_sm_resp)
            factory.setDeliveryReportCallback(self.delivery_report)
            factory.setDeliverSMCallback(self.deliver_sm)
            factory.setSendFailureCallback(self.send_failure)
            log.msg(factory.defaults)
            reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.update_error_handlers({
            "mess_tempfault": self.mess_tempfault,
            "conn_throttle": self.conn_throttle,
            })

        # Start the consumer
        return self._setup_message_consumer()

    def handle_outbound_message(self, message):
        log.msg("Consumed outgoing message", message)
        log.msg("Unacknowledged message count: %s" % (
            self.esme_client.get_unacked_count()))
        #self.conn_throttle(unacked=self.esme_client.get_unacked_count())
        sequence_number = self.send_smpp(message)
        self.r_set_last_sequence(sequence_number)
        self.r_set_id_for_sequence(sequence_number,
                                   message.payload.get("message_id"))

    def esme_disconnected(self):
        log.msg("ESME Disconnected")

    def r_get_id_for_sequence(self, sequence_number):
        return self.r_server.get("%s#%s" % (self.r_prefix, sequence_number))

    def r_delete_for_sequence(self, sequence_number):
        return self.r_server.delete("%s#%s" % (self.r_prefix, sequence_number))

    def r_set_id_for_sequence(self, sequence_number, id):
        self.r_server.set("%s#%s" % (self.r_prefix, sequence_number), id)

    def r_get_last_sequence(self):
        return self.r_server.get("%s_%s#last_sequence_number" % (
            self.r_prefix, self.smpp_offset))

    def r_set_last_sequence(self, sequence_number):
        self.r_server.set("%s_%s#last_sequence_number" % (
            self.r_prefix, self.smpp_offset),
                sequence_number)

    def submit_sm_resp(self, *args, **kwargs):
        transport_msg_id = kwargs['message_id']
        sent_sms_id = self.r_get_id_for_sequence(kwargs['sequence_number'])
        self.r_delete_for_sequence(kwargs['sequence_number'])
        log.msg("Mapping transport_msg_id=%s to sent_sms_id=%s" % (
            transport_msg_id, sent_sms_id))
        log.msg("PUBLISHING ACK: (%s -> %s)" % (sent_sms_id, transport_msg_id))
        return self.publish_ack(
            user_message_id=sent_sms_id,
            sent_message_id=transport_msg_id)

    def delivery_status(self, state):
        if state in [
                "DELIVRD"
                ]:
            return "delivered"
        if state in [
                "REJECTD"
                ]:
            return "failed"
        return "pending"

    def delivery_report(self, *args, **kwargs):
        transport_metadata = {
                "message": kwargs['delivery_report'],
                "date": datetime.strptime(
                    kwargs['delivery_report']['done_date'], "%y%m%d%H%M%S")
                }
        delivery_status = self.delivery_status(
            kwargs['delivery_report']['stat']),
        message_id = kwargs['delivery_report']['id'],
        log.msg("PUBLISHING DELIV REPORT: %s %s" % (message_id,
                                                    delivery_status))
        return self.publish_delivery_report(
            user_message_id=message_id,
            delivery_status=delivery_status,
            transport_metadata=transport_metadata)

    def deliver_sm(self, *args, **kwargs):
        message = dict(
                message_id=kwargs.get('message_id'),
                to_addr=kwargs.get('destination_addr'),
                from_addr=kwargs.get('source_addr'),
                transport_type='sms',
                content=kwargs.get('short_message'))
        log.msg("PUBLISHING INBOUND: %s" % (message,))
        return self.publish_message(**message)

    def send_smpp(self, message):
        log.msg("Sending SMPP message: %s" % (message))
        # first do a lookup in our YAML to see if we've got a source_addr
        # defined for the given MT number, if not, trust the from_addr
        # in the message
        to_addr = message['to_addr']
        from_addr = message['from_addr']
        text = message['content']
        route = get_operator_number(to_addr,
                self.config.get('COUNTRY_CODE', ''),
                self.config.get('OPERATOR_PREFIX', {}),
                self.config.get('OPERATOR_NUMBER', {})) or from_addr
        sequence_number = self.esme_client.submit_sm(
                short_message=text.encode('utf-8'),
                destination_addr=str(to_addr),
                source_addr=route,
                )
        return sequence_number

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")
        return super(SmppTransport, self).stopWorker()

    def send_failure(self, message, exception, reason):
        """Send a failure report."""
        log.msg("Failed to send: %s reason: %s" % (message, reason))
        return super(SmppTransport, self).send_failure(message,
                                                       exception, reason)

    def mess_tempfault(self, *args, **kwargs):
        pdu = kwargs.get('pdu')
        sequence_number = pdu['header']['sequence_number']
        id = self.r_get_id_for_sequence(sequence_number)
        reason = pdu['header']['command_status']
        # TODO: Get real message here.
        self.send_failure(Message(id=id), reason)

    def conn_throttle(self, *args, **kwargs):
        log.msg("*********** conn_throttle: %s" % kwargs)
        unacked = kwargs.get('unacked', 0)
        if unacked > 100:
            # do something
            pass
        pdu = kwargs.get('pdu')
        if pdu:
            # do as above
            pass
<commit_msg>Delivery report handling in SMPP was broken, which was probably my fault. This fixes it.<commit_after># -*- test-case-name: vumi.transports.smpp.test_smpp -*-

from datetime import datetime

import redis
from twisted.python import log
from twisted.internet import reactor

# from vumi.service import Worker
from vumi.utils import get_operator_number, get_deploy_int
from vumi.transports.base import Transport
from vumi.transports.smpp.client import EsmeTransceiverFactory
from vumi.message import Message

class SmppTransport(Transport):
    """
    An SMPP transport.

    The SMPP transport has many configuration parameters. These are
    divided up into sections below.

    SMPP sequence number configuration options:

    :type smpp_increment: int
    :param smpp_increment:
        Increment for SMPP sequence number (must be >= number of
        SMPP workers on a single SMPP account).
    :type smpp_offset: int
    :param smpp_offset:
        Offset for this worker's SMPP sequence numbers (no duplicates
        on a single SMPP account and must be <= increment)

    SMPP server account configuration options:

    :type system_id: str
    :param system_id:
        User id used to connect to the SMPP server.
    :type password: str
    :param password:
        Password for the system id.
    :type system_type: str, optional
    :param system_type:
        Additional system metadata that is passed through to the SMPP server
        on connect.
    :type host: str
    :param host:
        Hostname of the SMPP server.
    :type port: int
    :param port:
        Port the SMPP server is listening on.
    :type initial_reconnect_delay: int, optional
    :param initial_reconnect_delay:
        Number of seconds to delay before reconnecting to the server after
        being disconnected. Default is 5s. Some WASPs, e.g. Clickatell,
        require a 30s delay before reconnecting. In these cases a 45s
        initial_reconnect_delay is recommended.

    SMPP protocol configuration options:

    :type interface_version: str, optional
    :param interface_version:
        SMPP protocol version. Default is '34' (i.e. version 3.4).
    :type dest_addr_ton:
    :param dest_addr_ton:
        Destination TON (type of number). Default .
    :type dest_addr_npi:
    :param dest_addr_npi:
        Destination NPI (number plan identifier). Default 1 (ISDN/E.164/E.163).
    :type registered_delivery:
    :param registered_delivery:
        Whether to ask for delivery reports. Default 1 (request delivery
        reports).

    The list of SMPP protocol configuration options given above is not
    exhaustive. Any other options specified are passed through to the
    python-smpp library PDU (protocol data unit) builder.

    Cellphone number routing options:

    :type COUNTRY_CODE: str, optional
    :param COUNTRY_CODE:
        Used to translate a leading zero in a destination MSISDN into a
        country code. Default '',
    :type OPERATOR_PREFIX: str, optional
    :param OPERATOR_PREFIX:
        Nested dictionary of prefix to network name mappings. Default {} (set
        network to 'UNKNOWN'). E.g. { '27': { '27761': 'NETWORK1' }}.
    :type OPERATOR_NUMBER:
    :param OPERATOR_NUMBER:
        Dictionary of source MSISDN to use for each network listed in
        OPERATOR_PREFIX. If a network is not listed, the source MSISDN
        specified by the message sender is used. Default {} (always used the
        from address specified by the message sender). E.g. { 'NETWORK1':
        '27761234567'}.
    """

    def setup_transport(self):
        log.msg("Starting the SmppTransport with %s" % self.config)

        # TODO: move this to a config file
        dbindex = get_deploy_int(self._amqp_client.vhost)
        self.smpp_offset = int(self.config['smpp_offset'])

        # Connect to Redis
        if not hasattr(self, 'r_server'):
            # Only set up redis if we don't have a test stub already
            self.r_server = redis.Redis("localhost", db=dbindex)
        self.r_prefix = "%(system_id)s@%(host)s:%(port)s" % self.config
        log.msg("Connected to Redis, prefix: %s" % self.r_prefix)
        last_sequence_number = int(self.r_get_last_sequence()
                                   or self.smpp_offset)
        log.msg("Last sequence_number: %s" % last_sequence_number)

        if not hasattr(self, 'esme_client'):
            # start the Smpp transport (if we don't have one)
            factory = EsmeTransceiverFactory(self.config,
                                             self._amqp_client.vumi_options)
            factory.loadDefaults(self.config)
            factory.setLastSequenceNumber(last_sequence_number)
            factory.setConnectCallback(self.esme_connected)
            factory.setDisconnectCallback(self.esme_disconnected)
            factory.setSubmitSMRespCallback(self.submit_sm_resp)
            factory.setDeliveryReportCallback(self.delivery_report)
            factory.setDeliverSMCallback(self.deliver_sm)
            factory.setSendFailureCallback(self.send_failure)
            log.msg(factory.defaults)
            reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.update_error_handlers({
            "mess_tempfault": self.mess_tempfault,
            "conn_throttle": self.conn_throttle,
            })

        # Start the consumer
        return self._setup_message_consumer()

    def handle_outbound_message(self, message):
        log.msg("Consumed outgoing message", message)
        log.msg("Unacknowledged message count: %s" % (
            self.esme_client.get_unacked_count()))
        #self.conn_throttle(unacked=self.esme_client.get_unacked_count())
        sequence_number = self.send_smpp(message)
        self.r_set_last_sequence(sequence_number)
        self.r_set_id_for_sequence(sequence_number,
                                   message.payload.get("message_id"))

    def esme_disconnected(self):
        log.msg("ESME Disconnected")

    def r_get_id_for_sequence(self, sequence_number):
        return self.r_server.get("%s#%s" % (self.r_prefix, sequence_number))

    def r_delete_for_sequence(self, sequence_number):
        return self.r_server.delete("%s#%s" % (self.r_prefix, sequence_number))

    def r_set_id_for_sequence(self, sequence_number, id):
        self.r_server.set("%s#%s" % (self.r_prefix, sequence_number), id)

    def r_get_last_sequence(self):
        return self.r_server.get("%s_%s#last_sequence_number" % (
            self.r_prefix, self.smpp_offset))

    def r_set_last_sequence(self, sequence_number):
        self.r_server.set("%s_%s#last_sequence_number" % (
            self.r_prefix, self.smpp_offset),
                sequence_number)

    def submit_sm_resp(self, *args, **kwargs):
        transport_msg_id = kwargs['message_id']
        sent_sms_id = self.r_get_id_for_sequence(kwargs['sequence_number'])
        self.r_delete_for_sequence(kwargs['sequence_number'])
        log.msg("Mapping transport_msg_id=%s to sent_sms_id=%s" % (
            transport_msg_id, sent_sms_id))
        log.msg("PUBLISHING ACK: (%s -> %s)" % (sent_sms_id, transport_msg_id))
        return self.publish_ack(
            user_message_id=sent_sms_id,
            sent_message_id=transport_msg_id)

    def delivery_status(self, state):
        if state in [
                "DELIVRD"
                ]:
            return "delivered"
        if state in [
                "REJECTD"
                ]:
            return "failed"
        return "pending"

    def delivery_report(self, *args, **kwargs):
        transport_metadata = {
                "message": kwargs['delivery_report'],
                "date": datetime.strptime(
                    kwargs['delivery_report']['done_date'], "%y%m%d%H%M%S")
                }
        delivery_status = self.delivery_status(
            kwargs['delivery_report']['stat'])
        message_id = kwargs['delivery_report']['id']
        log.msg("PUBLISHING DELIV REPORT: %s %s" % (message_id,
                                                    delivery_status))
        return self.publish_delivery_report(
            user_message_id=message_id,
            delivery_status=delivery_status,
            transport_metadata=transport_metadata)

    def deliver_sm(self, *args, **kwargs):
        message = dict(
                message_id=kwargs.get('message_id'),
                to_addr=kwargs.get('destination_addr'),
                from_addr=kwargs.get('source_addr'),
                transport_type='sms',
                content=kwargs.get('short_message'))
        log.msg("PUBLISHING INBOUND: %s" % (message,))
        return self.publish_message(**message)

    def send_smpp(self, message):
        log.msg("Sending SMPP message: %s" % (message))
        # first do a lookup in our YAML to see if we've got a source_addr
        # defined for the given MT number, if not, trust the from_addr
        # in the message
        to_addr = message['to_addr']
        from_addr = message['from_addr']
        text = message['content']
        route = get_operator_number(to_addr,
                self.config.get('COUNTRY_CODE', ''),
                self.config.get('OPERATOR_PREFIX', {}),
                self.config.get('OPERATOR_NUMBER', {})) or from_addr
        sequence_number = self.esme_client.submit_sm(
                short_message=text.encode('utf-8'),
                destination_addr=str(to_addr),
                source_addr=route,
                )
        return sequence_number

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")
        return super(SmppTransport, self).stopWorker()

    def send_failure(self, message, exception, reason):
        """Send a failure report."""
        log.msg("Failed to send: %s reason: %s" % (message, reason))
        return super(SmppTransport, self).send_failure(message,
                                                       exception, reason)

    def mess_tempfault(self, *args, **kwargs):
        pdu = kwargs.get('pdu')
        sequence_number = pdu['header']['sequence_number']
        id = self.r_get_id_for_sequence(sequence_number)
        reason = pdu['header']['command_status']
        # TODO: Get real message here.
        self.send_failure(Message(id=id), reason)

    def conn_throttle(self, *args, **kwargs):
        log.msg("*********** conn_throttle: %s" % kwargs)
        unacked = kwargs.get('unacked', 0)
        if unacked > 100:
            # do something
            pass
        pdu = kwargs.get('pdu')
        if pdu:
            # do as above
            pass
<|endoftext|>