<commit_before># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# 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.
#

"""
A luigi file system client that wraps around snakebite

Originally written by Alan Brenner <alan@magnetic.com> github.com/alanbbr
"""

from luigi.contrib.hdfs import config as hdfs_config
from luigi.contrib.hdfs import clients as hdfs_clients
from luigi import six
import luigi.contrib.target
import logging
import datetime
import os

logger = logging.getLogger('luigi-interface')

class SnakebiteHdfsClient(hdfs_clients.HdfsClient):
    """
    This client uses Spotify's snakebite client whenever possible.
    """

    def __init__(self):
        super(SnakebiteHdfsClient, self).__init__()
        try:
            from snakebite.client import Client
            self._bite = None
            self.pid = -1
        except Exception as err:    # IGNORE:broad-except
            raise RuntimeError("You must specify namenode_host and namenode_port "
                               "in the [hdfs] section of your luigi config in "
                               "order to use luigi's snakebite support", err)

    def __new__(cls):
        try:
            from snakebite.client import Client
            this = super(SnakebiteHdfsClient, cls).__new__(cls)
            return this
        except ImportError:
            logger.warning("Failed to load snakebite.client. Using HdfsClient.")
            return hdfs_clients.HdfsClient()

    @staticmethod
    def list_path(path):
        if isinstance(path, list) or isinstance(path, tuple):
            return path
        if isinstance(path, str) or isinstance(path, unicode):
            return [path, ]
        return [str(path), ]

    def get_bite(self):
        """
        If Luigi has forked, we have a different PID, and need to reconnect.
        """
        config = hdfs_config.hdfs()
        if self.pid != os.getpid() or not self._bite:
            client_kwargs = dict(filter(
                lambda k_v: k_v[1] is not None and k_v[1] != '', six.iteritems({
                    'hadoop_version': config.client_version,
                    'effective_user': config.effective_user,
                })
            ))
            if config.snakebite_autoconfig:
                """
                This is fully backwards compatible with the vanilla Client and can be used for a non HA cluster as well.
                This client tries to read ``${HADOOP_PATH}/conf/hdfs-site.xml`` to get the address of the namenode.
                The behaviour is the same as Client.
                """
                from snakebite.client import AutoConfigClient
                self._bite = AutoConfigClient(**client_kwargs)
            else:
                from snakebite.client import Client
                self._bite = Client(config.namenode_host, config.namenode_port, **client_kwargs)
        return self._bite

    def exists(self, path):
        """
        Use snakebite.test to check file existence.

        :param path: path to test
        :type path: string
        :return: boolean, True if path exists in HDFS
        """
        try:
            return self.get_bite().test(path, exists=True)
        except Exception as err:    # IGNORE:broad-except
            raise hdfs_clients.HDFSCliError("snakebite.test", -1, str(err), repr(err))

    def rename(self, path, dest):
        """
        Use snakebite.rename, if available.

        :param path: source file(s)
        :type path: either a string or sequence of strings
        :param dest: destination file (single input) or directory (multiple)
        :type dest: string
        :return: list of renamed items
        """
        parts = dest.rstrip('/').split('/')
        if len(parts) > 1:
            dir_path = '/'.join(parts[0:-1])
            if not self.exists(dir_path):
                self.mkdir(dir_path, parents=True)
        return list(self.get_bite().rename(self.list_path(path), dest))

    def rename_dont_move(self, path, dest):
        """
        Use snakebite.rename_dont_move, if available.

        :param path: source path (single input)
        :type path: string
        :param dest: destination path
        :type dest: string
        :return: True if succeeded
        :raises: snakebite.errors.FileAlreadyExistsException
        """
        from snakebite.errors import FileAlreadyExistsException
        try:
            self.get_bite().rename2(path, dest, overwriteDest=False)
            return True
        except FileAlreadyExistsException:
            return False

    def remove(self, path, recursive=True, skip_trash=False):
        """
        Use snakebite.delete, if available.

        :param path: delete-able file(s) or directory(ies)
        :type path: either a string or a sequence of strings
        :param recursive: delete directories trees like \*nix: rm -r
        :type recursive: boolean, default is True
        :param skip_trash: do or don't move deleted items into the trash first
        :type skip_trash: boolean, default is False (use trash)
        :return: list of deleted items
        """
        return list(self.get_bite().delete(self.list_path(path), recurse=recursive))

    def chmod(self, path, permissions, recursive=False):
        """
        Use snakebite.chmod, if available.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param permissions: \*nix style permission number
        :type permissions: octal
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        if type(permissions) == str:
            permissions = int(permissions, 8)
        return list(self.get_bite().chmod(self.list_path(path),
                                          permissions, recursive))

    def chown(self, path, owner, group, recursive=False):
        """
        Use snakebite.chown/chgrp, if available.

        One of owner or group must be set. Just setting group calls chgrp.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param owner: new owner, can be blank
        :type owner: string
        :param group: new group, can be blank
        :type group: string
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        bite = self.get_bite()
        if owner:
            if group:
                return all(bite.chown(self.list_path(path), "%s:%s" % (owner, group),
                                      recurse=recursive))
            return all(bite.chown(self.list_path(path), owner, recurse=recursive))
        return list(bite.chgrp(self.list_path(path), group, recurse=recursive))

    def count(self, path):
        """
        Use snakebite.count, if available.

        :param path: directory to count the contents of
        :type path: string
        :return: dictionary with content_size, dir_count and file_count keys
        """
        try:
            res = self.get_bite().count(self.list_path(path)).next()
            dir_count = res['directoryCount']
            file_count = res['fileCount']
            content_size = res['spaceConsumed']
        except StopIteration:
            dir_count = file_count = content_size = 0
        return {'content_size': content_size, 'dir_count': dir_count,
                'file_count': file_count}

    def get(self, path, local_destination):
        """
        Use snakebite.copyToLocal, if available.

        :param path: HDFS file
        :type path: string
        :param local_destination: path on the system running Luigi
        :type local_destination: string
        """
        return list(self.get_bite().copyToLocal(self.list_path(path),
                                                local_destination))

    def mkdir(self, path, parents=True, mode=0o755, raise_if_exists=False):
        """
        Use snakebite.mkdir, if available.

        Snakebite's mkdir method allows control over full path creation, so by
        default, tell it to build a full path to work like ``hadoop fs -mkdir``.

        :param path: HDFS path to create
        :type path: string
        :param parents: create any missing parent directories
        :type parents: boolean, default is True
        :param mode: \*nix style owner/group/other permissions
        :type mode: octal, default 0755
        """
        result = list(self.get_bite().mkdir(self.list_path(path),
                                            create_parent=parents, mode=mode))
        if raise_if_exists and "ile exists" in result[0].get('error', ''):
            raise luigi.target.FileAlreadyExists("%s exists" % (path, ))
        return result

    def listdir(self, path, ignore_directories=False, ignore_files=False,
                include_size=False, include_type=False, include_time=False,
                recursive=False):
        """
        Use snakebite.ls to get the list of items in a directory.

        :param path: the directory to list
        :type path: string
        :param ignore_directories: if True, do not yield directory entries
        :type ignore_directories: boolean, default is False
        :param ignore_files: if True, do not yield file entries
        :type ignore_files: boolean, default is False
        :param include_size: include the size in bytes of the current item
        :type include_size: boolean, default is False (do not include)
        :param include_type: include the type (d or f) of the current item
        :type include_type: boolean, default is False (do not include)
        :param include_time: include the last modification time of the current item
        :type include_time: boolean, default is False (do not include)
        :param recursive: list subdirectory contents
        :type recursive: boolean, default is False (do not recurse)
        :return: yield with a string, or if any of the include_* settings are
            true, a tuple starting with the path, and include_* items in order
        """
        bite = self.get_bite()
        for entry in bite.ls(self.list_path(path), recurse=recursive):
            if ignore_directories and entry['file_type'] == 'd':
                continue
            if ignore_files and entry['file_type'] == 'f':
                continue
            rval = [entry['path'], ]
            if include_size:
                rval.append(entry['length'])
            if include_type:
                rval.append(entry['file_type'])
            if include_time:
                rval.append(datetime.datetime.fromtimestamp(entry['modification_time'] / 1000))
            if len(rval) > 1:
                yield tuple(rval)
            else:
                yield rval[0]
<commit_msg>Don't make SnakebiteClient fallback to HdfsClient<commit_after># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# 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.
#

"""
A luigi file system client that wraps around snakebite

Originally written by Alan Brenner <alan@magnetic.com> github.com/alanbbr
"""

from luigi.contrib.hdfs import config as hdfs_config
from luigi.contrib.hdfs import clients as hdfs_clients
from luigi import six
import luigi.contrib.target
import logging
import datetime
import os

logger = logging.getLogger('luigi-interface')

