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

from IPython.nbformat import (
    from_dict,
)
from IPython.utils.traitlets import (
    Bool,
    Integer,
)
from IPython.html.services.contents.manager import ContentsManager
from tornado import web

from .api_utils import (
    base_model,
    base_directory_model,
    DUMMY_CREATED_DATE,
    from_b64,
    outside_root_to_404,
    reads_base64,
    to_api_path,
    to_b64,
    writes_base64,
)
from .checkpoints import PostgresCheckpoints
from .constants import UNLIMITED
from .error import (
    DirectoryNotEmpty,
    FileExists,
    FileTooLarge,
    NoSuchDirectory,
    NoSuchFile,
    PathOutsideRoot,
)
from .managerbase import PostgresManagerMixin
from .query import (
    delete_file,
    delete_directory,
    dir_exists,
    ensure_directory,
    get_directory,
    get_file,
    purge_user,
    rename_file,
    save_file,
)

class PostgresContentsManager(PostgresManagerMixin, ContentsManager):
    """
    ContentsManager that persists to a postgres database rather than to the
    local filesystem.
    """
    max_file_size_bytes = Integer(
        default_value=UNLIMITED,
        config=True,
        help="Maximum size in bytes of a file that will be saved.",
    )

    create_directory_on_startup = Bool(
        config=True,
        help="Create a user for user_id automatically?",
    )

    def _checkpoints_class_default(self):
        return PostgresCheckpoints

    def _checkpoints_kwargs_default(self):
        kw = super(PostgresContentsManager, self)._checkpoints_kwargs_default()
        kw.update(
            {
                'db_url': self.db_url,
                'user_id': self.user_id,
                'create_user_on_startup': self.create_user_on_startup,
            }
        )
        return kw

    def _create_directory_on_startup_default(self):
        return self.create_user_on_startup

    def __init__(self, *args, **kwargs):
        super(PostgresContentsManager, self).__init__(*args, **kwargs)
        if self.create_directory_on_startup:
            self.ensure_root_directory()

    def ensure_root_directory(self):
        with self.engine.begin() as db:
            ensure_directory(db, self.user_id, '')

    def purge_db(self):
        """
        Clear all matching our user_id.
        """
        with self.engine.begin() as db:
            purge_user(db, self.user_id)

    @outside_root_to_404
    def guess_type(self, path, allow_directory=True):
        """
        Guess the type of a file.

        If allow_directory is False, don't consider the possibility that the
        file is a directory.
        """
        if path.endswith('.ipynb'):
            return 'notebook'
        elif allow_directory and self.dir_exists(path):
            return 'directory'
        else:
            return 'file'

    # Begin ContentsManager API.
    @outside_root_to_404
    def dir_exists(self, path):
        with self.engine.begin() as db:
            return dir_exists(db, self.user_id, path)

    def is_hidden(self, path):
        return False

    @outside_root_to_404
    def file_exists(self, path):
        with self.engine.begin() as db:
            try:
                get_file(db, self.user_id, path, include_content=False)
                return True
            except NoSuchFile:
                return False

    @outside_root_to_404
    def get(self, path, content=True, type=None, format=None):
        if type is None:
            type = self.guess_type(path)
        try:
            fn = {
                'notebook': self._get_notebook,
                'directory': self._get_directory,
                'file': self._get_file,
            }[type]
        except KeyError:
            raise ValueError("Unknown type passed: '{}'".format(type))
        return fn(path=path, content=content, format=format)

    def _get_notebook(self, path, content, format):
        """
        Get a notebook from the database.
        """
        with self.engine.begin() as db:
            try:
                record = get_file(db, self.user_id, path, content)
            except NoSuchFile:
                self.no_such_entity(path)

        return self._notebook_model_from_db(record, content)

    def _notebook_model_from_db(self, record, content):
        """
        Build a notebook model from database record.
        """
        path = to_api_path(record['parent_name'] + record['name'])
        model = base_model(path)
        model['type'] = 'notebook'
        model['last_modified'] = model['created'] = record['created_at']
        if content:
            content = reads_base64(record['content'])
            self.mark_trusted_cells(content, path)
            model['content'] = content
            model['format'] = 'json'
            self.validate_notebook_model(model)
        return model

    def _get_directory(self, path, content, format):
        """
        Get a directory from the database.
        """
        with self.engine.begin() as db:
            try:
                record = get_directory(
                    db, self.user_id, path, content
                )
            except NoSuchDirectory:
                if self.file_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400("Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)

        return self._directory_model_from_db(record, content)

    def _convert_file_records(self, file_records):
        """
        Apply _notebook_model_from_db or _file_model_from_db to each entry
        in file_records, depending on the result of `guess_type`.
        """
        for record in file_records:
            type_ = self.guess_type(record['name'], allow_directory=False)
            if type_ == 'notebook':
                yield self._notebook_model_from_db(record, False)
            elif type_ == 'file':
                yield self._file_model_from_db(record, False, None)
            else:
                self.do_500("Unknown file type %s" % type_)

    def _directory_model_from_db(self, record, content):
        """
        Build a directory model from database directory record.
        """
        model = base_directory_model(to_api_path(record['name']))
        if content:
            model['format'] = 'json'
            model['content'] = list(
                chain(
                    self._convert_file_records(record['files']),
                    (
                        self._directory_model_from_db(subdir, False)
                        for subdir in record['subdirs']
                    ),
                )
            )
        return model

    def _file_model_from_db(self, record, content, format):
        """
        Build a file model from database record.
        """
        # TODO: Most of this is shared with _notebook_model_from_db.
        path = to_api_path(record['parent_name'] + record['name'])
        model = base_model(path)
        model['type'] = 'file'
        model['last_modified'] = model['created'] = record['created_at']
        if content:
            bcontent = record['content']
            model['content'], model['format'], model['mimetype'] = from_b64(
                path,
                bcontent,
                format,
            )
        return model

    def _get_file(self, path, content, format):
        with self.engine.begin() as db:
            try:
                record = get_file(db, self.user_id, path, content)
            except NoSuchFile:
                if self.dir_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400(u"Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)
        return self._file_model_from_db(record, content, format)

    def _save_notebook(self, db, model, path):
        """
        Save a notebook.

        Returns a validation message.
        """
        nb_contents = from_dict(model['content'])
        self.check_and_sign(nb_contents, path)
        save_file(
            db,
            self.user_id,
            path,
            writes_base64(nb_contents),
            self.max_file_size_bytes,
        )
        # It's awkward that this writes to the model instead of returning.
        self.validate_notebook_model(model)
        return model.get('message')

    def _save_file(self, db, model, path):
        """
        Save a non-notebook file.
        """
        save_file(
            db,
            self.user_id,
            path,
            to_b64(model['content'], model.get('format', None)),
            self.max_file_size_bytes,
        )
        return None

    def _save_directory(self, db, path):
        """
        'Save' a directory.
        """
        ensure_directory(db, self.user_id, path)

    @outside_root_to_404
    def save(self, model, path):
        if 'type' not in model:
            raise web.HTTPError(400, u'No model type provided')
        if 'content' not in model and model['type'] != 'directory':
            raise web.HTTPError(400, u'No file content provided')

        path = path.strip('/')

        # Almost all of this is duplicated with FileContentsManager :(.
        self.log.debug("Saving %s", path)
        if model['type'] not in ('file', 'directory', 'notebook'):
            self.do_400("Unhandled contents type: %s" % model['type'])
        try:
            with self.engine.begin() as db:
                if model['type'] == 'notebook':
                    validation_message = self._save_notebook(db, model, path)
                elif model['type'] == 'file':
                    validation_message = self._save_file(db, model, path)
                else:
                    validation_message = self._save_directory(db, path)
        except (web.HTTPError, PathOutsideRoot):
            raise
        except FileTooLarge:
            self.file_too_large(path)
        except Exception as e:
            self.log.error(u'Error while saving file: %s %s',
                           path, e, exc_info=True)
            self.do_500(
                u'Unexpected error while saving file: %s %s' % (path, e)
            )

        # TODO: Consider not round-tripping to the database again here.
        model = self.get(path, type=model['type'], content=False)
        if validation_message is not None:
            model['message'] = validation_message
        return model

    @outside_root_to_404
    def rename_file(self, old_path, path):
        """
        Rename object from old_path to path.
        """
        with self.engine.begin() as db:
            try:
                rename_file(db, self.user_id, old_path, path)
            except FileExists:
                self.already_exists(path)

    def _delete_non_directory(self, path):
        with self.engine.begin() as db:
            deleted_count = delete_file(db, self.user_id, path)
            if not deleted_count:
                self.no_such_entity(path)

    def _delete_directory(self, path):
        with self.engine.begin() as db:
            try:
                deleted_count = delete_directory(db, self.user_id, path)
            except DirectoryNotEmpty:
                self.not_empty(path)
            if not deleted_count:
                self.no_such_entity(path)

    @outside_root_to_404
    def delete_file(self, path):
        """
        Delete object corresponding to path.
        """
        if self.file_exists(path):
            self._delete_non_directory(path)
        elif self.dir_exists(path):
            self._delete_directory(path)
        else:
            self.no_such_entity(path)
<commit_msg>MAINT: Unused import.<commit_after>#
# Copyright 2014 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
PostgreSQL implementation of IPython ContentsManager API.
"""
from __future__ import unicode_literals
from itertools import chain

from IPython.nbformat import (
    from_dict,
)
from IPython.utils.traitlets import (
    Bool,
    Integer,
)
from IPython.html.services.contents.manager import ContentsManager
from tornado import web

from .api_utils import (
    base_model,
    base_directory_model,
    from_b64,
    outside_root_to_404,
    reads_base64,
    to_api_path,
    to_b64,
    writes_base64,
)
from .checkpoints import PostgresCheckpoints
from .constants import UNLIMITED
from .error import (
    DirectoryNotEmpty,
    FileExists,
    FileTooLarge,
    NoSuchDirectory,
    NoSuchFile,
    PathOutsideRoot,
)
from .managerbase import PostgresManagerMixin
from .query import (
    delete_file,
    delete_directory,
    dir_exists,
    ensure_directory,
    get_directory,
    get_file,
    purge_user,
    rename_file,
    save_file,
)

class PostgresContentsManager(PostgresManagerMixin, ContentsManager):
    """
    ContentsManager that persists to a postgres database rather than to the
    local filesystem.
    """
    max_file_size_bytes = Integer(
        default_value=UNLIMITED,
        config=True,
        help="Maximum size in bytes of a file that will be saved.",
    )

    create_directory_on_startup = Bool(
        config=True,
        help="Create a user for user_id automatically?",
    )

    def _checkpoints_class_default(self):
        return PostgresCheckpoints

    def _checkpoints_kwargs_default(self):
        kw = super(PostgresContentsManager, self)._checkpoints_kwargs_default()
        kw.update(
            {
                'db_url': self.db_url,
                'user_id': self.user_id,
                'create_user_on_startup': self.create_user_on_startup,
            }
        )
        return kw

    def _create_directory_on_startup_default(self):
        return self.create_user_on_startup

    def __init__(self, *args, **kwargs):
        super(PostgresContentsManager, self).__init__(*args, **kwargs)
        if self.create_directory_on_startup:
            self.ensure_root_directory()

    def ensure_root_directory(self):
        with self.engine.begin() as db:
            ensure_directory(db, self.user_id, '')

    def purge_db(self):
        """
        Clear all matching our user_id.
        """
        with self.engine.begin() as db:
            purge_user(db, self.user_id)

    @outside_root_to_404
    def guess_type(self, path, allow_directory=True):
        """
        Guess the type of a file.

        If allow_directory is False, don't consider the possibility that the
        file is a directory.
        """
        if path.endswith('.ipynb'):
            return 'notebook'
        elif allow_directory and self.dir_exists(path):
            return 'directory'
        else:
            return 'file'

    # Begin ContentsManager API.
    @outside_root_to_404
    def dir_exists(self, path):
        with self.engine.begin() as db:
            return dir_exists(db, self.user_id, path)

    def is_hidden(self, path):
        return False

    @outside_root_to_404
    def file_exists(self, path):
        with self.engine.begin() as db:
            try:
                get_file(db, self.user_id, path, include_content=False)
                return True
            except NoSuchFile:
                return False

    @outside_root_to_404
    def get(self, path, content=True, type=None, format=None):
        if type is None:
            type = self.guess_type(path)
        try:
            fn = {
                'notebook': self._get_notebook,
                'directory': self._get_directory,
                'file': self._get_file,
            }[type]
        except KeyError:
            raise ValueError("Unknown type passed: '{}'".format(type))
        return fn(path=path, content=content, format=format)

    def _get_notebook(self, path, content, format):
        """
        Get a notebook from the database.
        """
        with self.engine.begin() as db:
            try:
                record = get_file(db, self.user_id, path, content)
            except NoSuchFile:
                self.no_such_entity(path)

        return self._notebook_model_from_db(record, content)

    def _notebook_model_from_db(self, record, content):
        """
        Build a notebook model from database record.
        """
        path = to_api_path(record['parent_name'] + record['name'])
        model = base_model(path)
        model['type'] = 'notebook'
        model['last_modified'] = model['created'] = record['created_at']
        if content:
            content = reads_base64(record['content'])
            self.mark_trusted_cells(content, path)
            model['content'] = content
            model['format'] = 'json'
            self.validate_notebook_model(model)
        return model

    def _get_directory(self, path, content, format):
        """
        Get a directory from the database.
        """
        with self.engine.begin() as db:
            try:
                record = get_directory(
                    db, self.user_id, path, content
                )
            except NoSuchDirectory:
                if self.file_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400("Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)

        return self._directory_model_from_db(record, content)

    def _convert_file_records(self, file_records):
        """
        Apply _notebook_model_from_db or _file_model_from_db to each entry
        in file_records, depending on the result of `guess_type`.
        """
        for record in file_records:
            type_ = self.guess_type(record['name'], allow_directory=False)
            if type_ == 'notebook':
                yield self._notebook_model_from_db(record, False)
            elif type_ == 'file':
                yield self._file_model_from_db(record, False, None)
            else:
                self.do_500("Unknown file type %s" % type_)

    def _directory_model_from_db(self, record, content):
        """
        Build a directory model from database directory record.
        """
        model = base_directory_model(to_api_path(record['name']))
        if content:
            model['format'] = 'json'
            model['content'] = list(
                chain(
                    self._convert_file_records(record['files']),
                    (
                        self._directory_model_from_db(subdir, False)
                        for subdir in record['subdirs']
                    ),
                )
            )
        return model

    def _file_model_from_db(self, record, content, format):
        """
        Build a file model from database record.
        """
        # TODO: Most of this is shared with _notebook_model_from_db.
        path = to_api_path(record['parent_name'] + record['name'])
        model = base_model(path)
        model['type'] = 'file'
        model['last_modified'] = model['created'] = record['created_at']
        if content:
            bcontent = record['content']
            model['content'], model['format'], model['mimetype'] = from_b64(
                path,
                bcontent,
                format,
            )
        return model

    def _get_file(self, path, content, format):
        with self.engine.begin() as db:
            try:
                record = get_file(db, self.user_id, path, content)
            except NoSuchFile:
                if self.dir_exists(path):
                    # TODO: It's awkward/expensive to have to check this to
                    # return a 400 instead of 404. Consider just 404ing.
                    self.do_400(u"Wrong type: %s" % path)
                else:
                    self.no_such_entity(path)
        return self._file_model_from_db(record, content, format)

    def _save_notebook(self, db, model, path):
        """
        Save a notebook.

        Returns a validation message.
        """
        nb_contents = from_dict(model['content'])
        self.check_and_sign(nb_contents, path)
        save_file(
            db,
            self.user_id,
            path,
            writes_base64(nb_contents),
            self.max_file_size_bytes,
        )
        # It's awkward that this writes to the model instead of returning.
        self.validate_notebook_model(model)
        return model.get('message')

    def _save_file(self, db, model, path):
        """
        Save a non-notebook file.
        """
        save_file(
            db,
            self.user_id,
            path,
            to_b64(model['content'], model.get('format', None)),
            self.max_file_size_bytes,
        )
        return None

    def _save_directory(self, db, path):
        """
        'Save' a directory.
        """
        ensure_directory(db, self.user_id, path)

    @outside_root_to_404
    def save(self, model, path):
        if 'type' not in model:
            raise web.HTTPError(400, u'No model type provided')
        if 'content' not in model and model['type'] != 'directory':
            raise web.HTTPError(400, u'No file content provided')

        path = path.strip('/')

        # Almost all of this is duplicated with FileContentsManager :(.
        self.log.debug("Saving %s", path)
        if model['type'] not in ('file', 'directory', 'notebook'):
            self.do_400("Unhandled contents type: %s" % model['type'])
        try:
            with self.engine.begin() as db:
                if model['type'] == 'notebook':
                    validation_message = self._save_notebook(db, model, path)
                elif model['type'] == 'file':
                    validation_message = self._save_file(db, model, path)
                else:
                    validation_message = self._save_directory(db, path)
        except (web.HTTPError, PathOutsideRoot):
            raise
        except FileTooLarge:
            self.file_too_large(path)
        except Exception as e:
            self.log.error(u'Error while saving file: %s %s',
                           path, e, exc_info=True)
            self.do_500(
                u'Unexpected error while saving file: %s %s' % (path, e)
            )

        # TODO: Consider not round-tripping to the database again here.
        model = self.get(path, type=model['type'], content=False)
        if validation_message is not None:
            model['message'] = validation_message
        return model

    @outside_root_to_404
    def rename_file(self, old_path, path):
        """
        Rename object from old_path to path.
        """
        with self.engine.begin() as db:
            try:
                rename_file(db, self.user_id, old_path, path)
            except FileExists:
                self.already_exists(path)

    def _delete_non_directory(self, path):
        with self.engine.begin() as db:
            deleted_count = delete_file(db, self.user_id, path)
            if not deleted_count:
                self.no_such_entity(path)

    def _delete_directory(self, path):
        with self.engine.begin() as db:
            try:
                deleted_count = delete_directory(db, self.user_id, path)
            except DirectoryNotEmpty:
                self.not_empty(path)
            if not deleted_count:
                self.no_such_entity(path)

    @outside_root_to_404
    def delete_file(self, path):
        """
        Delete object corresponding to path.
        """
        if self.file_exists(path):
            self._delete_non_directory(path)
        elif self.dir_exists(path):
            self._delete_directory(path)
        else:
            self.no_such_entity(path)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

from collections import defaultdict
from itertools import chain

from pgmpy.models import BayesianModel
from pgmpy.models import MarkovModel
from pgmpy.models import FactorGraph
from pgmpy.models import JunctionTree
from pgmpy.models import DynamicBayesianNetwork
from pgmpy.exceptions import ModelError

class Inference:
    """
    Base class for all inference algorithms.

    Converts BayesianModel and MarkovModel to a uniform representation so that inference
    algorithms can be applied. Also it checks if all the associated CPDs / Factors are
    consistent with the model.

    Initialize inference for a model.

    Parameters
    ----------
    model: pgmpy.models.BayesianModel or pgmpy.models.MarkovModel or pgmpy.models.NoisyOrModel
        model for which to initialize the inference object.

    Examples
    --------
    >>> from pgmpy.inference import Inference
    >>> from pgmpy.models import BayesianModel
    >>> from pgmpy.factors import TabularCPD
    >>> student = BayesianModel([('diff', 'grade'), ('intel', 'grade')])
    >>> diff_cpd = TabularCPD('diff', 2, [[0.2, 0.8]])
    >>> intel_cpd = TabularCPD('intel', 2, [[0.3, 0.7]])
    >>> grade_cpd = TabularCPD('grade', 3, [[0.1, 0.1, 0.1, 0.1],
    ...                                     [0.1, 0.1, 0.1, 0.1],
    ...                                     [0.8, 0.8, 0.8, 0.8]],
    ...                        evidence=['diff', 'intel'], evidence_card=[2, 2])
    >>> student.add_cpds(diff_cpd, intel_cpd, grade_cpd)
    >>> model = Inference(student)

    >>> from pgmpy.models import MarkovModel
    >>> from pgmpy.factors import Factor
    >>> import numpy as np
    >>> student = MarkovModel([('Alice', 'Bob'), ('Bob', 'Charles'),
    ...                        ('Charles', 'Debbie'), ('Debbie', 'Alice')])
    >>> factor_a_b = Factor(['Alice', 'Bob'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_b_c = Factor(['Bob', 'Charles'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_c_d = Factor(['Charles', 'Debbie'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_d_a = Factor(['Debbie', 'Alice'], cardinality=[2, 2], value=np.random.rand(4))
    >>> student.add_factors(factor_a_b, factor_b_c, factor_c_d, factor_d_a)
    >>> model = Inference(student)
    """

    def __init__(self, model):
        if not model.check_model():
            raise ModelError("Model of type {!r} not valid".format(type(model).__name__))

        if isinstance(model, JunctionTree):
            self.variables = set(chain(*model.nodes()))
        else:
            self.variables = model.nodes()

        self.cardinality = {}
        self.factors = defaultdict(list)

        if isinstance(model, BayesianModel):
            for node in model.nodes():
                cpd = model.get_cpds(node)
                cpd_as_factor = cpd.to_factor()
                self.cardinality[node] = cpd.variable_card

                for var in cpd.variables:
                    self.factors[var].append(cpd_as_factor)

        elif isinstance(model, (MarkovModel, FactorGraph, JunctionTree)):
            self.cardinality = model.cardinalities

            for factor in model.get_factors():
                for var in factor.variables:
                    self.factors[var].append(factor)

        elif isinstance(model, DynamicBayesianNetwork):
            self.start_bayesian_model = BayesianModel(model.get_intra_edges(0))
            self.start_bayesian_model.add_cpds(*model.get_cpds(time_slice=0))
            cpd_inter = [model.get_cpds(node) for node in model.get_interface_nodes(1)]
            self.interface_nodes = model.get_interface_nodes(0)
            self.one_and_half_model = BayesianModel(model.get_inter_edges() + model.get_intra_edges(1))
            self.one_and_half_model.add_cpds(*(model.get_cpds(time_slice=1) + cpd_inter))
<commit_msg>updated use of cardinality in inference/base.py<commit_after>#!/usr/bin/env python3

from collections import defaultdict
from itertools import chain

from pgmpy.models import BayesianModel
from pgmpy.models import MarkovModel
from pgmpy.models import FactorGraph
from pgmpy.models import JunctionTree
from pgmpy.models import DynamicBayesianNetwork
from pgmpy.exceptions import ModelError

class Inference:
    """
    Base class for all inference algorithms.

    Converts BayesianModel and MarkovModel to a uniform representation so that inference
    algorithms can be applied. Also it checks if all the associated CPDs / Factors are
    consistent with the model.

    Initialize inference for a model.

    Parameters
    ----------
    model: pgmpy.models.BayesianModel or pgmpy.models.MarkovModel or pgmpy.models.NoisyOrModel
        model for which to initialize the inference object.

    Examples
    --------
    >>> from pgmpy.inference import Inference
    >>> from pgmpy.models import BayesianModel
    >>> from pgmpy.factors import TabularCPD
    >>> student = BayesianModel([('diff', 'grade'), ('intel', 'grade')])
    >>> diff_cpd = TabularCPD('diff', 2, [[0.2, 0.8]])
    >>> intel_cpd = TabularCPD('intel', 2, [[0.3, 0.7]])
    >>> grade_cpd = TabularCPD('grade', 3, [[0.1, 0.1, 0.1, 0.1],
    ...                                     [0.1, 0.1, 0.1, 0.1],
    ...                                     [0.8, 0.8, 0.8, 0.8]],
    ...                        evidence=['diff', 'intel'], evidence_card=[2, 2])
    >>> student.add_cpds(diff_cpd, intel_cpd, grade_cpd)
    >>> model = Inference(student)

    >>> from pgmpy.models import MarkovModel
    >>> from pgmpy.factors import Factor
    >>> import numpy as np
    >>> student = MarkovModel([('Alice', 'Bob'), ('Bob', 'Charles'),
    ...                        ('Charles', 'Debbie'), ('Debbie', 'Alice')])
    >>> factor_a_b = Factor(['Alice', 'Bob'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_b_c = Factor(['Bob', 'Charles'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_c_d = Factor(['Charles', 'Debbie'], cardinality=[2, 2], value=np.random.rand(4))
    >>> factor_d_a = Factor(['Debbie', 'Alice'], cardinality=[2, 2], value=np.random.rand(4))
    >>> student.add_factors(factor_a_b, factor_b_c, factor_c_d, factor_d_a)
    >>> model = Inference(student)
    """

    def __init__(self, model):
        if not model.check_model():
            raise ModelError("Model of type {!r} not valid".format(type(model).__name__))

        if isinstance(model, JunctionTree):
            self.variables = set(chain(*model.nodes()))
        else:
            self.variables = model.nodes()

        self.cardinality = {}
        self.factors = defaultdict(list)

        if isinstance(model, BayesianModel):
            for node in model.nodes():
                cpd = model.get_cpds(node)
                cpd_as_factor = cpd.to_factor()
                self.cardinality[node] = cpd.variable_card

                for var in cpd.variables:
                    self.factors[var].append(cpd_as_factor)

        elif isinstance(model, (MarkovModel, FactorGraph, JunctionTree)):
            self.cardinality = model.get_cardinality()

            for factor in model.get_factors():
                for var in factor.variables:
                    self.factors[var].append(factor)

        elif isinstance(model, DynamicBayesianNetwork):
            self.start_bayesian_model = BayesianModel(model.get_intra_edges(0))
            self.start_bayesian_model.add_cpds(*model.get_cpds(time_slice=0))
            cpd_inter = [model.get_cpds(node) for node in model.get_interface_nodes(1)]
            self.interface_nodes = model.get_interface_nodes(0)
            self.one_and_half_model = BayesianModel(model.get_inter_edges() + model.get_intra_edges(1))
            self.one_and_half_model.add_cpds(*(model.get_cpds(time_slice=1) + cpd_inter))
<|endoftext|>


<commit_before>"""
python-phabricator
------------------
>>> api = phabricator.Phabricator()
>>> api.user.whoami().userName
'example'

For more endpoints, see https://secure.phabricator.com/conduit/

"""
try:
    __version__ = __import__('pkg_resources') \
        .get_distribution('phabricator').version
except:
    __version__ = 'unknown'

import collections
import copy
import hashlib
import json
import os.path
import re
import socket
import time

from ._compat import (
    MutableMapping, iteritems, string_types, httplib, urlparse, urlencode,
)

__all__ = ['Phabricator']

# Default phabricator interfaces
INTERFACES = json.loads(open(os.path.join(os.path.dirname(__file__), 'interfaces.json'), 'r').read())

# Load arc config
ARC_CONFIGS = [
    # System config
    os.path.join(os.environ['ProgramData'], 'Phabricator', 'Arcanist', 'config') if os.name == 'nt' else
    os.path.join('/etc', 'arcconfig'),

    # User config
    os.path.join(os.environ['AppData'] if os.name == 'nt' else os.path.expanduser('~'), '.arcrc'),

    # Project config
    os.path.join(os.getcwd(), '.arcconfig'),

    # Local project config
    os.path.join(os.getcwd(), '.git', 'arc', 'config'),
]

ARCRC = {}
for conf in ARC_CONFIGS:
    if os.path.exists(conf):
        ARCRC.update(json.load(open(conf, 'r')))

# Map Phabricator types to Python types
PARAM_TYPE_MAP = {
    # int types
    'int': int,
    'uint': int,
    'revisionid': int,
    'revision_id': int,
    'diffid': int,
    'diff_id': int,
    'id': int,
    'enum': int,

    # bool types
    'bool': bool,

    # dict types
    'map': dict,
    'dict': dict,

    # list types
    'list': list,

    # tuple types
    'pair': tuple,

    # str types
    'str': string_types,
    'string': string_types,
    'phid': string_types,
    'guids': string_types,
    'type': string_types,
}

STR_RE = re.compile(r'([a-zA-Z_]+)')

def map_param_type(param_type):
    """
    Perform param type mapping
    This requires a bit of logic since this isn't standardized.
    If a type doesn't map, assume str
    """
    m = STR_RE.match(param_type)
    main_type = m.group(0)

    if main_type in ('list', 'array'):
        info = param_type.replace(' ', '').split('<', 1)

        # Handle no sub-type: "required list"
        sub_type = info[1] if len(info) > 1 else 'str'

        # Handle list of pairs: "optional list<pair<callsign, path>>"
        sub_match = STR_RE.match(sub_type)
        sub_type = sub_match.group(0).lower()

        return [PARAM_TYPE_MAP.setdefault(sub_type, string_types)]

    return PARAM_TYPE_MAP.setdefault(main_type, string_types)

def parse_interfaces(interfaces):
    """
    Parse the conduit.query json dict response
    This performs the logic of parsing the non-standard params dict
        and then returning a dict Resource can understand
    """
    parsed_interfaces = collections.defaultdict(dict)

    for m, d in iteritems(interfaces):
        app, func = m.split('.', 1)

        method = parsed_interfaces[app][func] = {}

        # Make default assumptions since these aren't provided by Phab
        method['formats'] = ['json', 'human']
        method['method'] = 'POST'

        method['optional'] = {}
        method['required'] = {}

        for name, type_info in iteritems(dict(d['params'])):
            # Usually in the format: <optionality> <param_type>
            info_pieces = type_info.split(' ', 1)

            # If optionality isn't specified, assume required
            if info_pieces[0] not in ('optional', 'required'):
                optionality = 'required'
                param_type = info_pieces[0]
            # Just make an optional string for "ignored" params
            elif info_pieces[0] == 'ignored':
                optionality = 'optional'
                param_type = 'string'
            else:
                optionality = info_pieces[0]
                param_type = info_pieces[1]

            # This isn't validated by the client
            if param_type.startswith('nonempty'):
                optionality = 'required'
                param_type = param_type[9:]

            method[optionality][name] = map_param_type(param_type)

    return dict(parsed_interfaces)

class InterfaceNotDefined(NotImplementedError):
    pass

class APIError(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message

    def __str__(self):
        return '%s: %s' % (self.code, self.message)

class InvalidAccessToken(APIError):
    pass

class Result(MutableMapping):
    def __init__(self, response):
        self.response = response

    def __getitem__(self, key):
        return self.response[key]

    __getattr__ = __getitem__

    def __setitem__(self, key, value):
        self.response[key] = value

    def __delitem__(self, key):
        del self.response[key]

    def __iter__(self):
        return iter(self.response)

    def __len__(self):
        return len(self.response.keys())

    def __repr__(self):
        return '<%s: %s>' % (type(self).__name__, repr(self.response))

class Resource(object):
    def __init__(self, api, interface=None, endpoint=None, method=None):
        self.api = api
        self.interface = interface or copy.deepcopy(parse_interfaces(INTERFACES))
        self.endpoint = endpoint
        self.method = method

    def __getattr__(self, attr):
        if attr in getattr(self, '__dict__'):
            return getattr(self, attr)
        interface = self.interface
        if attr not in interface:
            interface[attr] = {}
        return Resource(self.api, interface[attr], attr, self.endpoint)

    def __call__(self, **kwargs):
        return self._request(**kwargs)

    def _request(self, **kwargs):
        # Check for missing variables
        resource = self.interface

        def validate_kwarg(key, target):
            # Always allow list
            if isinstance(key, list):
                return all([validate_kwarg(x, target[0]) for x in key])
            return isinstance(key, tuple(target) if isinstance(target, list) else target)

        for key, val in resource.get('required', {}).items():
            if key not in [x.split(':')[0] for x in kwargs.keys()]:
                raise ValueError('Missing required argument: %s' % key)
            if isinstance(kwargs.get(key), list) and not isinstance(val, list):
                raise ValueError('Wrong argument type: %s is not a list' % key)
            elif not validate_kwarg(kwargs.get(key), val):
                if isinstance(val, list):
                    raise ValueError('Wrong argument type: %s is not a list of %ss' % (key, val[0]))
                raise ValueError('Wrong argument type: %s is not a %s' % (key, val))

        conduit = self.api._conduit

        if conduit:
            # Already authenticated, add session key to json data
            kwargs['__conduit__'] = conduit
        elif self.method == 'conduit' and self.endpoint == 'connect':
            # Not authenticated, requesting new session key
            token = str(int(time.time()))
            kwargs['authToken'] = token
            kwargs['authSignature'] = self.api.generate_hash(token)
        else:
            # Authorization is required, silently auth the user
            self.api.connect()
            kwargs['__conduit__'] = self.api._conduit

        url = urlparse.urlparse(self.api.host)
        if url.scheme == 'https':
            conn = httplib.HTTPSConnection(url.netloc, timeout=self.api.timeout)
        else:
            conn = httplib.HTTPConnection(url.netloc, timeout=self.api.timeout)

        path = url.path + '%s.%s' % (self.method, self.endpoint)

        headers = {
            'User-Agent': 'python-phabricator/%s' % str(self.api.clientVersion),
            'Content-Type': 'application/x-www-form-urlencoded'
        }

        body = urlencode({
            "params": json.dumps(kwargs),
            "output": self.api.response_format
        })

        # TODO: Use HTTP "method" from interfaces.json
        conn.request('POST', path, body, headers)
        response = conn.getresponse()
        data = self._parse_response(response.read())

        return Result(data['result'])

    def _parse_response(self, data):
        # Process the response back to python
        parsed = self.api.formats[self.api.response_format](data)

        # Errors return 200, so check response content for exception
        if parsed['error_code']:
            raise APIError(parsed['error_code'], parsed['error_info'])

        return parsed

class Phabricator(Resource):
    formats = {
        'json': lambda x: json.loads(x),
    }

    def __init__(self, username=None, certificate=None, host=None,
            timeout=5, response_format='json', token=None, **kwargs):

        # Set values in ~/.arcrc as defaults
        if ARCRC:
            self.host = host if host else ARCRC['hosts'].keys()[0]
            if token or ARCRC['hosts'][self.host].has_key('token'):
                self.token = token if token else ARCRC['hosts'][self.host]['token']
            else:
                self.username = username if username else ARCRC['hosts'][self.host]['user']
                self.certificate = certificate if certificate else ARCRC['hosts'][self.host]['cert']
        else:
            self.host = host
            self.token = token
            self.username = username
            self.certificate = certificate

        self.timeout = timeout
        self.response_format = response_format
        self.client = 'python-phabricator'
        self.clientVersion = 1
        self.clientDescription = socket.gethostname() + ':python-phabricator'
        self._conduit = None

        super(Phabricator, self).__init__(self)

    def _request(self, **kwargs):
        raise SyntaxError('You cannot call the Conduit API without a resource.')

    def connect(self):
        if self.token:
            self._conduit = {
                'token': self.token
            }
            return

        auth = Resource(api=self, method='conduit', endpoint='connect')

        response = auth(user=self.username, host=self.host,
                client=self.client, clientVersion=self.clientVersion)

        self._conduit = {
            'sessionKey': response.sessionKey,
            'connectionID': response.connectionID
        }

    def generate_hash(self, token):
        source_string = (token + self.api.certificate).encode('utf-8')
        return hashlib.sha1(source_string).hexdigest()

    def update_interfaces(self):
        query = Resource(api=self, method='conduit', endpoint='query')

        interfaces = query()

        self.interface = parse_interfaces(interfaces)
<commit_msg>Safer config and interface loading mechanisms using `with open() as f:`<commit_after>"""
python-phabricator
------------------
>>> api = phabricator.Phabricator()
>>> api.user.whoami().userName
'example'

For more endpoints, see https://secure.phabricator.com/conduit/

"""
try:
    __version__ = __import__('pkg_resources') \
        .get_distribution('phabricator').version
except:
    __version__ = 'unknown'

import collections
import copy
import hashlib
import json
import os.path
import re
import socket
import time

from ._compat import (
    MutableMapping, iteritems, string_types, httplib, urlparse, urlencode,
)

__all__ = ['Phabricator']

ON_WINDOWS = os.name == 'nt'
CURRENT_DIR = os.getcwd()

# Default Phabricator interfaces
INTERFACES = {}
with open(os.path.join(os.path.dirname(__file__), 'interfaces.json')) as fobj:
    INTERFACES = json.load(fobj)

# Load arc config
ARC_CONFIGS = (
    # System config
    os.path.join(
        os.environ['ProgramData'],
        'Phabricator',
        'Arcanist',
        'config'
    ) if ON_WINDOWS else os.path.join('/etc', 'arcconfig'),

    # User config
    os.path.join(
        os.environ['AppData'] if ON_WINDOWS else os.path.expanduser('~'),
        '.arcrc'
    ),

    # Project config
    os.path.join(CURRENT_DIR, '.arcconfig'),

    # Local project config
    os.path.join(CURRENT_DIR, '.git', 'arc', 'config'),
)

ARCRC = {}
for conf in ARC_CONFIGS:
    if os.path.exists(conf):
        with open(conf, 'r') as fobj:
            ARCRC.update(json.load(fobj))

# Map Phabricator types to Python types
PARAM_TYPE_MAP = {
    # int types
    'int': int,
    'uint': int,
    'revisionid': int,
    'revision_id': int,
    'diffid': int,
    'diff_id': int,
    'id': int,
    'enum': int,

    # bool types
    'bool': bool,

    # dict types
    'map': dict,
    'dict': dict,

    # list types
    'list': list,

    # tuple types
    'pair': tuple,

    # str types
    'str': string_types,
    'string': string_types,
    'phid': string_types,
    'guids': string_types,
    'type': string_types,
}

STR_RE = re.compile(r'([a-zA-Z_]+)')

def map_param_type(param_type):
    """
    Perform param type mapping
    This requires a bit of logic since this isn't standardized.
    If a type doesn't map, assume str
    """
    m = STR_RE.match(param_type)
    main_type = m.group(0)

    if main_type in ('list', 'array'):
        info = param_type.replace(' ', '').split('<', 1)

        # Handle no sub-type: "required list"
        sub_type = info[1] if len(info) > 1 else 'str'

        # Handle list of pairs: "optional list<pair<callsign, path>>"
        sub_match = STR_RE.match(sub_type)
        sub_type = sub_match.group(0).lower()

        return [PARAM_TYPE_MAP.setdefault(sub_type, string_types)]

    return PARAM_TYPE_MAP.setdefault(main_type, string_types)

def parse_interfaces(interfaces):
    """
    Parse the conduit.query json dict response
    This performs the logic of parsing the non-standard params dict
        and then returning a dict Resource can understand
    """
    parsed_interfaces = collections.defaultdict(dict)

    for m, d in iteritems(interfaces):
        app, func = m.split('.', 1)

        method = parsed_interfaces[app][func] = {}

        # Make default assumptions since these aren't provided by Phab
        method['formats'] = ['json', 'human']
        method['method'] = 'POST'

        method['optional'] = {}
        method['required'] = {}

        for name, type_info in iteritems(dict(d['params'])):
            # Usually in the format: <optionality> <param_type>
            info_pieces = type_info.split(' ', 1)

            # If optionality isn't specified, assume required
            if info_pieces[0] not in ('optional', 'required'):
                optionality = 'required'
                param_type = info_pieces[0]
            # Just make an optional string for "ignored" params
            elif info_pieces[0] == 'ignored':
                optionality = 'optional'
                param_type = 'string'
            else:
                optionality = info_pieces[0]
                param_type = info_pieces[1]

            # This isn't validated by the client
            if param_type.startswith('nonempty'):
                optionality = 'required'
                param_type = param_type[9:]

            method[optionality][name] = map_param_type(param_type)

    return dict(parsed_interfaces)

class InterfaceNotDefined(NotImplementedError):
    pass

class APIError(Exception):
    def __init__(self, code, message):
        self.code = code
        self.message = message

    def __str__(self):
        return '%s: %s' % (self.code, self.message)

class InvalidAccessToken(APIError):
    pass

class Result(MutableMapping):
    def __init__(self, response):
        self.response = response

    def __getitem__(self, key):
        return self.response[key]

    __getattr__ = __getitem__

    def __setitem__(self, key, value):
        self.response[key] = value

    def __delitem__(self, key):
        del self.response[key]

    def __iter__(self):
        return iter(self.response)

    def __len__(self):
        return len(self.response.keys())

    def __repr__(self):
        return '<%s: %s>' % (type(self).__name__, repr(self.response))

class Resource(object):
    def __init__(self, api, interface=None, endpoint=None, method=None):
        self.api = api
        self.interface = interface or copy.deepcopy(parse_interfaces(INTERFACES))
        self.endpoint = endpoint
        self.method = method

    def __getattr__(self, attr):
        if attr in getattr(self, '__dict__'):
            return getattr(self, attr)
        interface = self.interface
        if attr not in interface:
            interface[attr] = {}
        return Resource(self.api, interface[attr], attr, self.endpoint)

    def __call__(self, **kwargs):
        return self._request(**kwargs)

    def _request(self, **kwargs):
        # Check for missing variables
        resource = self.interface

        def validate_kwarg(key, target):
            # Always allow list
            if isinstance(key, list):
                return all([validate_kwarg(x, target[0]) for x in key])
            return isinstance(key, tuple(target) if isinstance(target, list) else target)

        for key, val in resource.get('required', {}).items():
            if key not in [x.split(':')[0] for x in kwargs.keys()]:
                raise ValueError('Missing required argument: %s' % key)
            if isinstance(kwargs.get(key), list) and not isinstance(val, list):
                raise ValueError('Wrong argument type: %s is not a list' % key)
            elif not validate_kwarg(kwargs.get(key), val):
                if isinstance(val, list):
                    raise ValueError('Wrong argument type: %s is not a list of %ss' % (key, val[0]))
                raise ValueError('Wrong argument type: %s is not a %s' % (key, val))

        conduit = self.api._conduit

        if conduit:
            # Already authenticated, add session key to json data
            kwargs['__conduit__'] = conduit
        elif self.method == 'conduit' and self.endpoint == 'connect':
            # Not authenticated, requesting new session key
            token = str(int(time.time()))
            kwargs['authToken'] = token
            kwargs['authSignature'] = self.api.generate_hash(token)
        else:
            # Authorization is required, silently auth the user
            self.api.connect()
            kwargs['__conduit__'] = self.api._conduit

        url = urlparse.urlparse(self.api.host)
        if url.scheme == 'https':
            conn = httplib.HTTPSConnection(url.netloc, timeout=self.api.timeout)
        else:
            conn = httplib.HTTPConnection(url.netloc, timeout=self.api.timeout)

        path = url.path + '%s.%s' % (self.method, self.endpoint)

        headers = {
            'User-Agent': 'python-phabricator/%s' % str(self.api.clientVersion),
            'Content-Type': 'application/x-www-form-urlencoded'
        }

        body = urlencode({
            "params": json.dumps(kwargs),
            "output": self.api.response_format
        })

        # TODO: Use HTTP "method" from interfaces.json
        conn.request('POST', path, body, headers)
        response = conn.getresponse()
        data = self._parse_response(response.read())

        return Result(data['result'])

    def _parse_response(self, data):
        # Process the response back to python
        parsed = self.api.formats[self.api.response_format](data)

        # Errors return 200, so check response content for exception
        if parsed['error_code']:
            raise APIError(parsed['error_code'], parsed['error_info'])

        return parsed

class Phabricator(Resource):
    formats = {
        'json': lambda x: json.loads(x),
    }

    def __init__(self, username=None, certificate=None, host=None,
            timeout=5, response_format='json', token=None, **kwargs):

        # Set values in ~/.arcrc as defaults
        if ARCRC:
            self.host = host if host else ARCRC['hosts'].keys()[0]
            if token or ARCRC['hosts'][self.host].has_key('token'):
                self.token = token if token else ARCRC['hosts'][self.host]['token']
            else:
                self.username = username if username else ARCRC['hosts'][self.host]['user']
                self.certificate = certificate if certificate else ARCRC['hosts'][self.host]['cert']
        else:
            self.host = host
            self.token = token
            self.username = username
            self.certificate = certificate

        self.timeout = timeout
        self.response_format = response_format
        self.client = 'python-phabricator'
        self.clientVersion = 1
        self.clientDescription = socket.gethostname() + ':python-phabricator'
        self._conduit = None

        super(Phabricator, self).__init__(self)

    def _request(self, **kwargs):
        raise SyntaxError('You cannot call the Conduit API without a resource.')

    def connect(self):
        if self.token:
            self._conduit = {
                'token': self.token
            }
            return

        auth = Resource(api=self, method='conduit', endpoint='connect')

        response = auth(user=self.username, host=self.host,
                client=self.client, clientVersion=self.clientVersion)

        self._conduit = {
            'sessionKey': response.sessionKey,
            'connectionID': response.connectionID
        }

    def generate_hash(self, token):
        source_string = (token + self.api.certificate).encode('utf-8')
        return hashlib.sha1(source_string).hexdigest()

    def update_interfaces(self):
        query = Resource(api=self, method='conduit', endpoint='query')

        interfaces = query()

        self.interface = parse_interfaces(interfaces)
<|endoftext|>


<commit_before>"""
Views (in ``vtk``) for end points and fibers.
"""
import vtk
from matplotlib.colors import colorConverter
import numpy as np

_vtk_major_version = vtk.vtkVersion.GetVTKMajorVersion()

class ViewFiber():
    """
    A fiber is represented as a vtkPolyData, and attached to a vtk renderer.
    """
    def __init__(self, fiber, ren, **kwargs):
        self.init_draw(fiber, ren, **kwargs)
        self.draw(fiber)

    def init_draw(self, fiber, ren, **kwargs):
        # Convert color to rgb.
        my_color = kwargs.get('color', 'b')
        label = kwargs.get('label', '')
        rgb_color = 255 * colorConverter.to_rgba(my_color)

        # Create a vtkPoints instance.
        points = vtk.vtkPoints()
        self.points = points
        poly_line = vtk.vtkPolyLine()
        nb_points = fiber.get_nb_points()
        poly_line.GetPointIds().SetNumberOfIds(nb_points)
        for i in range(nb_points):
            poly_line.GetPointIds().SetId(i, i)
            self.points.InsertNextPoint(fiber.get_points()[i])
        cells = vtk.vtkCellArray()
        cells.InsertNextCell(poly_line)
        poly_data = vtk.vtkPolyData()
        poly_data.SetPoints(points)
        poly_data.SetLines(cells)

        # The tube is wrapped around the generated streamline.
        stream_tube = vtk.vtkTubeFilter()
        self.tube = stream_tube
        stream_tube.SetInput(poly_data)
        stream_tube.SetRadius(fiber.get_radius())
        stream_tube.SetNumberOfSides(12)
        stream_tube.SetCapping(True)
        map_stream_tube = vtk.vtkPolyDataMapper()
        map_stream_tube.SetInputConnection(stream_tube.GetOutputPort())
        actor = vtk.vtkActor()
        actor.SetMapper(map_stream_tube)
        actor.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        ren.AddActor(actor)
        self.actor = actor

        # Add label to the figure
        text_source = vtk.vtkVectorText()
        text_source.SetText(label)
        text_source.Update()
        text_mapper = vtk.vtkPolyDataMapper()
        text_mapper.SetInputConnection(text_source.GetOutputPort())

        text_actor1 = vtk.vtkActor()
        text_actor1.SetMapper(text_mapper)
        text_actor1.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        text_actor1.SetScale(1.0)
        center1 = text_actor1.GetCenter()
        text_actor1.RotateX(90)
        text_actor1.AddPosition(fiber.get_points()[-1] * 1.2 - center1)

        text_actor2 = vtk.vtkActor()
        text_actor2.SetMapper(text_mapper)
        text_actor2.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        text_actor2.SetScale(1.0)
        center2 = text_actor2.GetCenter()
        text_actor2.RotateX(90)
        text_actor2.AddPosition(fiber.get_points()[0] * 1.2 - center2)

        ren.AddActor(text_actor1)
        ren.AddActor(text_actor2)

    def get_actor(self):
        return self.actor

    def draw(self, fiber):
        nb_points = fiber.get_nb_points()
        for i in range(nb_points):
            self.points.SetPoint(i, fiber.get_points()[i])

        self.actor.GetProperty().SetOpacity(1.0)

        # The tube is wrapped around the generated streamline.
        self.tube.SetRadius(fiber.get_radius())

    def notify(self, fiber):
        """If the internal structure of a fiber has changed, we need to modify the
	vtk objects accordingly."""
        self.draw(fiber)

class ViewIsotropicRegion():
    """
    An isotropic region (sphere) is represented as a vtkPolyData, and attached
    to a vtk renderer.
    """
    def __init__(self, isotropic_region, ren, **kwargs):
        self.init_draw(isotropic_region.center, isotropic_region.radius, ren, **kwargs)

    def init_draw(self, center, radius, ren, **kwargs):
        # Convert color to rgb.
        my_color = kwargs.get('color', 'b')
        label = kwargs.get('label', '')
        rgb_color = 255 * colorConverter.to_rgba(my_color)

        sphere = vtk.vtkSphereSource()
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(20)
        sphere.SetPhiResolution(40)

        sphere_mapper = vtk.vtkPolyDataMapper()
        sphere_mapper.SetInputConnection(sphere.GetOutputPort())

        sphere_actor = vtk.vtkActor()
        sphere_actor.SetMapper(sphere_mapper)
        sphere_actor.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])

        transform = vtk.vtkTransform()
        transform.Translate(center[0], center[1], center[2])
        sphere_actor.SetUserMatrix(transform.GetMatrix())

        ren.AddActor(sphere_actor)
        self.actor = sphere_actor

        # Add label to the figure
        text_source = vtk.vtkVectorText()
        text_source.SetText(label)
        text_source.Update()
        text_mapper = vtk.vtkPolyDataMapper()
        text_mapper.SetInputConnection(text_source.GetOutputPort())
        text_actor1 = vtk.vtkActor()
        text_actor1.SetMapper(text_mapper)
        text_actor1.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        text_actor1.SetScale(0.05)
        center1 = text_actor1.GetCenter()
        text_actor1.RotateX(90)
        text_actor1.AddPosition(center)
        text_actor1.AddPosition([0., 0., radius + 0.1])
        ren.AddActor(text_actor1)

    def get_actor(self):
        return self.actor

class ViewSphere():
    """Simply adds a translucent sphere (representing cortical surface)."""
    def __init__(self, radius, ren, **kwargs):
        self.init_draw(radius, ren, **kwargs)

    def init_draw(self, radius, ren, **kwargs):
        # Convert color to rgb.
        opacity = kwargs.get('opacity', 0.5)
        rgb_color = [0.9, 0.9, 1.0]

        sphere = vtk.vtkSphereSource()
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(20)
        sphere.SetPhiResolution(40)

        sphere_mapper = vtk.vtkPolyDataMapper()
        sphere_mapper.SetInputConnection(sphere.GetOutputPort())

        sphere_actor = vtk.vtkActor()
        sphere_actor.SetMapper(sphere_mapper)
        sphere_actor.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        sphere_actor.GetProperty().SetOpacity(opacity)

        ren.AddActor(sphere_actor)
<commit_msg>RF/PEP8/DOC: Removed unused call to vtkVersion, and clarified a bit the DOCs. Also fixed some line widths as per PEP8.<commit_after>"""
Views (in ``vtk``) for fibers and isotropic regions.
"""
import vtk
from matplotlib.colors import colorConverter
import numpy as np

class ViewFiber():
    """
    A fiber is represented as a vtkPolyData, and attached to a vtk renderer.
    """
    def __init__(self, fiber, ren, **kwargs):
        self.init_draw(fiber, ren, **kwargs)
        self.draw(fiber)

    def init_draw(self, fiber, ren, **kwargs):
        # Convert color to rgb.
        my_color = kwargs.get('color', 'b')
        label = kwargs.get('label', '')
        rgb_color = 255 * colorConverter.to_rgba(my_color)

        # Create a vtkPoints instance.
        points = vtk.vtkPoints()
        self.points = points
        poly_line = vtk.vtkPolyLine()
        nb_points = fiber.get_nb_points()
        poly_line.GetPointIds().SetNumberOfIds(nb_points)
        for i in range(nb_points):
            poly_line.GetPointIds().SetId(i, i)
            self.points.InsertNextPoint(fiber.get_points()[i])
        cells = vtk.vtkCellArray()
        cells.InsertNextCell(poly_line)
        poly_data = vtk.vtkPolyData()
        poly_data.SetPoints(points)
        poly_data.SetLines(cells)

        # The tube is wrapped around the generated streamline.
        stream_tube = vtk.vtkTubeFilter()
        self.tube = stream_tube
        stream_tube.SetInput(poly_data)
        stream_tube.SetRadius(fiber.get_radius())
        stream_tube.SetNumberOfSides(12)
        stream_tube.SetCapping(True)
        map_stream_tube = vtk.vtkPolyDataMapper()
        map_stream_tube.SetInputConnection(stream_tube.GetOutputPort())
        actor = vtk.vtkActor()
        actor.SetMapper(map_stream_tube)
        actor.GetProperty().SetColor(rgb_color[0], rgb_color[1], rgb_color[2])
        ren.AddActor(actor)
        self.actor = actor

        # Add label to the figure
        text_source = vtk.vtkVectorText()
        text_source.SetText(label)
        text_source.Update()
        text_mapper = vtk.vtkPolyDataMapper()
        text_mapper.SetInputConnection(text_source.GetOutputPort())

        text_actor1 = vtk.vtkActor()
        text_actor1.SetMapper(text_mapper)
        text_actor1.GetProperty().SetColor(rgb_color[0],
                                           rgb_color[1],
                                           rgb_color[2])
        text_actor1.SetScale(1.0)
        center1 = text_actor1.GetCenter()
        text_actor1.RotateX(90)
        text_actor1.AddPosition(fiber.get_points()[-1] * 1.2 - center1)

        text_actor2 = vtk.vtkActor()
        text_actor2.SetMapper(text_mapper)
        text_actor2.GetProperty().SetColor(rgb_color[0],
                                           rgb_color[1],
                                           rgb_color[2])
        text_actor2.SetScale(1.0)
        center2 = text_actor2.GetCenter()
        text_actor2.RotateX(90)
        text_actor2.AddPosition(fiber.get_points()[0] * 1.2 - center2)

        ren.AddActor(text_actor1)
        ren.AddActor(text_actor2)

    def get_actor(self):
        return self.actor

    def draw(self, fiber):
        nb_points = fiber.get_nb_points()
        for i in range(nb_points):
            self.points.SetPoint(i, fiber.get_points()[i])

        self.actor.GetProperty().SetOpacity(1.0)

        # The tube is wrapped around the generated streamline.
        self.tube.SetRadius(fiber.get_radius())

    def notify(self, fiber):
        """
        If the internal structure of a fiber has changed, we need to modify
        the vtk objects accordingly."""
        self.draw(fiber)

class ViewIsotropicRegion():
    """
    An isotropic region (sphere) is represented as a vtkPolyData, and attached
    to a vtk renderer.
    """
    def __init__(self, isotropic_region, ren, **kwargs):
        self.init_draw(isotropic_region.center, isotropic_region.radius, ren,
                       **kwargs)

    def init_draw(self, center, radius, ren, **kwargs):
        # Convert color to rgb.
        my_color = kwargs.get('color', 'b')
        label = kwargs.get('label', '')
        rgb_color = 255 * colorConverter.to_rgba(my_color)

        sphere = vtk.vtkSphereSource()
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(20)
        sphere.SetPhiResolution(40)

        sphere_mapper = vtk.vtkPolyDataMapper()
        sphere_mapper.SetInputConnection(sphere.GetOutputPort())

        sphere_actor = vtk.vtkActor()
        sphere_actor.SetMapper(sphere_mapper)
        sphere_actor.GetProperty().SetColor(rgb_color[0],
                                            rgb_color[1],
                                            rgb_color[2])

        transform = vtk.vtkTransform()
        transform.Translate(center[0], center[1], center[2])
        sphere_actor.SetUserMatrix(transform.GetMatrix())

        ren.AddActor(sphere_actor)
        self.actor = sphere_actor

        # Add label to the figure
        text_source = vtk.vtkVectorText()
        text_source.SetText(label)
        text_source.Update()
        text_mapper = vtk.vtkPolyDataMapper()
        text_mapper.SetInputConnection(text_source.GetOutputPort())
        text_actor1 = vtk.vtkActor()
        text_actor1.SetMapper(text_mapper)
        text_actor1.GetProperty().SetColor(rgb_color[0],
                                           rgb_color[1],
                                           rgb_color[2])
        text_actor1.SetScale(0.05)
        center1 = text_actor1.GetCenter()
        text_actor1.RotateX(90)
        text_actor1.AddPosition(center)
        text_actor1.AddPosition([0., 0., radius + 0.1])
        ren.AddActor(text_actor1)

    def get_actor(self):
        return self.actor

class ViewSphere():
    """
    Simply adds a translucent sphere (representing cortical surface).
    """
    def __init__(self, radius, ren, **kwargs):
        self.init_draw(radius, ren, **kwargs)

    def init_draw(self, radius, ren, **kwargs):
        # Convert color to rgb.
        opacity = kwargs.get('opacity', 0.5)
        rgb_color = [0.9, 0.9, 1.0]

        sphere = vtk.vtkSphereSource()
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(20)
        sphere.SetPhiResolution(40)

        sphere_mapper = vtk.vtkPolyDataMapper()
        sphere_mapper.SetInputConnection(sphere.GetOutputPort())

        sphere_actor = vtk.vtkActor()
        sphere_actor.SetMapper(sphere_mapper)
        sphere_actor.GetProperty().SetColor(rgb_color[0],
                                            rgb_color[1],
                                            rgb_color[2])
        sphere_actor.GetProperty().SetOpacity(opacity)

        ren.AddActor(sphere_actor)
<|endoftext|>


<commit_before># coding: utf-8
from pianette.utils import Debug
from flask import Flask, render_template, request
from threading import Thread

app = Flask(__name__, template_folder='../templates')

import logging
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

@app.route('/', methods = ['GET'])
def home():
    return render_template('index.html', url_root=request.url_root)

@app.route('/console/play', methods = ['POST'])
def console_play():
    command = "console.play %s" % (request.form.get('command'))
    app.pianette.inputcmds(command, source="api")
    return ('1', 200)

class PianetteApi:

    def __init__(self, configobj=None, pianette=None, **kwargs):
        super().__init__(**kwargs)
        self.configobj = configobj
        app.pianette = pianette

        Debug.println("INFO", "Starting API thread")
        t = Thread(target=self.startApi)
        t.daemon = True
        t.start()

    def startApi(self):
        app.run(debug=False, threaded=True)
<commit_msg>Add external IP for Flask<commit_after># coding: utf-8
from pianette.utils import Debug
from flask import Flask, render_template, request
from threading import Thread

app = Flask(__name__, template_folder='../templates')

import logging
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

@app.route('/', methods = ['GET'])
def home():
    return render_template('index.html', url_root=request.url_root)

@app.route('/console/play', methods = ['POST'])
def console_play():
    command = "console.play %s" % (request.form.get('command'))
    app.pianette.inputcmds(command, source="api")
    return ('1', 200)

class PianetteApi:

    def __init__(self, configobj=None, pianette=None, **kwargs):
        super().__init__(**kwargs)
        self.configobj = configobj
        app.pianette = pianette

        Debug.println("INFO", "Starting API thread")
        t = Thread(target=self.startApi)
        t.daemon = True
        t.start()

    def startApi(self):
        app.run(debug=False, threaded=True, host=0.0.0.0)
<|endoftext|>


<commit_before>"""Stock ticker module for ppbot.
@package ppbot

Given a stock symbol, will lookup the current trading price.
Uses Google's Finance API

@syntax stock <symbol>

"""
import urllib2
import string
import json
from xml.dom.minidom import parseString

from piebot.modules import *

class Stock(Module):

    def __init__(self, *args, **kwargs):
        """Constructor"""

        Module.__init__(self, kwargs=kwargs)

        self.url = "http://www.google.com/finance/info?infotype=infoquoteall&q=%s"

    def _register_events(self):
        """Register module commands."""

        self.add_command('stock')

    def stock(self, event):
        """Action to react/respond to user calls."""

        if self.num_args == 1:
            # need to fetch the weather and parse it
            symbol = event['args'][0]
            try:
                stock_info = self.lookup_symbol(symbol)
                # stylize the message output
                try:
                    change = float(stock_info['ecp'])
                except:
                    change = float(stock_info['cp'])
                if change > 0:
                    color = "\x033"
                elif change < 0:
                    color = "\x034"
                else:
                    color = ""
                stock_info['color'] = color
		try:
                    message1 = "%(name)s (%(e)s:%(t)s) - %(el)s (%(color)s%(ec)s\x0f,%(color)s %(ecp)s%%\x0f) - 52week high/low: (%(hi52)s/%(lo52)s) - MktCap: %(mc)s - P/E: %(pe)s" % (stock_info)
                except:
                    message1 = "%(name)s (%(e)s:%(t)s) - %(l)s (%(color)s%(c)s\x0f,%(color)s %(cp)s%%\x0f) - 52week high/low: (%(hi52)s/%(lo52)s) - MktCap: %(mc)s - P/E: %(pe)s" % (stock_info)
                # send the messages
                self.reply(message1)
            except:
                self.reply('Could not find symbol "%s"' % symbol)
                import traceback
                traceback.print_exc()
        else:
            self.syntax_message(event['nick'], '.stock <symbol>')

    def lookup_symbol(self, symbol):
        """Connects to google's secret finance API and parses the receiving json for the stock info."""

        # make the parser, and send the xml to be parsed
        result = urllib2.urlopen(self.url % symbol).read()
        stock = json.loads(result[4:])
        return stock[0]
<commit_msg>Display and differentiate between after hours and day change<commit_after>"""Stock ticker module for ppbot.
@package ppbot

Given a stock symbol, will lookup the current trading price.
Uses Google's Finance API

@syntax stock <symbol>

"""
import urllib2
import string
import json
from xml.dom.minidom import parseString

from piebot.modules import *

class Stock(Module):

    def __init__(self, *args, **kwargs):
        """Constructor"""

        Module.__init__(self, kwargs=kwargs)

        self.url = "http://www.google.com/finance/info?infotype=infoquoteall&q=%s"

    def _register_events(self):
        """Register module commands."""

        self.add_command('stock')
        self.add_event('pubmsg', 'parse_message')

    def stock(self, event):
        """Action to react/respond to user calls."""

        if self.num_args == 1:
            # need to fetch the weather and parse it
            symbol = event['args'][0]

            try:
                self.reply(self.lookup_symbol(symbol))
            except:
                self.reply('Could not find symbol "%s"' % symbol)
                import traceback
                traceback.print_exc()
        else:
            self.syntax_message(event['nick'], '.stock <symbol>')

    def parse_message(self, event):
        match = re.search('\$([a-zA-Z]{2,10})', event['message'])

        if match:
            try:
                self.reply(self.lookup_symbol(match.group(1)))
            except:
                pass

    def lookup_symbol(self, symbol):
        """Given a stock symbol, returns a formatted string to be displayed"""

        stock_info = self.request_symbol(symbol)
        # stylize the message output
        change = float(stock_info['cp'])
        color = "\x033" if change > 0 else "\x034" if change < 0 else ""
        stock_info['color'] = color

        try:
            ah_change = float(stock_info['ecp'])
            color = "\x033" if ah_change > 0 else "\x034" if ah_change < 0 else ""
            stock_info['ah_color'] = color
            stock_info['afterhours'] = "%(el)s (%(ah_color)s%(ec)s\x0f,%(ah_color)s %(ecp)s%%\x0f)" % stock_info
        except KeyError:
            pass

        try:
            message1 = "%(name)s (%(e)s:%(t)s) - Now: %(afterhours)s - Today: %(l)s (%(color)s%(c)s\x0f,%(color)s %(cp)s%%\x0f) - 52week high/low: (%(hi52)s/%(lo52)s) - MktCap: %(mc)s - P/E: %(pe)s" % (stock_info)
        except:
            message1 = "%(name)s (%(e)s:%(t)s) - %(l)s (%(color)s%(c)s\x0f,%(color)s %(cp)s%%\x0f) - 52week high/low: (%(hi52)s/%(lo52)s) - MktCap: %(mc)s - P/E: %(pe)s" % (stock_info)

        return message1

    def request_symbol(self, symbol):
        """Connects to google's secret finance API and parses the receiving json for the stock info."""

        # make the parser, and send the xml to be parsed
        result = urllib2.urlopen(self.url % symbol).read()
        result = result.replace('\\x', '\\u00')
        stock = json.loads(result[4:])
        return stock[0]
<|endoftext|>


<commit_before># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Alexey Anisenkov, anisyonk@cern.ch, 2018-2019
# - Paul Nilsson, paul.nilsson@cern.ch, 2019

"""
The implementation of data structure to host queuedata settings.

The main reasons for such incapsulation are to
 - apply in one place all data validation actions (for attributes and values)
 - introduce internal information schema (names of attribues) to remove dependency
 with data structrure, formats, names from external sources (e.g. AGIS/CRIC)

This module should be standalone as much as possible and even does not depend
on the configuration settings
(for that purposed `PilotConfigProvider` can be user to customize data)

:author: Alexey Anisenkov
:contact: anisyonk@cern.ch
:date: January 2018
"""

import re

from .basedata import BaseData

import logging
logger = logging.getLogger(__name__)

class QueueData(BaseData):
    """
        High-level object to host all queuedata settings associated to given PandaQueue
    """

    # ## put explicit list of all the attributes with comments for better inline-documentation by sphinx
    # ## FIX ME LATER: use proper doc format
    # ## incomplete list of attributes .. to be extended once becomes used

    name = ""       # Name of Panda Queue
    resource = ""   # Name of Panda Resource
    appdir = ""     #
    catchall = ""   #

    platform = ""     # cmtconfig value
    container_options = ""  # singularity only options? to be reviewed and forced to be a dict (support options for other containers?)
    container_type = {}  # dict of container names by user as a key

    copytools = None
    acopytools = None

    ## allowed protocol schemas for requested copytool/activity
    ## if passed value (per activity) is a list, then given schemas will be used for all allowed copytools
    ## in case of dict-based value, it specifies allowed schemas per copytool for given activity
    ## e.g. {'pr':['root', 'srm'], 'pw':['webdav'], 'default':['root']}
    ##      {'pr': {'gfalcopy':['webdav'], 'pw':{'lsm':['root']}}}
    acopytools_schemas = {}

    astorages = None
    aprotocols = None

    state = None  # AGIS PQ state, e.g. ACTIVE
    status = ""   # PQ status, e.g. online
    site = None   # ATLAS Site name

    direct_access_lan = False  # Prefer remote io (True) or use only copy2scratch method (False) for stage-in over LAN
    direct_access_wan = False  # Prefer remote io (True) or use only copy2scratch method (False) for stage-in over WAN

    allow_lan = True  # Allow LAN access (whatever method) for stage-in
    allow_wan = False  # Allow WAN access (whatever method) for stage-in

    use_pcache = False

    maxwdir = 0    # in MB
    maxrss = 0

    timefloor = 0  # The maximum time during which the pilot is allowed to start a new job, in seconds
    corecount = 1  #

    maxtime = 0  # maximum allowed lifetime for pilot to run on the resource (0 will be ignored, fallback to default)

    pledgedcpu = 0  #
    es_stageout_gap = 0  ## time gap value in seconds for ES stageout

    # specify the type of attributes for proper data validation and casting
    _keys = {int: ['timefloor', 'maxwdir', 'pledgedcpu', 'es_stageout_gap',
                   'corecount', 'maxrss', 'maxtime'],
             str: ['name', 'type', 'appdir', 'catchall', 'platform', 'container_options', 'container_type',
                   'resource', 'state', 'status', 'site'],
             dict: ['copytools', 'acopytools', 'astorages', 'aprotocols', 'acopytools_schemas'],
             bool: ['allow_lan', 'allow_wan', 'direct_access_lan', 'direct_access_wan', 'use_pcache']
             }

    def __init__(self, data):
        """
            :param data: input dictionary of queue data settings
        """

        self.load(data)

        # DEBUG
        #import pprint
        #logger.debug('initialize QueueData from raw:\n%s' % pprint.pformat(data))
        logger.debug('Final parsed QueueData content:\n%s' % self)

    def load(self, data):
        """
            Construct and initialize data from ext source
            :param data: input dictionary of queue data settings
        """

        # the translation map of the queue data attributes from external data to internal schema
        # 'internal_name':('ext_name1', 'extname2_if_any')
        # 'internal_name2':'ext_name3'

        # first defined ext field will be used
        # if key is not explicitly specified then ext name will be used as is
        ## fix me later to proper internal names if need

        kmap = {
            'name': 'nickname',
            'resource': 'panda_resource',
            'platform': 'cmtconfig',
            'site': ('atlas_site', 'gstat'),
            'es_stageout_gap': 'zip_time_gap',
        }

        self._load_data(data, kmap)

    def resolve_allowed_schemas(self, activity, copytool=None):
        """
            Resolve list of allowed schemas for given activity and requested copytool based on `acopytools_schemas` settings
            :param activity: str or ordered list of transfer activity names to resolve acopytools related data
            :return: list of protocol schemes
        """

        if not activity:
            activity = 'default'
        try:
            if isinstance(activity, basestring):  # Python 2
                activity = [activity]
        except Exception:
            if isinstance(activity, str):  # Python 3
                activity = [activity]

        if 'default' not in activity:
            activity = activity + ['default']

        adat = {}
        for aname in activity:
            adat = self.acopytools_schemas.get(aname)
            if adat:
                break
        if not adat:
            return []

        if not isinstance(adat, dict):
            adat = {'default': adat}

        if not copytool or copytool not in adat:
            copytool = 'default'

        return adat.get(copytool) or []

    def clean(self):
        """
            Validate and finally clean up required data values (required object properties) if need
            :return: None
        """

        # validate es_stageout_gap value
        if not self.es_stageout_gap:
            is_opportunistic = self.pledgedcpu and self.pledgedcpu == -1
            self.es_stageout_gap = 600 if is_opportunistic else 7200  ## 10 munites for opportunistic or 5 hours for normal resources

        # validate container_options: extract from the catchall if not set
        if not self.container_options and self.catchall:  ## container_options is considered for the singularity container, FIX ME LATER IF NEED
            # expected format
            # of catchall = "singularity_options=\'-B /etc/grid-security/certificates,/cvmfs,${workdir} --contain\'"
            pattern = re.compile("singularity_options=['\"]?([^'\"]+)['\"]?")  ### FIX ME LATER: move to proper args parsing via shlex at Job class
            found = re.findall(pattern, self.catchall)
            if found:
                self.container_options = found[0]
                logger.info('container_options extracted from catchall: %s' % self.container_options)

        # verify container_options: add the workdir if missing
        if self.container_options:
            if "${workdir}" not in self.container_options and " --contain" in self.container_options:  ## reimplement with shlex later
                self.container_options = self.container_options.replace(" --contain", ",${workdir} --contain")
                logger.info("Note: added missing ${workdir} to container_options/singularity_options: %s" % self.container_options)

        pass

    ## custom function pattern to apply extra validation to the key values
    ##def clean__keyname(self, raw, value):
    ##  :param raw: raw value passed from ext source as input
    ##  :param value: preliminary cleaned and casted to proper type value
    ##
    ##    return value

    def clean__timefloor(self, raw, value):
        """
            Verify and validate value for the timefloor key (convert to seconds)
        """

        return value * 60

    def clean__container_type(self, raw, value):
        """
            Parse and prepare value for the container_type key
            Expected raw data in format 'container_name:user_name;'
            E.g. container_type = 'singularity:pilot;docker:wrapper'

            :return: dict of container names by user as a key
        """

        ret = {}
        val = value or ''
        for e in val.split(';'):
            dat = e.split(':')
            if len(dat) == 2:
                name, user = dat[0].strip(), dat[1].strip()
                ret[user] = name

        return ret

    def clean__container_options(self, raw, value):
        """
            Verify and validate value for the container_options key (remove bad values)
        """

        return value if value.lower() not in ['none'] else ''

    def clean__corecount(self, raw, value):
        """
            Verify and validate value for the corecount key (set to 1 if not set)
        """

        return value if value else 1
<commit_msg>Added is_cvmfs<commit_after># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Alexey Anisenkov, anisyonk@cern.ch, 2018-2019
# - Paul Nilsson, paul.nilsson@cern.ch, 2019

"""
The implementation of data structure to host queuedata settings.

The main reasons for such incapsulation are to
 - apply in one place all data validation actions (for attributes and values)
 - introduce internal information schema (names of attribues) to remove dependency
 with data structrure, formats, names from external sources (e.g. AGIS/CRIC)

This module should be standalone as much as possible and even does not depend
on the configuration settings
(for that purposed `PilotConfigProvider` can be user to customize data)

:author: Alexey Anisenkov
:contact: anisyonk@cern.ch
:date: January 2018
"""

import re

from .basedata import BaseData

import logging
logger = logging.getLogger(__name__)

class QueueData(BaseData):
    """
        High-level object to host all queuedata settings associated to given PandaQueue
    """

    # ## put explicit list of all the attributes with comments for better inline-documentation by sphinx
    # ## FIX ME LATER: use proper doc format
    # ## incomplete list of attributes .. to be extended once becomes used

    name = ""       # Name of Panda Queue
    resource = ""   # Name of Panda Resource
    appdir = ""     #
    catchall = ""   #

    platform = ""     # cmtconfig value
    container_options = ""  # singularity only options? to be reviewed and forced to be a dict (support options for other containers?)
    container_type = {}  # dict of container names by user as a key

    copytools = None
    acopytools = None

    ## allowed protocol schemas for requested copytool/activity
    ## if passed value (per activity) is a list, then given schemas will be used for all allowed copytools
    ## in case of dict-based value, it specifies allowed schemas per copytool for given activity
    ## e.g. {'pr':['root', 'srm'], 'pw':['webdav'], 'default':['root']}
    ##      {'pr': {'gfalcopy':['webdav'], 'pw':{'lsm':['root']}}}
    acopytools_schemas = {}

    astorages = None
    aprotocols = None

    state = None  # AGIS PQ state, e.g. ACTIVE
    status = ""   # PQ status, e.g. online
    site = None   # ATLAS Site name

    direct_access_lan = False  # Prefer remote io (True) or use only copy2scratch method (False) for stage-in over LAN
    direct_access_wan = False  # Prefer remote io (True) or use only copy2scratch method (False) for stage-in over WAN

    allow_lan = True  # Allow LAN access (whatever method) for stage-in
    allow_wan = False  # Allow WAN access (whatever method) for stage-in

    use_pcache = False

    maxwdir = 0    # in MB
    maxrss = 0

    timefloor = 0  # The maximum time during which the pilot is allowed to start a new job, in seconds
    corecount = 1  #

    maxtime = 0  # maximum allowed lifetime for pilot to run on the resource (0 will be ignored, fallback to default)

    pledgedcpu = 0  #
    es_stageout_gap = 0  ## time gap value in seconds for ES stageout

    # specify the type of attributes for proper data validation and casting
    _keys = {int: ['timefloor', 'maxwdir', 'pledgedcpu', 'es_stageout_gap',
                   'corecount', 'maxrss', 'maxtime'],
             str: ['name', 'type', 'appdir', 'catchall', 'platform', 'container_options', 'container_type',
                   'resource', 'state', 'status', 'site'],
             dict: ['copytools', 'acopytools', 'astorages', 'aprotocols', 'acopytools_schemas'],
             bool: ['allow_lan', 'allow_wan', 'direct_access_lan', 'direct_access_wan', 'is_cvmfs', 'use_pcache']
             }

    def __init__(self, data):
        """
            :param data: input dictionary of queue data settings
        """

        self.load(data)

        # DEBUG
        #import pprint
        #logger.debug('initialize QueueData from raw:\n%s' % pprint.pformat(data))
        logger.debug('Final parsed QueueData content:\n%s' % self)

    def load(self, data):
        """
            Construct and initialize data from ext source
            :param data: input dictionary of queue data settings
        """

        # the translation map of the queue data attributes from external data to internal schema
        # 'internal_name':('ext_name1', 'extname2_if_any')
        # 'internal_name2':'ext_name3'

        # first defined ext field will be used
        # if key is not explicitly specified then ext name will be used as is
        ## fix me later to proper internal names if need

        kmap = {
            'name': 'nickname',
            'resource': 'panda_resource',
            'platform': 'cmtconfig',
            'site': ('atlas_site', 'gstat'),
            'es_stageout_gap': 'zip_time_gap',
        }

        self._load_data(data, kmap)

    def resolve_allowed_schemas(self, activity, copytool=None):
        """
            Resolve list of allowed schemas for given activity and requested copytool based on `acopytools_schemas` settings
            :param activity: str or ordered list of transfer activity names to resolve acopytools related data
            :return: list of protocol schemes
        """

        if not activity:
            activity = 'default'
        try:
            if isinstance(activity, basestring):  # Python 2
                activity = [activity]
        except Exception:
            if isinstance(activity, str):  # Python 3
                activity = [activity]

        if 'default' not in activity:
            activity = activity + ['default']

        adat = {}
        for aname in activity:
            adat = self.acopytools_schemas.get(aname)
            if adat:
                break
        if not adat:
            return []

        if not isinstance(adat, dict):
            adat = {'default': adat}

        if not copytool or copytool not in adat:
            copytool = 'default'

        return adat.get(copytool) or []

    def clean(self):
        """
            Validate and finally clean up required data values (required object properties) if need
            :return: None
        """

        # validate es_stageout_gap value
        if not self.es_stageout_gap:
            is_opportunistic = self.pledgedcpu and self.pledgedcpu == -1
            self.es_stageout_gap = 600 if is_opportunistic else 7200  ## 10 munites for opportunistic or 5 hours for normal resources

        # validate container_options: extract from the catchall if not set
        if not self.container_options and self.catchall:  ## container_options is considered for the singularity container, FIX ME LATER IF NEED
            # expected format
            # of catchall = "singularity_options=\'-B /etc/grid-security/certificates,/cvmfs,${workdir} --contain\'"
            pattern = re.compile("singularity_options=['\"]?([^'\"]+)['\"]?")  ### FIX ME LATER: move to proper args parsing via shlex at Job class
            found = re.findall(pattern, self.catchall)
            if found:
                self.container_options = found[0]
                logger.info('container_options extracted from catchall: %s' % self.container_options)

        # verify container_options: add the workdir if missing
        if self.container_options:
            if "${workdir}" not in self.container_options and " --contain" in self.container_options:  ## reimplement with shlex later
                self.container_options = self.container_options.replace(" --contain", ",${workdir} --contain")
                logger.info("Note: added missing ${workdir} to container_options/singularity_options: %s" % self.container_options)

        pass

    ## custom function pattern to apply extra validation to the key values
    ##def clean__keyname(self, raw, value):
    ##  :param raw: raw value passed from ext source as input
    ##  :param value: preliminary cleaned and casted to proper type value
    ##
    ##    return value

    def clean__timefloor(self, raw, value):
        """
            Verify and validate value for the timefloor key (convert to seconds)
        """

        return value * 60

    def clean__container_type(self, raw, value):
        """
            Parse and prepare value for the container_type key
            Expected raw data in format 'container_name:user_name;'
            E.g. container_type = 'singularity:pilot;docker:wrapper'

            :return: dict of container names by user as a key
        """

        ret = {}
        val = value or ''
        for e in val.split(';'):
            dat = e.split(':')
            if len(dat) == 2:
                name, user = dat[0].strip(), dat[1].strip()
                ret[user] = name

        return ret

    def clean__container_options(self, raw, value):
        """
            Verify and validate value for the container_options key (remove bad values)
        """

        return value if value.lower() not in ['none'] else ''

    def clean__corecount(self, raw, value):
        """
            Verify and validate value for the corecount key (set to 1 if not set)
        """

        return value if value else 1
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Mario Lassnig, mario.lassnig@cern.ch, 2017
# - Paul Nilsson, paul.nilsson@cern.ch, 2018-2019

from os import environ

# Pilot version
RELEASE = '2'   # released number should be fixed at 2 for Pilot 2
VERSION = '1'   # version number is '1' for first real Pilot 2 release, '0' until then, increased for bigger updates
REVISION = '1'  # revision number should be reset to '0' for every new version release, increased for small updates
BUILD = '8'   # build number should be reset to '1' for every new development cycle

SUCCESS = 0
FAILURE = 1

ERRNO_NOJOBS = 20

# Sorting order constants
UTILITY_BEFORE_PAYLOAD = 0
UTILITY_WITH_PAYLOAD = 1
UTILITY_AFTER_PAYLOAD_STARTED = 2
UTILITY_AFTER_PAYLOAD_FINISHED = 3
UTILITY_WITH_STAGEIN = 4

# Timing constants that allow for additional constants to be defined for values before the pilot is started, ie for
# wrapper timing purposes.
PILOT_START_TIME = 'PILOT_START_TIME'
PILOT_PRE_GETJOB = 'PILOT_PRE_GETJOB'
PILOT_POST_GETJOB = 'PILOT_POST_GETJOB'  # note: PILOT_POST_GETJOB corresponds to START_TIME in Pilot 1
PILOT_PRE_SETUP = 'PILOT_PRE_SETUP'
PILOT_POST_SETUP = 'PILOT_POST_SETUP'
PILOT_PRE_STAGEIN = 'PILOT_PRE_STAGEIN'
PILOT_POST_STAGEIN = 'PILOT_POST_STAGEIN'
PILOT_PRE_PAYLOAD = 'PILOT_PRE_PAYLOAD'
PILOT_POST_PAYLOAD = 'PILOT_POST_PAYLOAD'
PILOT_PRE_STAGEOUT = 'PILOT_PRE_STAGEOUT'
PILOT_POST_STAGEOUT = 'PILOT_POST_STAGEOUT'
PILOT_PRE_FINAL_UPDATE = 'PILOT_PRE_FINAL_UPDATE'
PILOT_POST_FINAL_UPDATE = 'PILOT_POST_FINAL_UPDATE'
PILOT_END_TIME = 'PILOT_END_TIME'
PILOT_KILL_SIGNAL = 'PILOT_KILL_SIGNAL'

# Keep track of log transfers
LOG_TRANSFER_NOT_DONE = 'NOT_DONE'
LOG_TRANSFER_IN_PROGRESS = 'IN_PROGRESS'
LOG_TRANSFER_DONE = 'DONE'
LOG_TRANSFER_FAILED = 'FAILED'

def get_pilot_version():
    """
    Return the current Pilot version string with the format <release>.<version>.<revision> (<build>).
    E.g. pilot_version = '2.1.3 (12)'
    :return: version string.
    """

    return '{release}.{version}.{revision} ({build})'.format(release=RELEASE,
                                                             version=VERSION,
                                                             revision=REVISION,
                                                             build=BUILD)

def get_rucio_client_version():
    """
    Return the current Rucio client version string using the environmental variable ATLAS_LOCAL_RUCIOCLIENTS_VERSION.
    If the environmental variable is not set, then an empty string will be returned.

    :return: $ATLAS_LOCAL_RUCIOCLIENTS_VERSION (string).
    """

    return environ.get('ATLAS_LOCAL_RUCIOCLIENTS_VERSION', '')
<commit_msg>Updated build number<commit_after>#!/usr/bin/env python
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Mario Lassnig, mario.lassnig@cern.ch, 2017
# - Paul Nilsson, paul.nilsson@cern.ch, 2018-2019

from os import environ

# Pilot version
RELEASE = '2'   # released number should be fixed at 2 for Pilot 2
VERSION = '1'   # version number is '1' for first real Pilot 2 release, '0' until then, increased for bigger updates
REVISION = '1'  # revision number should be reset to '0' for every new version release, increased for small updates
BUILD = '9'   # build number should be reset to '1' for every new development cycle

SUCCESS = 0
FAILURE = 1

ERRNO_NOJOBS = 20

# Sorting order constants
UTILITY_BEFORE_PAYLOAD = 0
UTILITY_WITH_PAYLOAD = 1
UTILITY_AFTER_PAYLOAD_STARTED = 2
UTILITY_AFTER_PAYLOAD_FINISHED = 3
UTILITY_WITH_STAGEIN = 4

# Timing constants that allow for additional constants to be defined for values before the pilot is started, ie for
# wrapper timing purposes.
PILOT_START_TIME = 'PILOT_START_TIME'
PILOT_PRE_GETJOB = 'PILOT_PRE_GETJOB'
PILOT_POST_GETJOB = 'PILOT_POST_GETJOB'  # note: PILOT_POST_GETJOB corresponds to START_TIME in Pilot 1
PILOT_PRE_SETUP = 'PILOT_PRE_SETUP'
PILOT_POST_SETUP = 'PILOT_POST_SETUP'
PILOT_PRE_STAGEIN = 'PILOT_PRE_STAGEIN'
PILOT_POST_STAGEIN = 'PILOT_POST_STAGEIN'
PILOT_PRE_PAYLOAD = 'PILOT_PRE_PAYLOAD'
PILOT_POST_PAYLOAD = 'PILOT_POST_PAYLOAD'
PILOT_PRE_STAGEOUT = 'PILOT_PRE_STAGEOUT'
PILOT_POST_STAGEOUT = 'PILOT_POST_STAGEOUT'
PILOT_PRE_FINAL_UPDATE = 'PILOT_PRE_FINAL_UPDATE'
PILOT_POST_FINAL_UPDATE = 'PILOT_POST_FINAL_UPDATE'
PILOT_END_TIME = 'PILOT_END_TIME'
PILOT_KILL_SIGNAL = 'PILOT_KILL_SIGNAL'

# Keep track of log transfers
LOG_TRANSFER_NOT_DONE = 'NOT_DONE'
LOG_TRANSFER_IN_PROGRESS = 'IN_PROGRESS'
LOG_TRANSFER_DONE = 'DONE'
LOG_TRANSFER_FAILED = 'FAILED'

def get_pilot_version():
    """
    Return the current Pilot version string with the format <release>.<version>.<revision> (<build>).
    E.g. pilot_version = '2.1.3 (12)'
    :return: version string.
    """

    return '{release}.{version}.{revision} ({build})'.format(release=RELEASE,
                                                             version=VERSION,
                                                             revision=REVISION,
                                                             build=BUILD)

def get_rucio_client_version():
    """
    Return the current Rucio client version string using the environmental variable ATLAS_LOCAL_RUCIOCLIENTS_VERSION.
    If the environmental variable is not set, then an empty string will be returned.

    :return: $ATLAS_LOCAL_RUCIOCLIENTS_VERSION (string).
    """

    return environ.get('ATLAS_LOCAL_RUCIOCLIENTS_VERSION', '')
<|endoftext|>