class SnakebiteHdfsClient(hdfs_clients.HdfsClient):
    """
    A hdfs client using snakebite. Since Snakebite has a python API, it'll be
    about 100 times faster than the hadoop cli client, which does shell out to
    a java program on each file system operation.
    """

    def __init__(self):
        super(SnakebiteHdfsClient, self).__init__()
        try:
            from snakebite.client import Client
            self._bite = None
            self.pid = -1
        except Exception as err:    # IGNORE:broad-except
            raise RuntimeError("You must specify namenode_host and namenode_port "
                               "in the [hdfs] section of your luigi config in "
                               "order to use luigi's snakebite support", err)

    @staticmethod
    def list_path(path):
        if isinstance(path, list) or isinstance(path, tuple):
            return path
        if isinstance(path, str) or isinstance(path, unicode):
            return [path, ]
        return [str(path), ]

    def get_bite(self):
        """
        If Luigi has forked, we have a different PID, and need to reconnect.
        """
        config = hdfs_config.hdfs()
        if self.pid != os.getpid() or not self._bite:
            client_kwargs = dict(filter(
                lambda k_v: k_v[1] is not None and k_v[1] != '', six.iteritems({
                    'hadoop_version': config.client_version,
                    'effective_user': config.effective_user,
                })
            ))
            if config.snakebite_autoconfig:
                """
                This is fully backwards compatible with the vanilla Client and can be used for a non HA cluster as well.
                This client tries to read ``${HADOOP_PATH}/conf/hdfs-site.xml`` to get the address of the namenode.
                The behaviour is the same as Client.
                """
                from snakebite.client import AutoConfigClient
                self._bite = AutoConfigClient(**client_kwargs)
            else:
                from snakebite.client import Client
                self._bite = Client(config.namenode_host, config.namenode_port, **client_kwargs)
        return self._bite

    def exists(self, path):
        """
        Use snakebite.test to check file existence.

        :param path: path to test
        :type path: string
        :return: boolean, True if path exists in HDFS
        """
        try:
            return self.get_bite().test(path, exists=True)
        except Exception as err:    # IGNORE:broad-except
            raise hdfs_clients.HDFSCliError("snakebite.test", -1, str(err), repr(err))

    def rename(self, path, dest):
        """
        Use snakebite.rename, if available.

        :param path: source file(s)
        :type path: either a string or sequence of strings
        :param dest: destination file (single input) or directory (multiple)
        :type dest: string
        :return: list of renamed items
        """
        parts = dest.rstrip('/').split('/')
        if len(parts) > 1:
            dir_path = '/'.join(parts[0:-1])
            if not self.exists(dir_path):
                self.mkdir(dir_path, parents=True)
        return list(self.get_bite().rename(self.list_path(path), dest))

    def rename_dont_move(self, path, dest):
        """
        Use snakebite.rename_dont_move, if available.

        :param path: source path (single input)
        :type path: string
        :param dest: destination path
        :type dest: string
        :return: True if succeeded
        :raises: snakebite.errors.FileAlreadyExistsException
        """
        from snakebite.errors import FileAlreadyExistsException
        try:
            self.get_bite().rename2(path, dest, overwriteDest=False)
            return True
        except FileAlreadyExistsException:
            return False

    def remove(self, path, recursive=True, skip_trash=False):
        """
        Use snakebite.delete, if available.

        :param path: delete-able file(s) or directory(ies)
        :type path: either a string or a sequence of strings
        :param recursive: delete directories trees like \*nix: rm -r
        :type recursive: boolean, default is True
        :param skip_trash: do or don't move deleted items into the trash first
        :type skip_trash: boolean, default is False (use trash)
        :return: list of deleted items
        """
        return list(self.get_bite().delete(self.list_path(path), recurse=recursive))

    def chmod(self, path, permissions, recursive=False):
        """
        Use snakebite.chmod, if available.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param permissions: \*nix style permission number
        :type permissions: octal
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        if type(permissions) == str:
            permissions = int(permissions, 8)
        return list(self.get_bite().chmod(self.list_path(path),
                                          permissions, recursive))

    def chown(self, path, owner, group, recursive=False):
        """
        Use snakebite.chown/chgrp, if available.

        One of owner or group must be set. Just setting group calls chgrp.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param owner: new owner, can be blank
        :type owner: string
        :param group: new group, can be blank
        :type group: string
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        bite = self.get_bite()
        if owner:
            if group:
                return all(bite.chown(self.list_path(path), "%s:%s" % (owner, group),
                                      recurse=recursive))
            return all(bite.chown(self.list_path(path), owner, recurse=recursive))
        return list(bite.chgrp(self.list_path(path), group, recurse=recursive))

    def count(self, path):
        """
        Use snakebite.count, if available.

        :param path: directory to count the contents of
        :type path: string
        :return: dictionary with content_size, dir_count and file_count keys
        """
        try:
            res = self.get_bite().count(self.list_path(path)).next()
            dir_count = res['directoryCount']
            file_count = res['fileCount']
            content_size = res['spaceConsumed']
        except StopIteration:
            dir_count = file_count = content_size = 0
        return {'content_size': content_size, 'dir_count': dir_count,
                'file_count': file_count}

    def get(self, path, local_destination):
        """
        Use snakebite.copyToLocal, if available.

        :param path: HDFS file
        :type path: string
        :param local_destination: path on the system running Luigi
        :type local_destination: string
        """
        return list(self.get_bite().copyToLocal(self.list_path(path),
                                                local_destination))

    def mkdir(self, path, parents=True, mode=0o755, raise_if_exists=False):
        """
        Use snakebite.mkdir, if available.

        Snakebite's mkdir method allows control over full path creation, so by
        default, tell it to build a full path to work like ``hadoop fs -mkdir``.

        :param path: HDFS path to create
        :type path: string
        :param parents: create any missing parent directories
        :type parents: boolean, default is True
        :param mode: \*nix style owner/group/other permissions
        :type mode: octal, default 0755
        """
        result = list(self.get_bite().mkdir(self.list_path(path),
                                            create_parent=parents, mode=mode))
        if raise_if_exists and "ile exists" in result[0].get('error', ''):
            raise luigi.target.FileAlreadyExists("%s exists" % (path, ))
        return result

    def listdir(self, path, ignore_directories=False, ignore_files=False,
                include_size=False, include_type=False, include_time=False,
                recursive=False):
        """
        Use snakebite.ls to get the list of items in a directory.

        :param path: the directory to list
        :type path: string
        :param ignore_directories: if True, do not yield directory entries
        :type ignore_directories: boolean, default is False
        :param ignore_files: if True, do not yield file entries
        :type ignore_files: boolean, default is False
        :param include_size: include the size in bytes of the current item
        :type include_size: boolean, default is False (do not include)
        :param include_type: include the type (d or f) of the current item
        :type include_type: boolean, default is False (do not include)
        :param include_time: include the last modification time of the current item
        :type include_time: boolean, default is False (do not include)
        :param recursive: list subdirectory contents
        :type recursive: boolean, default is False (do not recurse)
        :return: yield with a string, or if any of the include_* settings are
            true, a tuple starting with the path, and include_* items in order
        """
        bite = self.get_bite()
        for entry in bite.ls(self.list_path(path), recurse=recursive):
            if ignore_directories and entry['file_type'] == 'd':
                continue
            if ignore_files and entry['file_type'] == 'f':
                continue
            rval = [entry['path'], ]
            if include_size:
                rval.append(entry['length'])
            if include_type:
                rval.append(entry['file_type'])
            if include_time:
                rval.append(datetime.datetime.fromtimestamp(entry['modification_time'] / 1000))
            if len(rval) > 1:
                yield tuple(rval)
            else:
                yield rval[0]
<|endoftext|>


<commit_before>
from collections import OrderedDict
from mangrove.form_model.form_model import GEO_CODE_FIELD_NAME, LOCATION_TYPE_FIELD_NAME, GEO_CODE, ENTITY_TYPE_FIELD_CODE, REGISTRATION_FORM_CODE
from mangrove.form_model.location import Location
from mangrove.datastore import entity
from mangrove.utils.types import is_empty, is_not_empty

class FormSubmission(object):
    def __init__(self, form_model, form_answers, errors=None, location_tree=None):
        assert errors is None or type(errors) == OrderedDict
        assert form_answers is not None and type(form_answers) == OrderedDict
        assert form_model is not None

        self.form_model = form_model
        self._cleaned_data = form_answers
        entity_short_code = self.get_answer_for(form_model.entity_question.code)
        self.short_code = entity_short_code.lower() if entity_short_code is not None else None
        self.entity_type = self.get_entity_type(form_model)
        self.is_valid = (errors is None or len(errors) == 0)
        self.errors = errors
        self.data_record_id = None
        self.location_tree = location_tree

    @property
    def cleaned_data(self):
        return self._cleaned_data

    @property
    def form_code(self):
        return self.form_model.form_code

    @property
    def saved(self):
        return self.data_record_id is not None

    @property
    def is_registration(self):
        return self.form_model.is_entity_registration_form()

    def save_new(self, dbm):
        entity = self.create_entity(dbm)
        return self._save_data(entity)

    def save(self, dbm):
        return self.save_new(dbm)

    def update(self, dbm):
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)
        entity = self.get_entity(dbm)
        entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
        return self._save_data(entity)

    def _contains_geo_code(self, item):
        item_ = item[0]
        return item_ == GEO_CODE_FIELD_NAME

    def _get_event_time_value(self):
        return self.cleaned_data.get(self._get_event_time_code())

    def _get_event_time_code(self):
        event_time_field = self.form_model.event_time_question
        return event_time_field.code if event_time_field else None

    def _save_data(self, entity):
        submission_information = dict(form_code=self.form_code)
        self.data_record_id = entity.add_data(data=self._values, event_time=self._get_event_time_value(),
                                              submission=submission_information)
        return self.data_record_id

    def _to_three_tuple(self):
        return [(self.form_model._get_field_by_code(code).name, value, self.form_model._get_field_by_code(code).ddtype)
                for (code, value) in
                (self.cleaned_data.items())]

    def get_answer_for(self, code):
        for key in self._cleaned_data:
            if key.lower() == code.lower():
                return self._cleaned_data[key]
        return None

    @property
    def _values(self):
        return self._to_three_tuple()

    def create_entity(self, dbm):
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)

        return entity.create_entity(dbm=dbm, entity_type=self.entity_type,
                                    location=location_hierarchy,
                                    short_code=self.short_code,
                                    geometry=processed_geometry)

    def get_entity(self, dbm):
        return entity.get_by_short_code(dbm=dbm, short_code=self.short_code, entity_type=self.entity_type)

    def _get_field_code_by_name(self, field_name):
        field = self.form_model.get_field_by_name(name=field_name)
        return field.code if field is not None else None

    def get_location_field_code(self):
        return self._get_field_code_by_name(LOCATION_TYPE_FIELD_NAME)

    def get_geo_field_code(self):
        return self._get_field_code_by_name(GEO_CODE_FIELD_NAME)

    def get_entity_type(self, form_model):
        entity_type = self.form_model.entity_type
        return [e_type.lower() for e_type in entity_type] if is_not_empty(entity_type) else None

class DataFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors):
        super(DataFormSubmission, self).__init__(form_model, answers, errors)

    def create_entity(self, dbm):
        return entity.get_by_short_code(dbm, self.short_code, self.entity_type)

class GlobalRegistrationFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors, location_tree=None):
        super(GlobalRegistrationFormSubmission, self).__init__(form_model, answers, errors, location_tree=location_tree)

    def update_location_and_geo_code(self, dbm):
        existing_entity = self.get_entity(dbm)
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)
        existing_entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
        existing_entity.save()
        values = self._values
        if is_empty(filter(self._contains_geo_code, values)):
            self._cleaned_data[GEO_CODE] = processed_geometry['coordinates'] if processed_geometry is not None else None

    def get_entity_type(self, form_model):
        entity_type = self.get_answer_for(ENTITY_TYPE_FIELD_CODE)
        return [e_type.lower() for e_type in entity_type] if is_not_empty(entity_type) else None

    def void_existing_data_records(self, dbm, form_code=None):
        data_records = dbm.view.data_record_by_form_code(key=[REGISTRATION_FORM_CODE, self.short_code])
        for data_record in data_records:
            data_record_doc = data_record.value
            data_record_doc['void'] = True
            dbm.database.save(data_record_doc)

class EntityRegistrationFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors, location_tree=None):
        super(EntityRegistrationFormSubmission, self).__init__(form_model, answers, errors, location_tree=location_tree)

    # # get method is doing an update!!!!
    # def get_entity(self, dbm):
    #     location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
    #         self.cleaned_data)
    #     existing_entity = entity.get_by_short_code(dbm=dbm, short_code=self.short_code, entity_type=self.entity_type)
    #     existing_entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
    #     existing_entity.save()
    #     return existing_entity

    # soft deletes data records
    def void_existing_data_records(self, dbm, form_code):
        data_records = dbm.view.data_record_by_form_code(key=[form_code, self.short_code])
        for data_record in data_records:
            data_record_doc = data_record.value
            data_record_doc['void'] = True
            dbm.database.save(data_record_doc)

class FormSubmissionFactory(object):
    def get_form_submission(self, form_model, answers, errors=None, location_tree=None):
        if not form_model.is_entity_registration_form():
            return DataFormSubmission(form_model, answers, errors)
        elif form_model.is_global_registration_form():
            return GlobalRegistrationFormSubmission(form_model, answers, errors, location_tree=location_tree)
        else:
            return EntityRegistrationFormSubmission(form_model, answers, errors, location_tree=location_tree)<commit_msg>Ajay/Melanie | #000 | Removed unwanted update of Entity document with project questionnaire<commit_after>
from collections import OrderedDict
from mangrove.form_model.form_model import GEO_CODE_FIELD_NAME, LOCATION_TYPE_FIELD_NAME, GEO_CODE, ENTITY_TYPE_FIELD_CODE, REGISTRATION_FORM_CODE
from mangrove.form_model.location import Location
from mangrove.datastore import entity
from mangrove.utils.types import is_empty, is_not_empty

class FormSubmission(object):
    def __init__(self, form_model, form_answers, errors=None, location_tree=None):
        assert errors is None or type(errors) == OrderedDict
        assert form_answers is not None and type(form_answers) == OrderedDict
        assert form_model is not None

        self.form_model = form_model
        self._cleaned_data = form_answers
        entity_short_code = self.get_answer_for(form_model.entity_question.code)
        self.short_code = entity_short_code.lower() if entity_short_code is not None else None
        self.entity_type = self.get_entity_type(form_model)
        self.is_valid = (errors is None or len(errors) == 0)
        self.errors = errors
        self.data_record_id = None
        self.location_tree = location_tree

    @property
    def cleaned_data(self):
        return self._cleaned_data

    @property
    def form_code(self):
        return self.form_model.form_code

    @property
    def saved(self):
        return self.data_record_id is not None

    @property
    def is_registration(self):
        return self.form_model.is_entity_registration_form()

    def save_new(self, dbm):
        entity = self.create_entity(dbm)
        self.data_record_id = entity.id
        return entity.id

    def save(self, dbm):
        return self.save_new(dbm)

    def update(self, dbm):
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)
        entity = self.get_entity(dbm)
        entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
        return self._save_data(entity)

    def _contains_geo_code(self, item):
        item_ = item[0]
        return item_ == GEO_CODE_FIELD_NAME

    def _get_event_time_value(self):
        return self.cleaned_data.get(self._get_event_time_code())

    def _get_event_time_code(self):
        event_time_field = self.form_model.event_time_question
        return event_time_field.code if event_time_field else None

    def _save_data(self, entity):
        submission_information = dict(form_code=self.form_code)
        self.data_record_id = entity.add_data(data=self._values, event_time=self._get_event_time_value(),
                                              submission=submission_information)
        return self.data_record_id

    def _to_three_tuple(self):
        return [(self.form_model._get_field_by_code(code).name, value, self.form_model._get_field_by_code(code).ddtype)
                for (code, value) in
                (self.cleaned_data.items())]

    def get_answer_for(self, code):
        for key in self._cleaned_data:
            if key.lower() == code.lower():
                return self._cleaned_data[key]
        return None

    @property
    def _values(self):
        return self._to_three_tuple()

    def create_entity(self, dbm):
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)

        return entity.create_entity(dbm=dbm, entity_type=self.entity_type,
                                    location=location_hierarchy,
                                    short_code=self.short_code,
                                    geometry=processed_geometry)

    def get_entity(self, dbm):
        return entity.get_by_short_code(dbm=dbm, short_code=self.short_code, entity_type=self.entity_type)

    def _get_field_code_by_name(self, field_name):
        field = self.form_model.get_field_by_name(name=field_name)
        return field.code if field is not None else None

    def get_location_field_code(self):
        return self._get_field_code_by_name(LOCATION_TYPE_FIELD_NAME)

    def get_geo_field_code(self):
        return self._get_field_code_by_name(GEO_CODE_FIELD_NAME)

    def get_entity_type(self, form_model):
        entity_type = self.form_model.entity_type
        return [e_type.lower() for e_type in entity_type] if is_not_empty(entity_type) else None

class DataFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors):
        super(DataFormSubmission, self).__init__(form_model, answers, errors)

    def create_entity(self, dbm):
        return entity.get_by_short_code(dbm, self.short_code, self.entity_type)

class GlobalRegistrationFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors, location_tree=None):
        super(GlobalRegistrationFormSubmission, self).__init__(form_model, answers, errors, location_tree=location_tree)

    def update_location_and_geo_code(self, dbm):
        existing_entity = self.get_entity(dbm)
        location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
            self.cleaned_data)
        existing_entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
        existing_entity.save()
        values = self._values
        if is_empty(filter(self._contains_geo_code, values)):
            self._cleaned_data[GEO_CODE] = processed_geometry['coordinates'] if processed_geometry is not None else None

    def get_entity_type(self, form_model):
        entity_type = self.get_answer_for(ENTITY_TYPE_FIELD_CODE)
        return [e_type.lower() for e_type in entity_type] if is_not_empty(entity_type) else None

    def void_existing_data_records(self, dbm, form_code=None):
        data_records = dbm.view.data_record_by_form_code(key=[REGISTRATION_FORM_CODE, self.short_code])
        for data_record in data_records:
            data_record_doc = data_record.value
            data_record_doc['void'] = True
            dbm.database.save(data_record_doc)

class EntityRegistrationFormSubmission(FormSubmission):
    def __init__(self, form_model, answers, errors, location_tree=None):
        super(EntityRegistrationFormSubmission, self).__init__(form_model, answers, errors, location_tree=location_tree)

    # # get method is doing an update!!!!
    # def get_entity(self, dbm):
    #     location_hierarchy, processed_geometry = Location(self.location_tree, self.form_model).process_entity_creation(
    #         self.cleaned_data)
    #     existing_entity = entity.get_by_short_code(dbm=dbm, short_code=self.short_code, entity_type=self.entity_type)
    #     existing_entity.set_location_and_geo_code(location_hierarchy, processed_geometry)
    #     existing_entity.save()
    #     return existing_entity

    # soft deletes data records
    def void_existing_data_records(self, dbm, form_code):
        data_records = dbm.view.data_record_by_form_code(key=[form_code, self.short_code])
        for data_record in data_records:
            data_record_doc = data_record.value
            data_record_doc['void'] = True
            dbm.database.save(data_record_doc)

class FormSubmissionFactory(object):
    def get_form_submission(self, form_model, answers, errors=None, location_tree=None):
        if not form_model.is_entity_registration_form():
            return DataFormSubmission(form_model, answers, errors)
        elif form_model.is_global_registration_form():
            return GlobalRegistrationFormSubmission(form_model, answers, errors, location_tree=location_tree)
        else:
            return EntityRegistrationFormSubmission(form_model, answers, errors, location_tree=location_tree)<|endoftext|>


<commit_before># Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Set of utilities to add commands to a buildbot factory.

This is based on commands.py and adds skia-specific commands."""

from buildbot.process.properties import WithProperties
from buildbot.status.logfile import STDOUT
from master.factory import commands
from master.log_parser import retcode_command

import re

class SkiaBuildStep(retcode_command.ReturnCodeCommand):
  """ BuildStep wrapper for Skia. Allows us to define properties of BuildSteps
  to be used by ShouldDoStep. This is necessary because the properties referred
  to by BuildStep.getProperty() are scoped for the entire duration of the build.
  """
  def __init__(self, is_upload_step=False, is_rebaseline_step=False,
               get_props_from_stdout=None, **kwargs):
    """ Instantiates a new SkiaBuildStep.

    is_upload_step: boolean indicating whether this step should be skipped when
        the buildbot is not performing uploads.
    is_rebaseline_step: boolean indicating whether this step is required for
        rebaseline-only builds.
    get_props_from_stdout: optional dictionary. Keys are strings indicating
        build properties to set based on the output of this step. Values are
        strings containing regular expressions for parsing the property from
        the output of the step.
    """
    self._is_upload_step = is_upload_step
    self._is_rebaseline_step = is_rebaseline_step
    self._get_props_from_stdout = get_props_from_stdout

    # self._changed_props will be a dictionary containing the build properties
    # which were updated by this BuildStep. Those properties will be displayed
    # in the label for this step.
    self._changed_props = None

    return retcode_command.ReturnCodeCommand.__init__(self, **kwargs)

  def IsUploadStep(self):
    return self._is_upload_step

  def IsRebaselineStep(self):
    return self._is_rebaseline_step

  def commandComplete(self, cmd):
    """ Override of BuildStep's commandComplete method which allows us to parse
    build properties from the output of this step. """
    if self._get_props_from_stdout and cmd.rc == 0:
      log = cmd.logs['stdio']
      stdout = ''.join(log.getChunks([STDOUT], onlyText=True))
      self._changed_props = {}
      for property, regex in self._get_props_from_stdout.iteritems():
        matches = re.search(regex, stdout)
        if not matches:
          raise Exception('Unable to parse %s from stdout.' % property)
        groups = matches.groups()
        if len(groups) != 1:
          raise Exception('Multiple matches for "%s"' % regex)
        prop_value = groups[0]
        self.setProperty(property, prop_value, ''.join(self.description))
        self._changed_props[property] = prop_value
    retcode_command.ReturnCodeCommand.commandComplete(self, cmd)

  def getText(self, cmd, results):
    """ Override of BuildStep's getText method which appends any changed build
    properties to the description of the BuildStep. """
    text = self.description
    if self._changed_props:
      text.extend(['%s: %s' % (
          key, self._changed_props.get(key)) for key in self._changed_props])
    return text

def _HasProperty(step, property):
  """ Helper used by ShouldDoStep. Determine whether the given BuildStep has
  the requested property.

  step: an instance of BuildStep
  property: string, the property to test
  """
  try:
    step.getProperty(property)
    return True
  except:
    return False

def ShouldDoStep(step):
  """ At build time, use build properties to determine whether or not a step
  should be run or skipped.

  step: an instance of BuildStep which we may or may not run.
  """
  #return False
  print step.build.getProperties()
  if not isinstance(step, SkiaBuildStep):
    return True

  # If this step uploads results (and thus overwrites the most recently uploaded
  # results), only run it on scheduled builds (i.e. most recent revision) or if
  # the "force_upload" property was set.
  if step.IsUploadStep() and \
      not _HasProperty(step, 'scheduler') and \
      not _HasProperty(step, 'force_upload'):
    return False

  # When a commit consists of only new baseline images, we only need to run the
  # BuildSteps necessary for image verification, and only for the platform(s)
  # for which new baselines are provided.
  if  _HasProperty(step, 'branch') and \
      step.getProperty('branch') == 'gm-expected':
    if step.IsRebaselineStep():
      # This step is required for rebaselines, but do the associated commits
      # affect our platform?
      if _HasProperty(step, 'gm_image_subdir'):
        for change in step.build.allChanges():
          for file in change.asDict()['files']:
            if step.getProperty('gm_image_subdir') in file['name']:
              return True
    return False

  # Unless we have determined otherwise, run the step.
  return True

class SkiaCommands(commands.FactoryCommands):

  def __init__(self, factory, configuration, workdir, target_arch,
               default_timeout, target_platform, environment_variables):
    """Instantiates subclass of FactoryCommands appropriate for Skia.

    factory: a BaseFactory
    configuration: 'Debug' or 'Release'
    workdir: string indicating path within slave directory
    target_arch: string such as 'x64'
    default_timeout: default timeout for each command, in seconds
    target_platform: a string such as skia_factory.TARGET_PLATFORM_LINUX,
        to be passed into FactoryCommands.__init__()
    environment_variables: dictionary of environment variables that should
        be passed to all commands
    """
    commands.FactoryCommands.__init__(
        self, factory=factory, target=configuration,
        build_dir='', target_platform=target_platform)
    # Store some parameters that the subclass may want to use later.
    self.default_timeout = default_timeout
    self.environment_variables = environment_variables
    self.factory = factory
    self.target_arch = target_arch
    self.workdir = workdir
    # TODO(epoger): It would be better for this path to be specified by
    # an environment variable or some such, so that it is not dependent on the
    # path from CWD to slave/skia_slave_scripts... but for now, this will do.
    self._local_slave_script_dir = self.PathJoin(
        '..', '..', '..', '..', '..', '..', 'slave', 'skia_slave_scripts')

  def AddMergeIntoSvn(self, source_dir_path, dest_svn_url, merge_dir_path,
                      svn_username_file, svn_password_file,
                      commit_message=None, description='MergeIntoSvn',
                      timeout=None, is_rebaseline_step=False):
    """Adds a step that commits all files within a directory to a special SVN
    repository."""
    if not commit_message:
      commit_message = 'automated svn commit of %s step' % description
    args = ['--commit_message', commit_message,
            '--dest_svn_url', dest_svn_url,
            '--merge_dir_path', merge_dir_path,
            '--source_dir_path', source_dir_path,
            '--svn_password_file', svn_password_file,
            '--svn_username_file', svn_username_file,
            ]
    self.AddSlaveScript(script=self.PathJoin('utils', 'merge_into_svn.py'),
                        args=args, description=description, timeout=timeout,
                        is_upload_step=True,
                        is_rebaseline_step=is_rebaseline_step)

  def AddSlaveScript(self, script, args, description, timeout=None,
                     halt_on_failure=False, is_upload_step=False,
                     is_rebaseline_step=False, get_props_from_stdout=None,
                     workdir=None):
    """Run a slave-side Python script as its own build step."""
    if workdir:
      path_to_script = script
      use_workdir = workdir
    else:
      path_to_script = self.PathJoin(self._local_slave_script_dir, script)
      use_workdir = self.workdir
    self.AddRunCommand(command=['python', path_to_script] + args,
                       description=description, timeout=timeout,
                       halt_on_failure=halt_on_failure,
                       is_upload_step=is_upload_step,
                       is_rebaseline_step=is_rebaseline_step,
                       get_props_from_stdout=get_props_from_stdout,
                       workdir=use_workdir)

  def AddRunCommand(self, command, description='Run', timeout=None,
                    halt_on_failure=False, is_upload_step=False,
                    is_rebaseline_step=False, get_props_from_stdout=None,
                    workdir=None):
    """Runs an arbitrary command, perhaps a binary we built."""
    if not timeout:
      timeout = self.default_timeout
    self.factory.addStep(SkiaBuildStep,
                         is_upload_step=is_upload_step,
                         is_rebaseline_step=is_rebaseline_step,
                         get_props_from_stdout=get_props_from_stdout,
                         description=description, timeout=timeout,
                         command=command, workdir=workdir,
                         env=self.environment_variables,
                         haltOnFailure=halt_on_failure,
                         doStepIf=ShouldDoStep)

  def AddRunCommandList(self, command_list, description='Run', timeout=None,
                        halt_on_failure=False, is_upload_step=False,
                        is_rebaseline_step=False):
    """Runs a list of arbitrary commands."""
    # TODO(epoger): Change this so that build-step output shows each command
    # in the list separately--that will be a lot easier to follow.
    #
    # TODO(epoger): For now, this wraps the total command with WithProperties()
    # because *some* callers need it, and we can't use the string.join() command
    # to concatenate strings that have already been wrapped with
    # WithProperties().  Once I figure out how to make the build-step output
    # show each command separately, maybe I can remove this wrapper.
    self.AddRunCommand(command=WithProperties(' && '.join(command_list)),
                       description=description, timeout=timeout,
                       halt_on_failure=halt_on_failure,
                       is_upload_step=is_upload_step,
                       is_rebaseline_step=is_rebaseline_step)
<commit_msg>Fix housekeeping bot<commit_after># Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Set of utilities to add commands to a buildbot factory.

This is based on commands.py and adds skia-specific commands."""

from buildbot.process.properties import WithProperties
from buildbot.status.logfile import STDOUT
from master.factory import commands
from master.log_parser import retcode_command

import re

class SkiaBuildStep(retcode_command.ReturnCodeCommand):
  """ BuildStep wrapper for Skia. Allows us to define properties of BuildSteps
  to be used by ShouldDoStep. This is necessary because the properties referred
  to by BuildStep.getProperty() are scoped for the entire duration of the build.
  """
  def __init__(self, is_upload_step=False, is_rebaseline_step=False,
               get_props_from_stdout=None, **kwargs):
    """ Instantiates a new SkiaBuildStep.

    is_upload_step: boolean indicating whether this step should be skipped when
        the buildbot is not performing uploads.
    is_rebaseline_step: boolean indicating whether this step is required for
        rebaseline-only builds.
    get_props_from_stdout: optional dictionary. Keys are strings indicating
        build properties to set based on the output of this step. Values are
        strings containing regular expressions for parsing the property from
        the output of the step.
    """
    self._is_upload_step = is_upload_step
    self._is_rebaseline_step = is_rebaseline_step
    self._get_props_from_stdout = get_props_from_stdout

    # self._changed_props will be a dictionary containing the build properties
    # which were updated by this BuildStep. Those properties will be displayed
    # in the label for this step.
    self._changed_props = None

    return retcode_command.ReturnCodeCommand.__init__(self, **kwargs)

  def IsUploadStep(self):
    return self._is_upload_step

  def IsRebaselineStep(self):
    return self._is_rebaseline_step

  def commandComplete(self, cmd):
    """ Override of BuildStep's commandComplete method which allows us to parse
    build properties from the output of this step. """
    if self._get_props_from_stdout and cmd.rc == 0:
      log = cmd.logs['stdio']
      stdout = ''.join(log.getChunks([STDOUT], onlyText=True))
      self._changed_props = {}
      for property, regex in self._get_props_from_stdout.iteritems():
        matches = re.search(regex, stdout)
        if not matches:
          raise Exception('Unable to parse %s from stdout.' % property)
        groups = matches.groups()
        if len(groups) != 1:
          raise Exception('Multiple matches for "%s"' % regex)
        prop_value = groups[0]
        self.setProperty(property, prop_value, ''.join(self.description))
        self._changed_props[property] = prop_value
    retcode_command.ReturnCodeCommand.commandComplete(self, cmd)

  def getText(self, cmd, results):
    """ Override of BuildStep's getText method which appends any changed build
    properties to the description of the BuildStep. """
    text = self.description
    if self._changed_props:
      text.extend(['%s: %s' % (
          key, self._changed_props.get(key)) for key in self._changed_props])
    return text

def _HasProperty(step, property):
  """ Helper used by ShouldDoStep. Determine whether the given BuildStep has
  the requested property.

  step: an instance of BuildStep
  property: string, the property to test
  """
  try:
    step.getProperty(property)
    return True
  except:
    return False

def ShouldDoStep(step):
  """ At build time, use build properties to determine whether or not a step
  should be run or skipped.

  step: an instance of BuildStep which we may or may not run.
  """
  print step.build.getProperties()
  if not isinstance(step, SkiaBuildStep):
    return True

  # If this step uploads results (and thus overwrites the most recently uploaded
  # results), only run it on scheduled builds (i.e. most recent revision) or if
  # the "force_upload" property was set.
  if step.IsUploadStep() and \
      not _HasProperty(step, 'scheduler') and \
      not _HasProperty(step, 'force_upload'):
    return False

  # When a commit consists of only new baseline images, we only need to run the
  # BuildSteps necessary for image verification, and only for the platform(s)
  # for which new baselines are provided.
  if  _HasProperty(step, 'branch') and \
      step.getProperty('branch') == 'gm-expected':
    if step.IsRebaselineStep():
      # This step is required for rebaselines, but do the associated commits
      # affect our platform?
      if _HasProperty(step, 'gm_image_subdir'):
        for change in step.build.allChanges():
          for file in change.asDict()['files']:
            if step.getProperty('gm_image_subdir') in file['name']:
              return True
    return False

  # Unless we have determined otherwise, run the step.
  return True

class SkiaCommands(commands.FactoryCommands):

  def __init__(self, factory, configuration, workdir, target_arch,
               default_timeout, target_platform, environment_variables):
    """Instantiates subclass of FactoryCommands appropriate for Skia.

    factory: a BaseFactory
    configuration: 'Debug' or 'Release'
    workdir: string indicating path within slave directory
    target_arch: string such as 'x64'
    default_timeout: default timeout for each command, in seconds
    target_platform: a string such as skia_factory.TARGET_PLATFORM_LINUX,
        to be passed into FactoryCommands.__init__()
    environment_variables: dictionary of environment variables that should
        be passed to all commands
    """
    commands.FactoryCommands.__init__(
        self, factory=factory, target=configuration,
        build_dir='', target_platform=target_platform)
    # Store some parameters that the subclass may want to use later.
    self.default_timeout = default_timeout
    self.environment_variables = environment_variables
    self.factory = factory
    self.target_arch = target_arch
    self.workdir = workdir
    # TODO(epoger): It would be better for this path to be specified by
    # an environment variable or some such, so that it is not dependent on the
    # path from CWD to slave/skia_slave_scripts... but for now, this will do.
    self._local_slave_script_dir = self.PathJoin(
        '..', '..', '..', '..', '..', '..', 'slave', 'skia_slave_scripts')

  def AddMergeIntoSvn(self, source_dir_path, dest_svn_url, merge_dir_path,
                      svn_username_file, svn_password_file,
                      commit_message=None, description='MergeIntoSvn',
                      timeout=None, is_rebaseline_step=False):
    """Adds a step that commits all files within a directory to a special SVN
    repository."""
    if not commit_message:
      commit_message = 'automated svn commit of %s step' % description
    args = ['--commit_message', commit_message,
            '--dest_svn_url', dest_svn_url,
            '--merge_dir_path', merge_dir_path,
            '--source_dir_path', source_dir_path,
            '--svn_password_file', svn_password_file,
            '--svn_username_file', svn_username_file,
            ]
    self.AddSlaveScript(script=self.PathJoin('utils', 'merge_into_svn.py'),
                        args=args, description=description, timeout=timeout,
                        is_upload_step=True,
                        is_rebaseline_step=is_rebaseline_step)

  def AddSlaveScript(self, script, args, description, timeout=None,
                     halt_on_failure=False, is_upload_step=False,
                     is_rebaseline_step=False, get_props_from_stdout=None,
                     workdir=None):
    """Run a slave-side Python script as its own build step."""
    if workdir:
      path_to_script = script
      use_workdir = workdir
    else:
      path_to_script = self.PathJoin(self._local_slave_script_dir, script)
      use_workdir = self.workdir
    self.AddRunCommand(command=['python', path_to_script] + args,
                       description=description, timeout=timeout,
                       halt_on_failure=halt_on_failure,
                       is_upload_step=is_upload_step,
                       is_rebaseline_step=is_rebaseline_step,
                       get_props_from_stdout=get_props_from_stdout,
                       workdir=use_workdir)

  def AddRunCommand(self, command, description='Run', timeout=None,
                    halt_on_failure=False, is_upload_step=False,
                    is_rebaseline_step=False, get_props_from_stdout=None,
                    workdir=None):
    """Runs an arbitrary command, perhaps a binary we built."""
    if not timeout:
      timeout = self.default_timeout
    self.factory.addStep(SkiaBuildStep,
                         is_upload_step=is_upload_step,
                         is_rebaseline_step=is_rebaseline_step,
                         get_props_from_stdout=get_props_from_stdout,
                         description=description, timeout=timeout,
                         command=command, workdir=workdir or self.workdir,
                         env=self.environment_variables,
                         haltOnFailure=halt_on_failure,
                         doStepIf=ShouldDoStep)

  def AddRunCommandList(self, command_list, description='Run', timeout=None,
                        halt_on_failure=False, is_upload_step=False,
                        is_rebaseline_step=False):
    """Runs a list of arbitrary commands."""
    # TODO(epoger): Change this so that build-step output shows each command
    # in the list separately--that will be a lot easier to follow.
    #
    # TODO(epoger): For now, this wraps the total command with WithProperties()
    # because *some* callers need it, and we can't use the string.join() command
    # to concatenate strings that have already been wrapped with
    # WithProperties().  Once I figure out how to make the build-step output
    # show each command separately, maybe I can remove this wrapper.
    self.AddRunCommand(command=WithProperties(' && '.join(command_list)),
                       description=description, timeout=timeout,
                       halt_on_failure=halt_on_failure,
                       is_upload_step=is_upload_step,
                       is_rebaseline_step=is_rebaseline_step)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

u"""
Bot to enhance French healthcare professionals exchanges on Twitter.

Author: Jérome Pinguet.
License: GNU General Public License version 3 (see 'LICENSE' for details.)

DocTocTocBot started as a modification of the single file [retweet-bot](https://github.com/basti2342/retweet-bot)
by [Basti](https://github.com/basti2342) License: Mozilla Public License version 2.0
  * timelineIterator
  * savepoint
  * RT for loop with error & RT counters
"""

import hashlib
import logging
import os
import tweepy
from tweepy.error import TweepError
import unidecode

from bot.conf.cfg import getConfig
from bot.log.log import setup_logging
from bot.twitter import getAuth
from moderation.models import SocialUser

from .addstatusdj import addstatus_if_not_exist
from .tasks import handle_retweetroot, handle_question
from .twitter import get_api

setup_logging()
logger = logging.getLogger(__name__)

def has_greenlight(status):
        return isauthorized(status) \
            and not isreplacement(status) \
            and not isretweeted(status) \
            and not isselfstatus(status) \
            and not isquote(status)

def has_retweet_hashtag( status ):
    """ Returns True if the tweet contains a hashtag that is in the retweet_hashtag_list.
    Returns False otherwise.
    In this function a "keyword" is always a hashtag. We must remove the #
    (first character) before comparing the string with the text of the hashtag entity.
    """
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    #logger.debug("status in has_retweet_hashtag: %s", status)
    hashtags = status["entities"]["hashtags"]
    ismatch = False
    for hashtag in hashtags:
        for keyword in getConfig()["keyword_retweet_list"]:
            if keyword[1:].lower() == hashtag["text"].lower():
                ismatch = True
    return ismatch

def isreply( status ):
    "is status in reply to screen name or status or user?"
    logger.debug("in_reply_to_screen_name? %s" , status['in_reply_to_screen_name'])
    logger.debug("in_reply_to_status_id_str? %s" , status['in_reply_to_status_id_str'])
    logger.debug("in_reply_to_user_id_str? %s" , status['in_reply_to_user_id_str'])
    reply_screen = str(status['in_reply_to_screen_name'])
    reply_id = str(status['in_reply_to_status_id_str'])
    reply_user = str(status['in_reply_to_user_id_str'])
    isreply = not (reply_screen == "None" and
               reply_id == "None" and
               reply_user == "None")
    log = "is this status a reply? %s" % isreply
    logger.debug(log)
    return isreply

def isquestion ( status ):
    "Does this status text contain a question mark?"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    isquestion = '?' in status['full_text']
    return isquestion

def okrt( status ):
    """
    Should the bot retweet this doc(s)toctoc question?
    """
    return has_retweet_hashtag(status) and isauthorized(status)

def isretweeted(status):
    return status['retweeted']

def isselfstatus(status):
    if status['user']['id'] == getConfig()['settings']['bot_id']:
        logger.debug("self status, no RT")
        return False

def is_following_rules(status):
    """
    Does this status follow the structural rules?
    * Is Not a retweet
    * Is a question
    * Is not a reply
    """

    if isrt(status):
        return False

    if isreply(status):
        handle_retweetroot.apply_async(args=(status['id'],))
        return False

    if not isquestion(status):
        handle_question.apply_async(args=(status['id'],), countdown=10, expires=300)
        return False

    return True

def isrt( status ):
    "is this status a RT?"
    isrt = False
    isrt = "retweeted_status" in status
    logger.debug("is this status a retweet? %s" , isrt)
    return isrt

def isquote( status ):
    " is this a quoted status?"
    logger.debug("is this status a quoted status? %s" , status['is_quote_status'])
    return status['is_quote_status']

def isknown( status ):
    user_id = status['user']['id']
    if user_id not in SocialUser.objects.moderated_users():
        logger.debug("user %s UNKNOWN", user_id)
        return False
    logger.debug("user %s known", user_id)
    return True

def isauthorized(status):
    logger.debug(f"isauthorized(status): {status['user']['id'] in SocialUser.objects.authorized_users()}")
    return status['user']['id'] in SocialUser.objects.authorized_users()

def okbadlist( status ):
    "filter bad words"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    wordbadlist = set((u"remplacant",u"RT",u"remplaçant"))
    logger.debug("Does text contain bad words? %s", any(word in status['full_text'].split() for word in wordbadlist))
    return not any(word in status['full_text'].split() for word in wordbadlist)

def isreplacement( status ):
    "remove status asking for a replacement physician"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    text = status['full_text']
    #remplacementbadlist = frozenset([u"remplacement", u"rempla"])
    remplacantbadlist = frozenset(["rempla", "remplacant", "remplaçant"])
    monthlist = frozenset([u"janvier", u"fevrier", u"mars", u"avril", u"mai",
                          u"juin", u"juillet", u"aout", u"septembre",
                          u"octobre", u"novembre", u"decembre"])
    wordlist = frozenset(unidecode.unidecode(text).split())
    replacement = (bool(wordlist.intersection(monthlist)) and \
                     bool(wordlist.intersection(remplacantbadlist))) or \
                    (bool(wordlist.intersection(["du"])) and \
                     bool(wordlist.intersection(["au"])) and \
                     bool(wordlist.intersection(remplacantbadlist))) or \
                    (bool(wordlist.intersection(["cherche"])) and \
                     bool(wordlist.intersection(remplacantbadlist)))
    logger.debug("bool(replacement) == %s", bool(replacement))
    return bool(replacement)

def retweet(status_id) -> bool:
    api = tweepy.API(getAuth())
    try:
        api.retweet(status_id)
    except TweepError as e:
        logger.error(str(e))
        return False
    return True

def get_search_string():
    config = getConfig()
    keyword_list = config["keyword_retweet_list"]
    search_string = " OR ".join ( keyword_list )
    search_string = search_string + u" -filter:retweets"
    logger.debug("search_string: %s" % (search_string))
    return search_string

def last_id_file():
    hashedHashtag = hashlib.md5(get_search_string().encode('ascii')).hexdigest()
    last_id_filename = "last_id_hashtag_%s" % hashedHashtag
    rt_bot_path = os.path.dirname(os.path.abspath(__file__))
    last_id_file = os.path.join(rt_bot_path, last_id_filename)
    return last_id_file

def save_point():
    # retrieve last savepoint if available
    try:
        with open(last_id_file(), "r") as f:
            savepoint = f.read()
    except IOError:
        savepoint = ""
        logger.debug("No savepoint found. Bot is now searching for results")

    return savepoint

def timeline_iterator():
    # Tweet language (empty = all languages)
    config = getConfig()
    tweetLanguage = config["settings"]["tweet_language"]

    api = get_api()
    #return tweepy.Cursor(api.search, q=get_search_string(), since_id=save_point(), lang=tweetLanguage, tweet_mode='extended').items(config["settings"]["number_of_rt"])
    return tweepy.Cursor(api.search, q=get_search_string(), lang=tweetLanguage, tweet_mode='extended').items(config["settings"]["number_of_rt"])

def main():
    from .onstatus import triage
    timelineIterator = timeline_iterator()
    current_sid = None
    for tweet in timelineIterator:
        triage(tweet.id)
        #current_sid = tweet.id
        #addstatus_if_not_exist(tweet.id)
        #status = api.get_status(tweet.id,  tweet_mode='extended')
        #status_json = status._json
        #logger.debug(" raw_status: %s", status._json)
        #user = tweet.user
        #screenname = user.screen_name
        #logger.debug("screen name: %s", screenname)
        #logger.debug("text: %s", tweet.full_text.encode('utf-8'))
        #logger.debug("Is user authorized? %s", isauthorized(status_json))
        #logger.debug("isrt: %s", isrt(status_json))
        #logger.debug("status 'retweeted': %s", status_json['retweeted'])
        #logger.debug("is this a question? %s", isquestion(status_json))
        #if okrt(status_json) and is_following_rules(status_json):
        #    oklist.append(tweet)
        #    logger.debug(":) OK for RT")
        #else:
        #    logger.debug(":( not ok for RT")

    #if current_sid is not None:
    #    last_tweet_id = current_sid
    #else:
    #    last_tweet_id = save_point()

    # write last retweeted tweet id to file
    #with open(last_id_file(), "w") as f:
    #    f.write(str(last_tweet_id))

def retweet_lst(lst):
    api = get_api()
    tw_counter = 0
    err_counter = 0
    for status in lst:
        try:
            logger.debug("(%(date)s) %(name)s: %(message)s\n" % \
                  {"date": status.created_at,
                   "name": status.author.screen_name.encode('utf-8'),
                   "message": status.full_text.encode('utf-8')})

            api.retweet(status.id)
            tw_counter += 1
        except tweepy.error.TweepError as e:
            # just in case tweet got deleted in the meantime or already retweeted
            err_counter += 1
            logger.debug("error: %s", e)
            continue

    logger.info("Finished. %d Tweets retweeted, %d errors occured." % (tw_counter, err_counter))

if __name__ == '__main__':
    main()<commit_msg>comment out save_point() function<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

u"""
Bot to enhance French healthcare professionals exchanges on Twitter.

Author: Jérome Pinguet.
License: GNU General Public License version 3 (see 'LICENSE' for details.)

DocTocTocBot started as a modification of the single file [retweet-bot](https://github.com/basti2342/retweet-bot)
by [Basti](https://github.com/basti2342) License: Mozilla Public License version 2.0
  * timelineIterator
  * savepoint
  * RT for loop with error & RT counters
"""

import hashlib
import logging
import os
import tweepy
from tweepy.error import TweepError
import unidecode

from bot.conf.cfg import getConfig
from bot.log.log import setup_logging
from bot.twitter import getAuth
from moderation.models import SocialUser

from .addstatusdj import addstatus_if_not_exist
from .tasks import handle_retweetroot, handle_question
from .twitter import get_api

setup_logging()
logger = logging.getLogger(__name__)

def has_greenlight(status):
        return isauthorized(status) \
            and not isreplacement(status) \
            and not isretweeted(status) \
            and not isselfstatus(status) \
            and not isquote(status)

def has_retweet_hashtag( status ):
    """ Returns True if the tweet contains a hashtag that is in the retweet_hashtag_list.
    Returns False otherwise.
    In this function a "keyword" is always a hashtag. We must remove the #
    (first character) before comparing the string with the text of the hashtag entity.
    """
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    #logger.debug("status in has_retweet_hashtag: %s", status)
    hashtags = status["entities"]["hashtags"]
    ismatch = False
    for hashtag in hashtags:
        for keyword in getConfig()["keyword_retweet_list"]:
            if keyword[1:].lower() == hashtag["text"].lower():
                ismatch = True
    return ismatch

def isreply( status ):
    "is status in reply to screen name or status or user?"
    logger.debug("in_reply_to_screen_name? %s" , status['in_reply_to_screen_name'])
    logger.debug("in_reply_to_status_id_str? %s" , status['in_reply_to_status_id_str'])
    logger.debug("in_reply_to_user_id_str? %s" , status['in_reply_to_user_id_str'])
    reply_screen = str(status['in_reply_to_screen_name'])
    reply_id = str(status['in_reply_to_status_id_str'])
    reply_user = str(status['in_reply_to_user_id_str'])
    isreply = not (reply_screen == "None" and
               reply_id == "None" and
               reply_user == "None")
    log = "is this status a reply? %s" % isreply
    logger.debug(log)
    return isreply

def isquestion ( status ):
    "Does this status text contain a question mark?"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    isquestion = '?' in status['full_text']
    return isquestion

def okrt( status ):
    """
    Should the bot retweet this doc(s)toctoc question?
    """
    return has_retweet_hashtag(status) and isauthorized(status)

def isretweeted(status):
    return status['retweeted']

def isselfstatus(status):
    if status['user']['id'] == getConfig()['settings']['bot_id']:
        logger.debug("self status, no RT")
        return False

def is_following_rules(status):
    """
    Does this status follow the structural rules?
    * Is Not a retweet
    * Is a question
    * Is not a reply
    """

    if isrt(status):
        return False

    if isreply(status):
        handle_retweetroot.apply_async(args=(status['id'],))
        return False

    if not isquestion(status):
        handle_question.apply_async(args=(status['id'],), countdown=10, expires=300)
        return False

    return True

def isrt( status ):
    "is this status a RT?"
    isrt = False
    isrt = "retweeted_status" in status
    logger.debug("is this status a retweet? %s" , isrt)
    return isrt

def isquote( status ):
    " is this a quoted status?"
    logger.debug("is this status a quoted status? %s" , status['is_quote_status'])
    return status['is_quote_status']

def isknown( status ):
    user_id = status['user']['id']
    if user_id not in SocialUser.objects.moderated_users():
        logger.debug("user %s UNKNOWN", user_id)
        return False
    logger.debug("user %s known", user_id)
    return True

def isauthorized(status):
    logger.debug(f"isauthorized(status): {status['user']['id'] in SocialUser.objects.authorized_users()}")
    return status['user']['id'] in SocialUser.objects.authorized_users()

def okbadlist( status ):
    "filter bad words"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    wordbadlist = set((u"remplacant",u"RT",u"remplaçant"))
    logger.debug("Does text contain bad words? %s", any(word in status['full_text'].split() for word in wordbadlist))
    return not any(word in status['full_text'].split() for word in wordbadlist)

def isreplacement( status ):
    "remove status asking for a replacement physician"
    if 'extended_tweet' in status:
        status = status['extended_tweet']
    text = status['full_text']
    #remplacementbadlist = frozenset([u"remplacement", u"rempla"])
    remplacantbadlist = frozenset(["rempla", "remplacant", "remplaçant"])
    monthlist = frozenset([u"janvier", u"fevrier", u"mars", u"avril", u"mai",
                          u"juin", u"juillet", u"aout", u"septembre",
                          u"octobre", u"novembre", u"decembre"])
    wordlist = frozenset(unidecode.unidecode(text).split())
    replacement = (bool(wordlist.intersection(monthlist)) and \
                     bool(wordlist.intersection(remplacantbadlist))) or \
                    (bool(wordlist.intersection(["du"])) and \
                     bool(wordlist.intersection(["au"])) and \
                     bool(wordlist.intersection(remplacantbadlist))) or \
                    (bool(wordlist.intersection(["cherche"])) and \
                     bool(wordlist.intersection(remplacantbadlist)))
    logger.debug("bool(replacement) == %s", bool(replacement))
    return bool(replacement)

def retweet(status_id) -> bool:
    api = tweepy.API(getAuth())
    try:
        api.retweet(status_id)
    except TweepError as e:
        logger.error(str(e))
        return False
    return True

def get_search_string():
    config = getConfig()
    keyword_list = config["keyword_retweet_list"]
    search_string = " OR ".join ( keyword_list )
    search_string = search_string + u" -filter:retweets"
    logger.debug("search_string: %s" % (search_string))
    return search_string

def last_id_file():
    hashedHashtag = hashlib.md5(get_search_string().encode('ascii')).hexdigest()
    last_id_filename = "last_id_hashtag_%s" % hashedHashtag
    rt_bot_path = os.path.dirname(os.path.abspath(__file__))
    last_id_file = os.path.join(rt_bot_path, last_id_filename)
    return last_id_file

def save_point():
    # retrieve last savepoint if available
    pass
    """
    try:
        with open(last_id_file(), "r") as f:
            savepoint = f.read()
    except IOError:
        savepoint = ""
        logger.debug("No savepoint found. Bot is now searching for results")

    return savepoint
    """

def timeline_iterator():
    # Tweet language (empty = all languages)
    config = getConfig()
    tweetLanguage = config["settings"]["tweet_language"]

    api = get_api()
    #return tweepy.Cursor(api.search, q=get_search_string(), since_id=save_point(), lang=tweetLanguage, tweet_mode='extended').items(config["settings"]["number_of_rt"])
    return tweepy.Cursor(api.search, q=get_search_string(), lang=tweetLanguage, tweet_mode='extended').items(config["settings"]["number_of_rt"])

def main():
    from .onstatus import triage
    timelineIterator = timeline_iterator()
    current_sid = None
    for tweet in timelineIterator:
        triage(tweet.id)
        #current_sid = tweet.id
        #addstatus_if_not_exist(tweet.id)
        #status = api.get_status(tweet.id,  tweet_mode='extended')
        #status_json = status._json
        #logger.debug(" raw_status: %s", status._json)
        #user = tweet.user
        #screenname = user.screen_name
        #logger.debug("screen name: %s", screenname)
        #logger.debug("text: %s", tweet.full_text.encode('utf-8'))
        #logger.debug("Is user authorized? %s", isauthorized(status_json))
        #logger.debug("isrt: %s", isrt(status_json))
        #logger.debug("status 'retweeted': %s", status_json['retweeted'])
        #logger.debug("is this a question? %s", isquestion(status_json))
        #if okrt(status_json) and is_following_rules(status_json):
        #    oklist.append(tweet)
        #    logger.debug(":) OK for RT")
        #else:
        #    logger.debug(":( not ok for RT")

    #if current_sid is not None:
    #    last_tweet_id = current_sid
    #else:
    #    last_tweet_id = save_point()

    # write last retweeted tweet id to file
    #with open(last_id_file(), "w") as f:
    #    f.write(str(last_tweet_id))

def retweet_lst(lst):
    api = get_api()
    tw_counter = 0
    err_counter = 0
    for status in lst:
        try:
            logger.debug("(%(date)s) %(name)s: %(message)s\n" % \
                  {"date": status.created_at,
                   "name": status.author.screen_name.encode('utf-8'),
                   "message": status.full_text.encode('utf-8')})

            api.retweet(status.id)
            tw_counter += 1
        except tweepy.error.TweepError as e:
            # just in case tweet got deleted in the meantime or already retweeted
            err_counter += 1
            logger.debug("error: %s", e)
            continue

    logger.info("Finished. %d Tweets retweeted, %d errors occured." % (tw_counter, err_counter))

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


<commit_before>#!/router/bin/python

import zmq
import json
import re
from collections import namedtuple
from trex_stl_types import *
from utils.common import random_id_gen
import zlib
import struct

class bcolors:
    BLUE = '\033[94m'
    GREEN = '\033[32m'
    YELLOW = '\033[93m'
    RED = '\033[31m'
    MAGENTA = '\033[35m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# sub class to describe a batch
class BatchMessage(object):
    def __init__ (self, rpc_client):
        self.rpc_client = rpc_client
        self.batch_list = []

    def add (self, method_name, params={}):

        id, msg = self.rpc_client.create_jsonrpc_v2(method_name, params, encode = False)
        self.batch_list.append(msg)

    def invoke(self, block = False):
        if not self.rpc_client.connected:
            return RC_ERR("Not connected to server")

        msg = json.dumps(self.batch_list)

        return self.rpc_client.send_msg(msg)

# JSON RPC v2.0 client
class JsonRpcClient(object):

    MSG_COMPRESS_THRESHOLD = 4096
    MSG_COMPRESS_HEADER_MAGIC = 0xABE85CEA

    def __init__ (self, default_server, default_port, logger):
        self.logger = logger
        self.connected = False

        # default values
        self.port   = default_port
        self.server = default_server

        self.id_gen = random_id_gen()

    def get_connection_details (self):
        rc = {}
        rc['server'] = self.server
        rc['port']   = self.port

        return rc

    # pretty print for JSON
    def pretty_json (self, json_str, use_colors = True):
        pretty_str = json.dumps(json.loads(json_str), indent = 4, separators=(',', ': '), sort_keys = True)

        if not use_colors:
            return pretty_str

        try:
            # int numbers
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*[^.])',r'\1{0}\2{1}'.format(bcolors.BLUE, bcolors.ENDC), pretty_str)
            # float
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*\.[0-9]+)',r'\1{0}\2{1}'.format(bcolors.MAGENTA, bcolors.ENDC), pretty_str)
            # strings

            pretty_str = re.sub(r'([ ]*:[ ]+)("[^"]*")',r'\1{0}\2{1}'.format(bcolors.RED, bcolors.ENDC), pretty_str)
            pretty_str = re.sub(r"('[^']*')", r'{0}\1{1}'.format(bcolors.MAGENTA, bcolors.RED), pretty_str)
        except :
            pass

        return pretty_str

    def verbose_msg (self, msg):
        self.logger.log("\n\n[verbose] " + msg, level = self.logger.VERBOSE_HIGH)

    # batch messages
    def create_batch (self):
        return BatchMessage(self)

    def create_jsonrpc_v2 (self, method_name, params = {}, encode = True):
        msg = {}
        msg["jsonrpc"] = "2.0"
        msg["method"]  = method_name

        msg["params"] = params

        msg["id"] = self.id_gen.next()

        if encode:
            return id, json.dumps(msg)
        else:
            return id, msg

    def invoke_rpc_method (self, method_name, params = {}):
        if not self.connected:
            return RC_ERR("Not connected to server")

        id, msg = self.create_jsonrpc_v2(method_name, params)

        return self.send_msg(msg)

    def compress_msg (self, msg):
        # compress
        compressed = zlib.compress(msg)
        new_msg = struct.pack(">II", self.MSG_COMPRESS_HEADER_MAGIC, len(msg)) + compressed
        return new_msg

    def decompress_msg (self, msg):
        if len(msg) < 8:
            return None

        t = struct.unpack(">II", msg[:8])
        if (t[0] != self.MSG_COMPRESS_HEADER_MAGIC):
            return None

        x = zlib.decompress(msg[8:])
        if len(x) != t[1]:
            return None

        return x

    def send_msg (self, msg):
        # print before
        if self.logger.check_verbose(self.logger.VERBOSE_HIGH):
            self.verbose_msg("Sending Request To Server:\n\n" + self.pretty_json(msg) + "\n")

        if len(msg) > self.MSG_COMPRESS_THRESHOLD:
            response = self.send_raw_msg(self.compress_msg(msg))
            if response:
                response = self.decompress_msg(response)
        else:
            response = self.send_raw_msg(msg)

        if response == None:
            return RC_ERR("*** [RPC] - Failed to decode response from server")

        # print after
        if self.logger.check_verbose(self.logger.VERBOSE_HIGH):
            self.verbose_msg("Server Response:\n\n" + self.pretty_json(response) + "\n")

        # process response (batch and regular)

        response_json = json.loads(response)

        if isinstance(response_json, list):
            return self.process_batch_response(response_json)
        else:
            return self.process_single_response(response_json)

    # low level send of string message
    def send_raw_msg (self, msg):

        tries = 0
        while True:
            try:
                self.socket.send(msg)
                break
            except zmq.Again:
                tries += 1
                if tries > 5:
                    self.disconnect()
                    return RC_ERR("*** [RPC] - Failed to send message to server")

        tries = 0
        while True:
            try:
                response = self.socket.recv()
                break
            except zmq.Again:
                tries += 1
                if tries > 5:
                    self.disconnect()
                    return RC_ERR("*** [RPC] - Failed to get server response at {0}".format(self.transport))

        return response

    # processs a single response from server
    def process_single_response (self, response_json):

        if (response_json.get("jsonrpc") != "2.0"):
            return RC_ERR("Malformed Response ({0})".format(str(response_json)))

        # error reported by server
        if ("error" in response_json):
            if "specific_err" in response_json["error"]:
                return RC_ERR(response_json["error"]["specific_err"])
            else:
                return RC_ERR(response_json["error"]["message"])

        # if no error there should be a result
        if ("result" not in response_json):
            return RC_ERR("Malformed Response ({0})".format(str(response_json)))

        return RC_OK(response_json["result"])

    # process a batch response
    def process_batch_response (self, response_json):
        rc_batch = RC()

        for single_response in response_json:
            rc = self.process_single_response(single_response)
            rc_batch.add(rc)

        return rc_batch

    def disconnect (self):
        if self.connected:
            self.socket.close(linger = 0)
            self.context.destroy(linger = 0)
            self.connected = False
            return RC_OK()
        else:
            return RC_ERR("Not connected to server")

    def connect(self, server = None, port = None):
        if self.connected:
            self.disconnect()

        self.context = zmq.Context()

        self.server = (server if server else self.server)
        self.port = (port if port else self.port)

        #  Socket to talk to server
        self.transport = "tcp://{0}:{1}".format(self.server, self.port)

        self.socket = self.context.socket(zmq.REQ)
        try:
            self.socket.connect(self.transport)
        except zmq.error.ZMQError as e:
            return RC_ERR("ZMQ Error: Bad server or port name: " + str(e))

        self.socket.setsockopt(zmq.SNDTIMEO, 1000)
        self.socket.setsockopt(zmq.RCVTIMEO, 1000)

        self.connected = True

        rc = self.invoke_rpc_method('ping')
        if not rc:
            self.connected = False
            return rc

        return RC_OK()

    def reconnect(self):
        # connect using current values
        return self.connect()

        if not self.connected:
            return RC_ERR("Not connected to server")

        # reconnect
        return self.connect(self.server, self.port)

    def is_connected(self):
        return self.connected

    def __del__(self):
        self.logger.log("Shutting down RPC client\n")
        if hasattr(self, "context"):
            self.context.destroy(linger=0)

<commit_msg>timeout connection error message was broken<commit_after>#!/router/bin/python

import zmq
import json
import re
from collections import namedtuple
from trex_stl_types import *
from utils.common import random_id_gen
import zlib
import struct

class bcolors:
    BLUE = '\033[94m'
    GREEN = '\033[32m'
    YELLOW = '\033[93m'
    RED = '\033[31m'
    MAGENTA = '\033[35m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

# sub class to describe a batch
class BatchMessage(object):
    def __init__ (self, rpc_client):
        self.rpc_client = rpc_client
        self.batch_list = []

    def add (self, method_name, params={}):

        id, msg = self.rpc_client.create_jsonrpc_v2(method_name, params, encode = False)
        self.batch_list.append(msg)

    def invoke(self, block = False):
        if not self.rpc_client.connected:
            return RC_ERR("Not connected to server")

        msg = json.dumps(self.batch_list)

        return self.rpc_client.send_msg(msg)

# JSON RPC v2.0 client
class JsonRpcClient(object):

    MSG_COMPRESS_THRESHOLD = 4096
    MSG_COMPRESS_HEADER_MAGIC = 0xABE85CEA

    def __init__ (self, default_server, default_port, logger):
        self.logger = logger
        self.connected = False

        # default values
        self.port   = default_port
        self.server = default_server

        self.id_gen = random_id_gen()

    def get_connection_details (self):
        rc = {}
        rc['server'] = self.server
        rc['port']   = self.port

        return rc

    # pretty print for JSON
    def pretty_json (self, json_str, use_colors = True):
        pretty_str = json.dumps(json.loads(json_str), indent = 4, separators=(',', ': '), sort_keys = True)

        if not use_colors:
            return pretty_str

        try:
            # int numbers
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*[^.])',r'\1{0}\2{1}'.format(bcolors.BLUE, bcolors.ENDC), pretty_str)
            # float
            pretty_str = re.sub(r'([ ]*:[ ]+)(\-?[1-9][0-9]*\.[0-9]+)',r'\1{0}\2{1}'.format(bcolors.MAGENTA, bcolors.ENDC), pretty_str)
            # strings

            pretty_str = re.sub(r'([ ]*:[ ]+)("[^"]*")',r'\1{0}\2{1}'.format(bcolors.RED, bcolors.ENDC), pretty_str)
            pretty_str = re.sub(r"('[^']*')", r'{0}\1{1}'.format(bcolors.MAGENTA, bcolors.RED), pretty_str)
        except :
            pass

        return pretty_str

    def verbose_msg (self, msg):
        self.logger.log("\n\n[verbose] " + msg, level = self.logger.VERBOSE_HIGH)

    # batch messages
    def create_batch (self):
        return BatchMessage(self)

    def create_jsonrpc_v2 (self, method_name, params = {}, encode = True):
        msg = {}
        msg["jsonrpc"] = "2.0"
        msg["method"]  = method_name

        msg["params"] = params

        msg["id"] = self.id_gen.next()

        if encode:
            return id, json.dumps(msg)
        else:
            return id, msg

    def invoke_rpc_method (self, method_name, params = {}):
        if not self.connected:
            return RC_ERR("Not connected to server")

        id, msg = self.create_jsonrpc_v2(method_name, params)

        return self.send_msg(msg)

    def compress_msg (self, msg):
        # compress
        compressed = zlib.compress(msg)
        new_msg = struct.pack(">II", self.MSG_COMPRESS_HEADER_MAGIC, len(msg)) + compressed
        return new_msg

    def decompress_msg (self, msg):
        if len(msg) < 8:
            return None

        t = struct.unpack(">II", msg[:8])
        if (t[0] != self.MSG_COMPRESS_HEADER_MAGIC):
            return None

        x = zlib.decompress(msg[8:])
        if len(x) != t[1]:
            return None

        return x

    def send_msg (self, msg):
        # print before
        if self.logger.check_verbose(self.logger.VERBOSE_HIGH):
            self.verbose_msg("Sending Request To Server:\n\n" + self.pretty_json(msg) + "\n")

        if len(msg) > self.MSG_COMPRESS_THRESHOLD:
            response = self.send_raw_msg(self.compress_msg(msg))
            if response:
                response = self.decompress_msg(response)
        else:
            response = self.send_raw_msg(msg)

        if not response:
            return response

        # print after
        if self.logger.check_verbose(self.logger.VERBOSE_HIGH):
            self.verbose_msg("Server Response:\n\n" + self.pretty_json(response) + "\n")

        # process response (batch and regular)
        try:
            response_json = json.loads(response)
        except (TypeError, ValueError):
            return RC_ERR("*** [RPC] - Failed to decode response from server")

        if isinstance(response_json, list):
            return self.process_batch_response(response_json)
        else:
            return self.process_single_response(response_json)

    # low level send of string message
    def send_raw_msg (self, msg):

        tries = 0
        while True:
            try:
                self.socket.send(msg)
                break
            except zmq.Again:
                tries += 1
                if tries > 5:
                    self.disconnect()
                    return RC_ERR("*** [RPC] - Failed to send message to server")

        tries = 0
        while True:
            try:
                response = self.socket.recv()
                break
            except zmq.Again:
                tries += 1
                if tries > 5:
                    self.disconnect()
                    return RC_ERR("*** [RPC] - Failed to get server response at {0}".format(self.transport))

        return response

    # processs a single response from server
    def process_single_response (self, response_json):

        if (response_json.get("jsonrpc") != "2.0"):
            return RC_ERR("Malformed Response ({0})".format(str(response_json)))

        # error reported by server
        if ("error" in response_json):
            if "specific_err" in response_json["error"]:
                return RC_ERR(response_json["error"]["specific_err"])
            else:
                return RC_ERR(response_json["error"]["message"])

        # if no error there should be a result
        if ("result" not in response_json):
            return RC_ERR("Malformed Response ({0})".format(str(response_json)))

        return RC_OK(response_json["result"])

    # process a batch response
    def process_batch_response (self, response_json):
        rc_batch = RC()

        for single_response in response_json:
            rc = self.process_single_response(single_response)
            rc_batch.add(rc)

        return rc_batch

    def disconnect (self):
        if self.connected:
            self.socket.close(linger = 0)
            self.context.destroy(linger = 0)
            self.connected = False
            return RC_OK()
        else:
            return RC_ERR("Not connected to server")

    def connect(self, server = None, port = None):
        if self.connected:
            self.disconnect()

        self.context = zmq.Context()

        self.server = (server if server else self.server)
        self.port = (port if port else self.port)

        #  Socket to talk to server
        self.transport = "tcp://{0}:{1}".format(self.server, self.port)

        self.socket = self.context.socket(zmq.REQ)
        try:
            self.socket.connect(self.transport)
        except zmq.error.ZMQError as e:
            return RC_ERR("ZMQ Error: Bad server or port name: " + str(e))

        self.socket.setsockopt(zmq.SNDTIMEO, 1000)
        self.socket.setsockopt(zmq.RCVTIMEO, 1000)

        self.connected = True

        rc = self.invoke_rpc_method('ping')
        if not rc:
            self.connected = False
            return rc

        return RC_OK()

    def reconnect(self):
        # connect using current values
        return self.connect()

        if not self.connected:
            return RC_ERR("Not connected to server")

        # reconnect
        return self.connect(self.server, self.port)

    def is_connected(self):
        return self.connected

    def __del__(self):
        self.logger.log("Shutting down RPC client\n")
        if hasattr(self, "context"):
            self.context.destroy(linger=0)

<|endoftext|>