<commit_before>import os
import sys
import re
import pickle
import logging

from pysb.bng import BngInterfaceError

from indra.assemblers import PysbAssembler, pysb_assembler

logging.basicConfig(format='%(levelname)s: %(name)s - %(message)s',
                    level=logging.INFO)
logger = logging.getLogger('MSA')

from kqml import KQMLPerformative, KQMLList

from indra.sources.trips.processor import TripsProcessor
from indra import has_config

from bioagents import Bioagent

if has_config('INDRA_DB_REST_URL') and has_config('INDRA_DB_REST_API_KEY'):
    from indra.sources.indra_db_rest import get_statements, IndraDBRestError

    CAN_CHECK_STATEMENTS = True
else:
    logger.warning("Database web api not specified. Cannot get background.")
    CAN_CHECK_STATEMENTS = False

def _read_signor_afs():
    path = os.path.dirname(os.path.abspath(__file__)) + \
            '/../resources/signor_active_forms.pkl'
    with open(path, 'rb') as pkl_file:
        stmts = pickle.load(pkl_file)
    if isinstance(stmts, dict):
        signor_afs = []
        for _, stmt_list in stmts.items():
            signor_afs += stmt_list
    else:
        signor_afs = stmts
    return signor_afs

class MSA_Module(Bioagent):
    name = 'MSA'
    tasks = ['PHOSPHORYLATION-ACTIVATING', 'FIND-IMMEDIATE-RELATION']
    signor_afs = _read_signor_afs()

    def respond_phosphorylation_activating(self, content):
        """Return response content to phosphorylation_activating request."""
        if not CAN_CHECK_STATEMENTS:
            return self.make_failure(
                'NO_KNOWLEDGE_ACCESS',
                'Cannot access the database through the web api.'
                )
        heading = content.head()
        m = re.match('(\w+)-(\w+)', heading)
        if m is None:
            return self.make_failure('UNKNOWN_ACTION')
        action, polarity = [s.lower() for s in m.groups()]
        target_ekb = content.gets('target')
        if target_ekb is None or target_ekb == '':
            return self.make_failure('MISSING_TARGET')
        agent = self._get_agent(target_ekb)
        logger.debug('Found agent (target): %s.' % agent.name)
        residue = content.gets('residue')
        position = content.gets('position')
        related_result_dict = {}
        logger.info("Looking for statements with agent %s of type %s."
                    % (str(agent), 'ActiveForm'))
        for namespace, name in agent.db_refs.items():
            # TODO: Remove this eventually, as it is a temporary work-around.
            if namespace == 'FPLX':
                namespace = 'BE'
            logger.info("Checking namespace: %s" % namespace)
            stmts = get_statements(agents=['%s@%s' % (name, namespace)],
                                   stmt_type='ActiveForm')
            for s in stmts:
                if self._matching(s, residue, position, action, polarity):
                    related_result_dict[s.matches_key()] = s
        logger.info("Found %d matching statements." % len(related_result_dict))
        if not len(related_result_dict):
            return self.make_failure(
                'MISSING_MECHANISM',
                "Could not find statement matching phosphorylation activating "
                "%s, %s, %s, %s." % (agent.name, residue, position,
                                     'phosphorylation')
                )
        else:
            self.send_provenance_for_stmts(
                related_result_dict.values(),
                "Phosphorylation at %s%s activates %s." % (
                    residue,
                    position,
                    agent.name
                    )
                )
            msg = KQMLPerformative('SUCCESS')
            msg.set('is-activating', 'TRUE')
            return msg

    def respond_find_immediate_relation(self, content):
        """Find statements matching a query for FIND-IMMEDIATE-RELATION task."""
        agent_dict = dict.fromkeys(['subject', 'object'])
        for pos, loc in [('subject', 'source'), ('object', 'target')]:
            ekb = content.gets(loc)
            try:
                agent = self._get_agent(ekb)
                if agent is None or agent == 'None':
                    agent_dict[pos] = None
                else:
                    agent_dict[pos] = {'name': agent.name}
                    agent_dict[pos].update(agent.db_refs)
            except Exception as e:
                logger.error("Got exception while converting ekb for %s "
                             "(%s) into an agent." % (pos, ekb))
                logger.exception(e)
                return self.make_failure('MISSING_TARGET')
        stmt_type = content.gets('type')
        if stmt_type == 'unknown':
            stmt_type = None
        nl_question = ('{subject} {verb} of {object}'
                       .format(verb=stmt_type,
                               **{k: None if v is None else v['name']
                                  for k, v in agent_dict.items()}))
        logger.info("Got a query for %s." % nl_question)
        # Try to get related statements.
        try:
            input_dict = {'stmt_type': stmt_type}

            # Use the best available db ref for each agent.
            for pos, ref_dict in agent_dict.items():
                if ref_dict is None:
                    input_dict[pos] = None
                else:
                    for key in ['HGNC', 'FPLX', 'CHEBI', 'name', 'TEXT']:
                        if key in ref_dict.keys():
                            input_dict[pos] = ref_dict[key]

            # Actually get the statements.
            stmts = get_statements(**input_dict)
        except IndraDBRestError as e:
            logger.error("Failed to get statements.")
            logger.exception(e)
            return self.make_failure('MISSING_MECHANISM')

        # For now just list the statements in the provenance tab. Only captures
        # the top 5.
        try:
            self.send_display_stmts(stmts, nl_question)
        except Exception as e:
            logger.warning("Failed to send provenance.")
            logger.exception(e)

        # Assuming we haven't hit any errors yet, return SUCCESS
        resp = KQMLPerformative('SUCCESS')
        resp.set('relations-found', str(len(stmts)))
        return resp

    def send_display_stmts(self, stmts, nl_question):
        self.send_provenance_for_stmts(stmts, nl_question)
        resource = _make_sbgn(stmts)
        content = KQMLList('open-query-window')
        content.sets('cyld', '#1')
        content.sets('graph', resource)
        self.tell(content)

    @staticmethod
    def _get_agent(agent_ekb):
        tp = TripsProcessor(agent_ekb)
        terms = tp.tree.findall('TERM')
        if len(terms):
            term_id = terms[0].attrib['id']
            agent = tp._get_agent_by_id(term_id, None)
        else:
            agent = None
        return agent

    def _matching(self, stmt, residue, position, action, polarity):
        if stmt.is_active is not (polarity == 'activating'):
            return False
        matching_residues = any([
            m.residue == residue
            and m.position == position
            and m.mod_type == action
            for m in stmt.agent.mods])
        return matching_residues

def _make_sbgn(stmts):
    pa = PysbAssembler()
    pa.add_statements(stmts)
    pa.make_model()
    pa.add_default_initial_conditions(10)
    for m in pa.model.monomers:
        pysb_assembler.set_extended_initial_condition(pa.model, m, 0)
    try:
        sbgn_str = pa.export_model('sbgn')
    except BngInterfaceError:
        logger.error('Reaction network could not be generated for SBGN.')
        return None
    return sbgn_str

if __name__ == "__main__":
    MSA_Module(argv=sys.argv[1:])
<commit_msg>Try to use SBGN assembler (WIP)<commit_after>import os
import sys
import re
import pickle
import logging

from pysb.bng import BngInterfaceError

from indra.assemblers import PysbAssembler, pysb_assembler

logging.basicConfig(format='%(levelname)s: %(name)s - %(message)s',
                    level=logging.INFO)
logger = logging.getLogger('MSA')

from kqml import KQMLPerformative, KQMLList

from indra.sources.trips.processor import TripsProcessor
from indra import has_config

from bioagents import Bioagent

if has_config('INDRA_DB_REST_URL') and has_config('INDRA_DB_REST_API_KEY'):
    from indra.sources.indra_db_rest import get_statements, IndraDBRestError

    CAN_CHECK_STATEMENTS = True
else:
    logger.warning("Database web api not specified. Cannot get background.")
    CAN_CHECK_STATEMENTS = False

def _read_signor_afs():
    path = os.path.dirname(os.path.abspath(__file__)) + \
            '/../resources/signor_active_forms.pkl'
    with open(path, 'rb') as pkl_file:
        stmts = pickle.load(pkl_file)
    if isinstance(stmts, dict):
        signor_afs = []
        for _, stmt_list in stmts.items():
            signor_afs += stmt_list
    else:
        signor_afs = stmts
    return signor_afs

class MSA_Module(Bioagent):
    name = 'MSA'
    tasks = ['PHOSPHORYLATION-ACTIVATING', 'FIND-IMMEDIATE-RELATION']
    signor_afs = _read_signor_afs()

    def respond_phosphorylation_activating(self, content):
        """Return response content to phosphorylation_activating request."""
        if not CAN_CHECK_STATEMENTS:
            return self.make_failure(
                'NO_KNOWLEDGE_ACCESS',
                'Cannot access the database through the web api.'
                )
        heading = content.head()
        m = re.match('(\w+)-(\w+)', heading)
        if m is None:
            return self.make_failure('UNKNOWN_ACTION')
        action, polarity = [s.lower() for s in m.groups()]
        target_ekb = content.gets('target')
        if target_ekb is None or target_ekb == '':
            return self.make_failure('MISSING_TARGET')
        agent = self._get_agent(target_ekb)
        logger.debug('Found agent (target): %s.' % agent.name)
        residue = content.gets('residue')
        position = content.gets('position')
        related_result_dict = {}
        logger.info("Looking for statements with agent %s of type %s."
                    % (str(agent), 'ActiveForm'))
        for namespace, name in agent.db_refs.items():
            # TODO: Remove this eventually, as it is a temporary work-around.
            if namespace == 'FPLX':
                namespace = 'BE'
            logger.info("Checking namespace: %s" % namespace)
            stmts = get_statements(agents=['%s@%s' % (name, namespace)],
                                   stmt_type='ActiveForm')
            for s in stmts:
                if self._matching(s, residue, position, action, polarity):
                    related_result_dict[s.matches_key()] = s
        logger.info("Found %d matching statements." % len(related_result_dict))
        if not len(related_result_dict):
            return self.make_failure(
                'MISSING_MECHANISM',
                "Could not find statement matching phosphorylation activating "
                "%s, %s, %s, %s." % (agent.name, residue, position,
                                     'phosphorylation')
                )
        else:
            self.send_provenance_for_stmts(
                related_result_dict.values(),
                "Phosphorylation at %s%s activates %s." % (
                    residue,
                    position,
                    agent.name
                    )
                )
            msg = KQMLPerformative('SUCCESS')
            msg.set('is-activating', 'TRUE')
            return msg

    def respond_find_immediate_relation(self, content):
        """Find statements matching a query for FIND-IMMEDIATE-RELATION task."""
        agent_dict = dict.fromkeys(['subject', 'object'])
        for pos, loc in [('subject', 'source'), ('object', 'target')]:
            ekb = content.gets(loc)
            try:
                agent = self._get_agent(ekb)
                if agent is None or agent == 'None':
                    agent_dict[pos] = None
                else:
                    agent_dict[pos] = {'name': agent.name}
                    agent_dict[pos].update(agent.db_refs)
            except Exception as e:
                logger.error("Got exception while converting ekb for %s "
                             "(%s) into an agent." % (pos, ekb))
                logger.exception(e)
                return self.make_failure('MISSING_TARGET')
        stmt_type = content.gets('type')
        if stmt_type == 'unknown':
            stmt_type = None
        nl_question = ('{subject} {verb} of {object}'
                       .format(verb=stmt_type,
                               **{k: None if v is None else v['name']
                                  for k, v in agent_dict.items()}))
        logger.info("Got a query for %s." % nl_question)
        # Try to get related statements.
        try:
            input_dict = {'stmt_type': stmt_type}

            # Use the best available db ref for each agent.
            for pos, ref_dict in agent_dict.items():
                if ref_dict is None:
                    input_dict[pos] = None
                else:
                    for key in ['HGNC', 'FPLX', 'CHEBI', 'name', 'TEXT']:
                        if key in ref_dict.keys():
                            input_dict[pos] = ref_dict[key]

            # Actually get the statements.
            stmts = get_statements(**input_dict)
        except IndraDBRestError as e:
            logger.error("Failed to get statements.")
            logger.exception(e)
            return self.make_failure('MISSING_MECHANISM')

        # For now just list the statements in the provenance tab. Only captures
        # the top 5.
        try:
            self.send_display_stmts(stmts, nl_question)
        except Exception as e:
            logger.warning("Failed to send provenance.")
            logger.exception(e)

        # Assuming we haven't hit any errors yet, return SUCCESS
        resp = KQMLPerformative('SUCCESS')
        resp.set('relations-found', str(len(stmts)))
        return resp

    def send_display_stmts(self, stmts, nl_question):
        #self.send_provenance_for_stmts(stmts, nl_question)
        logger.info('Sending display statements')
        resource = _make_sbgn(stmts)
        logger.info(resource)
        content = KQMLList('open-query-window')
        content.sets('cyld', '#1')
        content.sets('graph', resource)
        self.tell(content)

    @staticmethod
    def _get_agent(agent_ekb):
        tp = TripsProcessor(agent_ekb)
        terms = tp.tree.findall('TERM')
        if len(terms):
            term_id = terms[0].attrib['id']
            agent = tp._get_agent_by_id(term_id, None)
        else:
            agent = None
        return agent

    def _matching(self, stmt, residue, position, action, polarity):
        if stmt.is_active is not (polarity == 'activating'):
            return False
        matching_residues = any([
            m.residue == residue
            and m.position == position
            and m.mod_type == action
            for m in stmt.agent.mods])
        return matching_residues

def _make_sbgn(stmts):
    sa = SBGNAssembler()
    sa.add_statements(stmts)
    sa.make_model()
    sbgn_str = sa.print_model()
    logger.info(sbgn_str)
    return sbgn_str

if __name__ == "__main__":
    MSA_Module(argv=sys.argv[1:])
<|endoftext|>


<commit_before><commit_msg>[BugFix: ZSTACK-43067] fix wrong ipxe.efi path<commit_after><|endoftext|>


<commit_before>import copy
import json
import logging
import os
import pprint
import shutil
from pathlib import Path

from i3configger import context, exc

log = logging.getLogger(__name__)
cliMainOverrideMap = {}
"""holds parsed args if started from cli and was initialized"""

DEFAULTS = {
    "main": {
        "target": "../config",
        "i3_refresh_msg": "reload",
        "status_command": "i3status",
        "log": None,
        "notify": False,
    },
    "bars": {
        "defaults": {"template": "tpl", "target": "..", "select": "default"},
        "targets": {},
    },
}

class I3configgerConfig:
    PARTIALS_NAME = "config.d"
    CONFIG_NAME = "i3configger.json"
    MESSAGES_NAME = ".messages.json"

    def __init__(self, load=True):
        i3configBasePath = get_i3wm_config_path()
        # FIXME rename partialsPath -> i3configgerPath
        self.partialsPath = i3configBasePath / self.PARTIALS_NAME
        self.configPath = self.partialsPath / self.CONFIG_NAME
        self.messagesPath = self.partialsPath / self.MESSAGES_NAME
        if load:
            self.load()

    def __str__(self):
        return f"{self.__class__.__name__}:\n{pprint.pformat(vars(self))}"

    def load(self):
        """Layered overrides `DEFAULTS` -> `i3configger.json`-> `args`"""
        cnfFromDefaults = copy.deepcopy(DEFAULTS)
        cnfFromFile = fetch(self.configPath)
        self.payload = context.merge(cnfFromDefaults, cnfFromFile)
        mainOverrides = dict(main=cliMainOverrideMap)
        self.payload = context.merge(self.payload, mainOverrides)
        targetPath = Path(self.payload["main"]["target"]).expanduser()
        if targetPath.is_absolute():
            self.targetPath = targetPath.resolve()
        else:
            self.targetPath = (self.partialsPath / targetPath).resolve()

    def get_bar_targets(self):
        """Create a resolved copy of the bar settings."""
        barTargets = {}
        barSettings = self.payload.get("bars", {})
        if not barSettings:
            return barTargets
        defaults = barSettings.get("defaults", {})
        for name, bar in barSettings["targets"].items():
            newBar = dict(bar)
            barTargets[name] = newBar
            for defaultKey, defaultValue in defaults.items():
                if defaultKey not in newBar:
                    newBar[defaultKey] = defaultValue
            container = Path(newBar["target"])
            if not container.is_absolute():
                container = (self.partialsPath / container).resolve()
            newBar["target"] = str(container)
        return barTargets

def fetch(path: Path) -> dict:
    if not path.exists():
        raise exc.ConfigError(f"file not found: {path}")
    # FIXME update to read_text() (also all other occurences)
    with path.open() as f:
        payload = json.load(f)
        log.debug(f"{path}:\n{pprint.pformat(payload)}")
        return payload

def freeze(path, obj):
    with open(path, "w") as f:
        json.dump(obj, f, sort_keys=True, indent=2)
    log.debug(f"froze {pprint.pformat(obj)} to {path}")

def ensure_i3_configger_sanity():
    i3wmConfigPath = get_i3wm_config_path()
    partialsPath = i3wmConfigPath / I3configgerConfig.PARTIALS_NAME
    if not partialsPath.exists():
        log.info(f"create new config folder at {partialsPath}")
        partialsPath.mkdir()
        for candidate in [i3wmConfigPath / "config", Path("etc/i3/config")]:
            if candidate.exists():
                log.info(f"populate config with {candidate}")
                shutil.copy2(candidate, partialsPath / "config.conf")
    configPath = partialsPath / I3configgerConfig.CONFIG_NAME
    if not configPath.exists():
        log.info(f"create default configuration at {configPath}")
        freeze(configPath, DEFAULTS)

def get_i3wm_config_path():
    """Use same search order like i3 (no system stuff though).

    see: https://github.com/i3/i3/blob/4.13/libi3/get_config_path.c#L31
    """
    candidates = [
        Path("~/.i3").expanduser(),
        Path(os.getenv("XDG_CONFIG_HOME", "~/.config/")).expanduser() / "i3",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise exc.ConfigError(
        f"can't find i3 config at the standard locations: {candidates}"
    )
<commit_msg>use path object properly<commit_after>import copy
import json
import logging
import os
import pprint
import shutil
from pathlib import Path

from i3configger import context, exc

log = logging.getLogger(__name__)
cliMainOverrideMap = {}
"""holds parsed args if started from cli and was initialized"""

DEFAULTS = {
    "main": {
        "target": "../config",
        "i3_refresh_msg": "reload",
        "status_command": "i3status",
        "log": None,
        "notify": False,
    },
    "bars": {
        "defaults": {"template": "tpl", "target": "..", "select": "default"},
        "targets": {},
    },
}

class I3configgerConfig:
    PARTIALS_NAME = "config.d"
    CONFIG_NAME = "i3configger.json"
    MESSAGES_NAME = ".messages.json"

    def __init__(self, load=True):
        i3configBasePath = get_i3wm_config_path()
        self.partialsPath = i3configBasePath / self.PARTIALS_NAME
        self.configPath = self.partialsPath / self.CONFIG_NAME
        self.messagesPath = self.partialsPath / self.MESSAGES_NAME
        if load:
            self.load()

    def __str__(self):
        return f"{self.__class__.__name__}:\n{pprint.pformat(vars(self))}"

    def load(self):
        """Layered overrides `DEFAULTS` -> `i3configger.json`-> `args`"""
        cnfFromDefaults = copy.deepcopy(DEFAULTS)
        cnfFromFile = fetch(self.configPath)
        self.payload = context.merge(cnfFromDefaults, cnfFromFile)
        mainOverrides = dict(main=cliMainOverrideMap)
        self.payload = context.merge(self.payload, mainOverrides)
        targetPath = Path(self.payload["main"]["target"]).expanduser()
        if targetPath.is_absolute():
            self.targetPath = targetPath.resolve()
        else:
            self.targetPath = (self.partialsPath / targetPath).resolve()

    def get_bar_targets(self):
        """Create a resolved copy of the bar settings."""
        barTargets = {}
        barSettings = self.payload.get("bars", {})
        if not barSettings:
            return barTargets
        defaults = barSettings.get("defaults", {})
        for name, bar in barSettings["targets"].items():
            newBar = dict(bar)
            barTargets[name] = newBar
            for defaultKey, defaultValue in defaults.items():
                if defaultKey not in newBar:
                    newBar[defaultKey] = defaultValue
            container = Path(newBar["target"])
            if not container.is_absolute():
                container = (self.partialsPath / container).resolve()
            newBar["target"] = str(container)
        return barTargets

def fetch(path: Path) -> dict:
    if not path.exists():
        raise exc.ConfigError(f"config not found at {path}")
    payload = json.loads(path.read_text())
    log.debug(f"{path}:\n{pprint.pformat(payload)}")
    return payload

def freeze(path, obj):
    path.write_text(json.dumps(obj, sort_keys=True, indent=2))
    log.debug(f"froze {pprint.pformat(obj)} to {path}")

def ensure_i3_configger_sanity():
    i3wmConfigPath = get_i3wm_config_path()
    partialsPath = i3wmConfigPath / I3configgerConfig.PARTIALS_NAME
    if not partialsPath.exists():
        log.info(f"create new config folder at {partialsPath}")
        partialsPath.mkdir()
        for candidate in [i3wmConfigPath / "config", Path("etc/i3/config")]:
            if candidate.exists():
                log.info(f"populate config with {candidate}")
                shutil.copy2(candidate, partialsPath / "config.conf")
    configPath = partialsPath / I3configgerConfig.CONFIG_NAME
    if not configPath.exists():
        log.info(f"create default configuration at {configPath}")
        freeze(configPath, DEFAULTS)

def get_i3wm_config_path():
    """Use same search order like i3 (no system stuff though).

    see: https://github.com/i3/i3/blob/4.13/libi3/get_config_path.c#L31
    """
    candidates = [
        Path("~/.i3").expanduser(),
        Path(os.getenv("XDG_CONFIG_HOME", "~/.config/")).expanduser() / "i3",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise exc.ConfigError(
        f"can't find i3 config at the standard locations: {candidates}"
    )
<|endoftext|>


<commit_before>from mongoengine import NotUniqueError
from pymongo.errors import DuplicateKeyError
from requests import get
import logging
from datetime import datetime

from manage_music.models import Track, Artist, Recording, Album

log = logging.getLogger(__name__)

def scraper_practice():
    log.info("Starting Scraper Practice")
    request_json = get_request_json("dakre", 1)
    favorite_tracks_pagination = request_json["pagination"]
    total_pages = favorite_tracks_pagination["total_pages"]
    for index in range(1, total_pages):
        request_json = get_request_json("dakre", index)
        all_favorite_tracks = request_json["favorite_tracks"]
        for favourite_track in all_favorite_tracks:
            # TODO: there's a way to grab the artist info when playing tracks
            # todo: investigate for aliases, description, homepage
            # https://8tracks.com/sets/197579681/next?player=sm&include=track%5Bfaved%2Bannotation%2Bartist_details%5D&mix_id=2183081&track_id=17451458&format=jsonh
            artist = save_artist(favourite_track)
            track = save_track(artist, favourite_track)
            save_recording(artist, favourite_track, track)
            save_album(favourite_track, track)
    log.info("Finished Scraper Practice")

def get_request_json(user, index):
    web_page_url = "https://8tracks.com/users/" + str(user) + "/favorite_tracks" + \
                  "?page=" + str(index) + "&format=jsonh"
    request = get(web_page_url)
    request_json = request.json()
    return request_json

def save_artist(favourite_track):
    artist = Artist()
    artist.name = favourite_track["performer"]
    # todo: append artist name as alias if different from name found
    try:
        artist.save()
    except (DuplicateKeyError, NotUniqueError):
        artist = Artist.objects(name=artist.name)
    return artist

def save_track(artist, favourite_track):
    artists = [artist]
    track = Track()
    track.name = favourite_track["name"]
    track.artists = artists
    try:
        track.save()
    except (DuplicateKeyError, NotUniqueError):
        track = Track.objects(name=track.name)
    return track

def save_recording(artist, favourite_track, track):
    artists = [artist]
    recording = Recording()
    recording.tracks = [track]
    recording.artists = artists
    recording.web_page = favourite_track["url"]
    try:
        recording.save()
    except (DuplicateKeyError, NotUniqueError):
        # todo: maybe try to get this proper from the database
        return recording
    return recording

def save_album(favourite_track, track):
    album = Album()
    album.name = favourite_track["release_name"]
    album.artists = track.artists
    album.tracks = [track]
    album.release_date = get_release_date(favourite_track["year"])
    try:
        album.save()
    except (DuplicateKeyError, NotUniqueError):
        # todo: maybe try to get this proper from the database
        return album
    return album

def get_release_date(integer_date):
    if len(str(integer_date)) == 4:
        return datetime.strptime(str(integer_date), "%Y")
    if len(str(integer_date)) == 6:
        return datetime.strptime(str(integer_date), "%Y%m")
    if len(str(integer_date)) == 8:
        return datetime.strptime(str(integer_date), "%Y%m%d")
<commit_msg>added print statements<commit_after>from mongoengine import NotUniqueError
from pymongo.errors import DuplicateKeyError
from requests import get
import logging
from datetime import datetime

from manage_music.models import Track, Artist, Recording, Album

log = logging.getLogger(__name__)

def scraper_practice():
    # log.info("Starting Scraper Practice")
    print("Starting Scraper Practice")
    request_json = get_request_json("dakre", 1)
    favorite_tracks_pagination = request_json["pagination"]
    total_pages = favorite_tracks_pagination["total_pages"]
    for index in range(1, total_pages):
        request_json = get_request_json("dakre", index)
        all_favorite_tracks = request_json["favorite_tracks"]
        for favourite_track in all_favorite_tracks:
            # TODO: there's a way to grab the artist info when playing tracks
            # todo: investigate for aliases, description, homepage
            # https://8tracks.com/sets/197579681/next?player=sm&include=track%5Bfaved%2Bannotation%2Bartist_details%5D&mix_id=2183081&track_id=17451458&format=jsonh
            artist = save_artist(favourite_track)
            track = save_track(artist, favourite_track)
            save_recording(artist, favourite_track, track)
            save_album(favourite_track, track)
    # log.info("Finished Scraper Practice")
    print("Finished Scraper Practice")

def get_request_json(user, index):
    web_page_url = "https://8tracks.com/users/" + str(user) + "/favorite_tracks" + \
                  "?page=" + str(index) + "&format=jsonh"
    request = get(web_page_url)
    request_json = request.json()
    return request_json

def save_artist(favourite_track):
    artist = Artist()
    artist.name = favourite_track["performer"]
    # todo: append artist name as alias if different from name found
    try:
        artist.save()
    except (DuplicateKeyError, NotUniqueError):
        artist = Artist.objects(name=artist.name)
    print(artist.name)
    return artist

def save_track(artist, favourite_track):
    artists = [artist]
    track = Track()
    track.name = favourite_track["name"]
    track.artists = artists
    try:
        track.save()
    except (DuplicateKeyError, NotUniqueError):
        track = Track.objects(name=track.name)
    print(track.name)
    return track

def save_recording(artist, favourite_track, track):
    artists = [artist]
    recording = Recording()
    recording.tracks = [track]
    recording.artists = artists
    recording.web_page = favourite_track["url"]
    try:
        print(recording.artists[0].name)
        recording.save()
    except (DuplicateKeyError, NotUniqueError):
        # todo: maybe try to get this proper from the database
        return recording
    return recording

def save_album(favourite_track, track):
    album = Album()
    album.name = favourite_track["release_name"]
    album.artists = track.artists
    album.tracks = [track]
    album.release_date = get_release_date(favourite_track["year"])
    try:
        album.save()
    except (DuplicateKeyError, NotUniqueError):
        # todo: maybe try to get this proper from the database
        return album
    return album

def get_release_date(integer_date):
    if len(str(integer_date)) == 4:
        return datetime.strptime(str(integer_date), "%Y")
    if len(str(integer_date)) == 6:
        return datetime.strptime(str(integer_date), "%Y%m")
    if len(str(integer_date)) == 8:
        return datetime.strptime(str(integer_date), "%Y%m%d")
<|endoftext|>


<commit_before>import urllib.request
import json
import time

from i3pystatus import IntervalModule
from i3pystatus.core.util import internet, require, user_open

class Bitcoin(IntervalModule):

    """
    This module fetches and displays current Bitcoin market prices and
    optionally monitors transactions to and from a list of user-specified
    wallet addresses. Market data is pulled from the BitcoinAverage Price
    Index API <https://bitcoinaverage.com> while transaction data is pulled
    from blockchain.info <https://blockchain.info/api/blockchain_api>.

    Available formatters:

    * {last_price}
    * {ask_price}
    * {bid_price}
    * {daily_average}
    * {volume}
    * {status}
    * {last_tx_type}
    * {last_tx_addr}
    * {last_tx_value}
    * {balance_btc}
    * {balance_fiat}

    """

    settings = (
        ("format", "Format string used for output."),
        ("currency", "Base fiat currency used for pricing."),
        ("wallet_addresses", "List of wallet address(es) to monitor."),
        ("color", "Standard color"),
        ("colorize", "Enable color change on price increase/decrease"),
        ("color_up", "Color for price increases"),
        ("color_down", "Color for price decreases"),
        ("leftclick", "URL to visit or command to run on left click"),
        ("rightclick", "URL to visit or command to run on right click"),
        ("interval", "Update interval."),
        "status"
    )
    format = "฿ {status}{last_price}"
    currency = "USD"
    wallet_addresses = ""
    color = "#FFFFFF"
    colorize = False
    color_up = "#00FF00"
    color_down = "#FF0000"
    leftclick = "electrum"
    rightclick = "https://bitcoinaverage.com/"
    interval = 600
    status = {
        "price_up": "▲",
        "price_down": "▼",
    }

    _price_prev = 0

    def _fetch_price_data(self):
        api = "https://api.bitcoinaverage.com/ticker/global/"
        url = "{}{}".format(api, self.currency.upper())
        return json.loads(urllib.request.urlopen(url).read().decode("utf-8"))

    def _fetch_blockchain_data(self):
        api = "https://blockchain.info/multiaddr?active="
        addresses = "|".join(self.wallet_addresses)
        url = "{}{}".format(api, addresses)
        return json.loads(urllib.request.urlopen(url).read().decode("utf-8"))

    @require(internet)
    def run(self):
        price_data = self._fetch_price_data()
        fdict = {
            "daily_average": price_data["24h_avg"],
            "ask_price": price_data["ask"],
            "bid_price": price_data["bid"],
            "last_price": price_data["last"],
            "volume": price_data["volume_btc"],
        }

        if self._price_prev and fdict["last_price"] > self._price_prev:
            color = self.color_up
            fdict["status"] = self.status["price_up"]
        elif self._price_prev and fdict["last_price"] > self._price_prev:
            color = self.color_down
            fdict["status"] = self.status["price_down"]
        else:
            color = self.color
            fdict["status"] = ""
        self._price_prev = fdict["last_price"]

        if not self.colorize:
            color = self.color

        if self.wallet_addresses:
            blockchain_data = self._fetch_blockchain_data()
            wallet_data = blockchain_data["wallet"]
            balance_btc = wallet_data["final_balance"]/100000000
            fdict["balance_btc"] = round(balance_btc, 2)
            balance_fiat = fdict["balance_btc"] * fdict["last_price"]
            fdict["balance_fiat"] = round(balance_fiat, 2)
            fdict["total_sent"] = wallet_data["total_sent"]
            fdict["total_received"] = wallet_data["total_received"]
            fdict["transactions"] = wallet_data["n_tx"]

            if fdict["transactions"]:
                last_tx = blockchain_data["txs"][0]
                fdict["last_tx_addr"] = last_tx["out"][0]["addr"]
                fdict["last_tx_value"] = last_tx["out"][0]["value"]/100000000
                if fdict["last_tx_addr"] in self.wallet_addresses:
                    fdict["last_tx_type"] = "recv"
                else:
                    fdict["last_tx_type"] = "sent"

        self.output = {
            "full_text": self.format.format(**fdict),
            "color": color,
        }

    def on_leftclick(self):
        user_open(self.leftclick)

    def on_rightclick(self):
        user_open(self.rightclick)
<commit_msg>Fix operator typo in price-comparison<commit_after>import urllib.request
import json
import time

from i3pystatus import IntervalModule
from i3pystatus.core.util import internet, require, user_open

class Bitcoin(IntervalModule):

    """
    This module fetches and displays current Bitcoin market prices and
    optionally monitors transactions to and from a list of user-specified
    wallet addresses. Market data is pulled from the BitcoinAverage Price
    Index API <https://bitcoinaverage.com> while transaction data is pulled
    from blockchain.info <https://blockchain.info/api/blockchain_api>.

    Available formatters:

    * {last_price}
    * {ask_price}
    * {bid_price}
    * {daily_average}
    * {volume}
    * {status}
    * {last_tx_type}
    * {last_tx_addr}
    * {last_tx_value}
    * {balance_btc}
    * {balance_fiat}

    """

    settings = (
        ("format", "Format string used for output."),
        ("currency", "Base fiat currency used for pricing."),
        ("wallet_addresses", "List of wallet address(es) to monitor."),
        ("color", "Standard color"),
        ("colorize", "Enable color change on price increase/decrease"),
        ("color_up", "Color for price increases"),
        ("color_down", "Color for price decreases"),
        ("leftclick", "URL to visit or command to run on left click"),
        ("rightclick", "URL to visit or command to run on right click"),
        ("interval", "Update interval."),
        "status"
    )
    format = "฿ {status}{last_price}"
    currency = "USD"
    wallet_addresses = ""
    color = "#FFFFFF"
    colorize = False
    color_up = "#00FF00"
    color_down = "#FF0000"
    leftclick = "electrum"
    rightclick = "https://bitcoinaverage.com/"
    interval = 600
    status = {
        "price_up": "▲",
        "price_down": "▼",
    }

    _price_prev = 0

    def _fetch_price_data(self):
        api = "https://api.bitcoinaverage.com/ticker/global/"
        url = "{}{}".format(api, self.currency.upper())
        return json.loads(urllib.request.urlopen(url).read().decode("utf-8"))

    def _fetch_blockchain_data(self):
        api = "https://blockchain.info/multiaddr?active="
        addresses = "|".join(self.wallet_addresses)
        url = "{}{}".format(api, addresses)
        return json.loads(urllib.request.urlopen(url).read().decode("utf-8"))

    @require(internet)
    def run(self):
        price_data = self._fetch_price_data()
        fdict = {
            "daily_average": price_data["24h_avg"],
            "ask_price": price_data["ask"],
            "bid_price": price_data["bid"],
            "last_price": price_data["last"],
            "volume": price_data["volume_btc"],
        }

        if self._price_prev and fdict["last_price"] > self._price_prev:
            color = self.color_up
            fdict["status"] = self.status["price_up"]
        elif self._price_prev and fdict["last_price"] < self._price_prev:
            color = self.color_down
            fdict["status"] = self.status["price_down"]
        else:
            color = self.color
            fdict["status"] = ""
        self._price_prev = fdict["last_price"]

        if not self.colorize:
            color = self.color

        if self.wallet_addresses:
            blockchain_data = self._fetch_blockchain_data()
            wallet_data = blockchain_data["wallet"]
            balance_btc = wallet_data["final_balance"]/100000000
            fdict["balance_btc"] = round(balance_btc, 2)
            balance_fiat = fdict["balance_btc"] * fdict["last_price"]
            fdict["balance_fiat"] = round(balance_fiat, 2)
            fdict["total_sent"] = wallet_data["total_sent"]
            fdict["total_received"] = wallet_data["total_received"]
            fdict["transactions"] = wallet_data["n_tx"]

            if fdict["transactions"]:
                last_tx = blockchain_data["txs"][0]
                fdict["last_tx_addr"] = last_tx["out"][0]["addr"]
                fdict["last_tx_value"] = last_tx["out"][0]["value"]/100000000
                if fdict["last_tx_addr"] in self.wallet_addresses:
                    fdict["last_tx_type"] = "recv"
                else:
                    fdict["last_tx_type"] = "sent"

        self.output = {
            "full_text": self.format.format(**fdict),
            "color": color,
        }

    def on_leftclick(self):
        user_open(self.leftclick)

    def on_rightclick(self):
        user_open(self.rightclick)
<|endoftext|>


<commit_before># Copyright 2020 The StackStorm Authors.
# Copyright 2019 Extreme Networks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Monkey patching should be done as early as possible.
# See http://eventlet.net/doc/patching.html#monkeypatching-the-standard-library
from __future__ import absolute_import

from st2common.util.monkey_patch import monkey_patch
monkey_patch()

import os
import signal
import sys

from st2actions.scheduler import config
from st2common import log as logging
from st2common.service_setup import teardown as common_teardown
from st2common.service_setup import setup as common_setup

__all__ = [
    'main'
]

LOG = logging.getLogger(__name__)

def _setup_sigterm_handler():
    def sigterm_handler(signum=None, frame=None):
        # This will cause SystemExit to be throw and allow for component cleanup.
        sys.exit(0)

    # Register a SIGTERM signal handler which calls sys.exit which causes SystemExit to
    # be thrown. We catch SystemExit and handle cleanup there.
    signal.signal(signal.SIGTERM, sigterm_handler)

def _setup():
    capabilities = {
        'name': 'scheduler',
        'type': 'passive'
    }
    common_setup(service='scheduler', config=config, setup_db=True, register_mq_exchanges=True,
                 register_signal_handlers=True, service_registry=True, capabilities=capabilities)

    _setup_sigterm_handler()

def _run_scheduler():
    LOG.info('(PID=%s) Scheduler started.', os.getpid())

    # Lazy load these so that decorator metrics are in place
    from st2actions.scheduler import (
        handler as scheduler_handler,
        entrypoint as scheduler_entrypoint
    )

    handler = scheduler_handler.get_handler()
    entrypoint = scheduler_entrypoint.get_scheduler_entrypoint()

    # TODO: Remove this try block for _cleanup_policy_delayed in v3.2.
    # This is a temporary cleanup to remove executions in deprecated policy-delayed status.
    try:
        handler._cleanup_policy_delayed()
    except Exception:
        LOG.exception('(PID=%s) Scheduler unable to perform migration cleanup.', os.getpid())

    # TODO: Remove this try block for _fix_missing_action_execution_id in v3.2.
    # This is a temporary fix to auto-populate action_execution_id.
    try:
        handler._fix_missing_action_execution_id()
    except Exception:
        LOG.exception('(PID=%s) Scheduler unable to populate action_execution_id.', os.getpid())

    try:
        handler.start()
        entrypoint.start()

        # Wait on handler first since entrypoint is more durable.
        handler.wait() or entrypoint.wait()
    except (KeyboardInterrupt, SystemExit):
        LOG.info('(PID=%s) Scheduler stopped.', os.getpid())

        errors = False

        try:
            handler.shutdown()
            entrypoint.shutdown()
        except:
            LOG.debug('Unable to shutdown scheduler.', exc_info=True)
            errors = True

        if errors:
            return 1
    except:
        LOG.exception('(PID=%s) Scheduler unexpectedly stopped.', os.getpid())

        try:
            handler.shutdown()
            entrypoint.shutdown()
        except:
            LOG.exception('Unable to shutdown scheduler.')

        return 1

    return 0

def _teardown():
    common_teardown()

def main():
    try:
        _setup()
        return _run_scheduler()
    except SystemExit as exit_code:
        sys.exit(exit_code)
    except:
        LOG.exception('(PID=%s) Scheduler quit due to exception.', os.getpid())
        return 1
    finally:
        _teardown()
<commit_msg>Resolved black formatting<commit_after># Copyright 2020 The StackStorm Authors.
# Copyright 2019 Extreme Networks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Monkey patching should be done as early as possible.
# See http://eventlet.net/doc/patching.html#monkeypatching-the-standard-library
from __future__ import absolute_import

from st2common.util.monkey_patch import monkey_patch

monkey_patch()

import os
import signal
import sys

from st2actions.scheduler import config
from st2common import log as logging
from st2common.service_setup import teardown as common_teardown
from st2common.service_setup import setup as common_setup

__all__ = ["main"]

LOG = logging.getLogger(__name__)

def _setup_sigterm_handler():
    def sigterm_handler(signum=None, frame=None):
        # This will cause SystemExit to be throw and allow for component cleanup.
        sys.exit(0)

    # Register a SIGTERM signal handler which calls sys.exit which causes SystemExit to
    # be thrown. We catch SystemExit and handle cleanup there.
    signal.signal(signal.SIGTERM, sigterm_handler)

def _setup():
    capabilities = {"name": "scheduler", "type": "passive"}
    common_setup(
        service="scheduler",
        config=config,
        setup_db=True,
        register_mq_exchanges=True,
        register_signal_handlers=True,
        service_registry=True,
        capabilities=capabilities,
    )

    _setup_sigterm_handler()

def _run_scheduler():
    LOG.info("(PID=%s) Scheduler started.", os.getpid())

    # Lazy load these so that decorator metrics are in place
    from st2actions.scheduler import (
        handler as scheduler_handler,
        entrypoint as scheduler_entrypoint,
    )

    handler = scheduler_handler.get_handler()
    entrypoint = scheduler_entrypoint.get_scheduler_entrypoint()

    # TODO: Remove this try block for _cleanup_policy_delayed in v3.2.
    # This is a temporary cleanup to remove executions in deprecated policy-delayed status.
    try:
        handler._cleanup_policy_delayed()
    except Exception:
        LOG.exception(
            "(PID=%s) Scheduler unable to perform migration cleanup.", os.getpid()
        )

    # TODO: Remove this try block for _fix_missing_action_execution_id in v3.2.
    # This is a temporary fix to auto-populate action_execution_id.
    try:
        handler._fix_missing_action_execution_id()
    except Exception:
        LOG.exception(
            "(PID=%s) Scheduler unable to populate action_execution_id.", os.getpid()
        )

    try:
        handler.start()
        entrypoint.start()

        # Wait on handler first since entrypoint is more durable.
        handler.wait() or entrypoint.wait()
    except (KeyboardInterrupt, SystemExit):
        LOG.info("(PID=%s) Scheduler stopped.", os.getpid())

        errors = False

        try:
            handler.shutdown()
            entrypoint.shutdown()
        except:
            LOG.debug("Unable to shutdown scheduler.", exc_info=True)
            errors = True

        if errors:
            return 1
    except:
        LOG.exception("(PID=%s) Scheduler unexpectedly stopped.", os.getpid())

        try:
            handler.shutdown()
            entrypoint.shutdown()
        except:
            LOG.exception("Unable to shutdown scheduler.")

        return 1

    return 0

def _teardown():
    common_teardown()

def main():
    try:
        _setup()
        return _run_scheduler()
    except SystemExit as exit_code:
        sys.exit(exit_code)
    except:
        LOG.exception("(PID=%s) Scheduler quit due to exception.", os.getpid())
        return 1
    finally:
        _teardown()
<|endoftext|>


<commit_before>from i3pystatus import IntervalModule
from psutil import virtual_memory
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_bar

class MemBar(IntervalModule, ColorRangeModule):
    """
    Shows memory load as a bar.

    .. rubric:: Available formatters

    * {used_mem_bar}

    Requires psutil and colour (from PyPI)
    """

    format = "{used_mem_bar}"
    color = "#00FF00"
    warn_color = "#FFFF00"
    alert_color = "#FF0000"
    warn_percentage = 50
    alert_percentage = 80
    multi_colors = False

    def init(self):
        self.colors = self.get_hex_color_range(self.color, self.alert_color, 100)

    settings = (
        ("format", "format string used for output."),
        ("warn_percentage", "minimal percentage for warn state"),
        ("alert_percentage", "minimal percentage for alert state"),
        ("color", "standard color"),
        ("warn_color",
         "defines the color used when warn percentage is exceeded"),
        ("alert_color",
         "defines the color used when alert percentage is exceeded"),
        ("multi_colors", "whether to use range of colors from 'color' to 'alert_color' based on memory usage."),
    )

    def run(self):
        memory_usage = virtual_memory()

        if self.multi_colors:
            color = self.get_gradient(memory_usage.percent, self.colors)
        elif memory_usage.percent >= self.alert_percentage:
            color = self.alert_color
        elif memory_usage.percent >= self.warn_percentage:
            color = self.warn_color
        else:
            color = self.color

        self.output = {
            "full_text": self.format.format(
                used_mem_bar=make_bar(memory_usage.percent)),
            "color": color
        }
<commit_msg>Add short_text for mem_bar<commit_after>from i3pystatus import IntervalModule
from psutil import virtual_memory
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_bar

class MemBar(IntervalModule, ColorRangeModule):
    """
    Shows memory load as a bar.

    .. rubric:: Available formatters

    * {used_mem_bar}

    Requires psutil and colour (from PyPI)
    """

    format = "{used_mem_bar}"
    color = "#00FF00"
    warn_color = "#FFFF00"
    alert_color = "#FF0000"
    warn_percentage = 50
    alert_percentage = 80
    multi_colors = False

    def init(self):
        self.colors = self.get_hex_color_range(self.color, self.alert_color, 100)

    settings = (
        ("format", "format string used for output."),
        ("warn_percentage", "minimal percentage for warn state"),
        ("alert_percentage", "minimal percentage for alert state"),
        ("color", "standard color"),
        ("warn_color",
         "defines the color used when warn percentage is exceeded"),
        ("alert_color",
         "defines the color used when alert percentage is exceeded"),
        ("multi_colors", "whether to use range of colors from 'color' to 'alert_color' based on memory usage."),
    )

    def run(self):
        memory_usage = virtual_memory()

        if self.multi_colors:
            color = self.get_gradient(memory_usage.percent, self.colors)
        elif memory_usage.percent >= self.alert_percentage:
            color = self.alert_color
        elif memory_usage.percent >= self.warn_percentage:
            color = self.warn_color
        else:
            color = self.color

        self.output = {
            "full_text": self.format.format(
                used_mem_bar=make_bar(memory_usage.percent)),
            "short_text": "",
            "color": color
        }
<|endoftext|>


<commit_before># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

import mock
import six
from st2common.persistence.action import Action, RunnerType
from st2common.persistence.reactor import Trigger
from st2common.models.db import action, reactor
from st2common.transport.publishers import PoolPublisher
from st2tests.fixturesloader import FixturesLoader
from tests import FunctionalTest

http_client = six.moves.http_client

TEST_FIXTURES = {
    'runners': ['testrunner1.json'],
    'actions': ['action1.json'],
    'triggers': ['trigger1.json'],
    'triggertypes': ['triggertype1.json']
}

FIXTURES_PACK = 'generic'

@mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
class TestRuleController(FunctionalTest):

    fixtures_loader = FixturesLoader()

    @classmethod
    def setUpClass(cls):
        super(TestRuleController, cls).setUpClass()
        models = TestRuleController.fixtures_loader.save_fixtures_to_db(
            fixtures_pack=FIXTURES_PACK, fixtures_dict=TEST_FIXTURES)
        TestRuleController.RUNNER_TYPE = models['runners']['testrunner1.json']
        TestRuleController.ACTION = models['actions']['action1.json']
        TestRuleController.TRIGGER = models['triggers']['trigger1.json']

        # Don't load rule into DB as that is what is being tested.
        TestRuleController.RULE_1 = TestRuleController.fixtures_loader.load_fixtures(
            fixtures_pack=FIXTURES_PACK,
            fixtures_dict={'rules': ['rule1.json']})['rules']['rule1.json']

    @classmethod
    def tearDownClass(cls):
        TestRuleController.fixtures_loader.delete_fixtures_from_db(
            fixtures_pack=FIXTURES_PACK, fixtures_dict=TEST_FIXTURES)
        super(TestRuleController, cls).setUpClass()

    def test_get_all(self):
        resp = self.app.get('/rules')
        self.assertEqual(resp.status_int, http_client.OK)

    def test_get_one(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        rule_id = self.__get_rule_id(post_resp)
        get_resp = self.__do_get_one(rule_id)
        self.assertEqual(get_resp.status_int, http_client.OK)
        self.assertEqual(self.__get_rule_id(get_resp), rule_id)
        self.__do_delete(rule_id)

    def test_get_one_fail(self):
        resp = self.app.get('/rules/1', expect_errors=True)
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    def test_post(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp.status_int, http_client.CREATED)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_post_duplicate(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp.status_int, http_client.CREATED)
        post_resp_2 = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp_2.status_int, http_client.CONFLICT)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_put(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        update_input = post_resp.json
        update_input['enabled'] = not update_input['enabled']
        put_resp = self.__do_put(self.__get_rule_id(post_resp), update_input)
        self.assertEqual(put_resp.status_int, http_client.OK)
        self.__do_delete(self.__get_rule_id(put_resp))

    def test_put_fail(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        update_input = post_resp.json
        # If the id in the URL is incorrect the update will fail since id in the body is ignored.
        put_resp = self.__do_put(1, update_input)
        self.assertEqual(put_resp.status_int, http_client.NOT_FOUND)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_delete(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        del_resp = self.__do_delete(self.__get_rule_id(post_resp))
        self.assertEqual(del_resp.status_int, http_client.NO_CONTENT)

    @staticmethod
    def __get_rule_id(resp):
        return resp.json['id']

    def __do_get_one(self, rule_id):
        return self.app.get('/rules/%s' % rule_id, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_post(self, rule):
        return self.app.post_json('/rules', rule, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_put(self, rule_id, rule):
        return self.app.put_json('/rules/%s' % rule_id, rule, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_delete(self, rule_id):
        return self.app.delete('/rules/%s' % rule_id)
<commit_msg>Remove unused imports.<commit_after># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

import mock
import six
from st2common.transport.publishers import PoolPublisher
from st2tests.fixturesloader import FixturesLoader
from tests import FunctionalTest

http_client = six.moves.http_client

TEST_FIXTURES = {
    'runners': ['testrunner1.json'],
    'actions': ['action1.json'],
    'triggers': ['trigger1.json'],
    'triggertypes': ['triggertype1.json']
}

FIXTURES_PACK = 'generic'

@mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
class TestRuleController(FunctionalTest):

    fixtures_loader = FixturesLoader()

    @classmethod
    def setUpClass(cls):
        super(TestRuleController, cls).setUpClass()
        models = TestRuleController.fixtures_loader.save_fixtures_to_db(
            fixtures_pack=FIXTURES_PACK, fixtures_dict=TEST_FIXTURES)
        TestRuleController.RUNNER_TYPE = models['runners']['testrunner1.json']
        TestRuleController.ACTION = models['actions']['action1.json']
        TestRuleController.TRIGGER = models['triggers']['trigger1.json']

        # Don't load rule into DB as that is what is being tested.
        TestRuleController.RULE_1 = TestRuleController.fixtures_loader.load_fixtures(
            fixtures_pack=FIXTURES_PACK,
            fixtures_dict={'rules': ['rule1.json']})['rules']['rule1.json']

    @classmethod
    def tearDownClass(cls):
        TestRuleController.fixtures_loader.delete_fixtures_from_db(
            fixtures_pack=FIXTURES_PACK, fixtures_dict=TEST_FIXTURES)
        super(TestRuleController, cls).setUpClass()

    def test_get_all(self):
        resp = self.app.get('/rules')
        self.assertEqual(resp.status_int, http_client.OK)

    def test_get_one(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        rule_id = self.__get_rule_id(post_resp)
        get_resp = self.__do_get_one(rule_id)
        self.assertEqual(get_resp.status_int, http_client.OK)
        self.assertEqual(self.__get_rule_id(get_resp), rule_id)
        self.__do_delete(rule_id)

    def test_get_one_fail(self):
        resp = self.app.get('/rules/1', expect_errors=True)
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    def test_post(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp.status_int, http_client.CREATED)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_post_duplicate(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp.status_int, http_client.CREATED)
        post_resp_2 = self.__do_post(TestRuleController.RULE_1)
        self.assertEqual(post_resp_2.status_int, http_client.CONFLICT)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_put(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        update_input = post_resp.json
        update_input['enabled'] = not update_input['enabled']
        put_resp = self.__do_put(self.__get_rule_id(post_resp), update_input)
        self.assertEqual(put_resp.status_int, http_client.OK)
        self.__do_delete(self.__get_rule_id(put_resp))

    def test_put_fail(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        update_input = post_resp.json
        # If the id in the URL is incorrect the update will fail since id in the body is ignored.
        put_resp = self.__do_put(1, update_input)
        self.assertEqual(put_resp.status_int, http_client.NOT_FOUND)
        self.__do_delete(self.__get_rule_id(post_resp))

    def test_delete(self):
        post_resp = self.__do_post(TestRuleController.RULE_1)
        del_resp = self.__do_delete(self.__get_rule_id(post_resp))
        self.assertEqual(del_resp.status_int, http_client.NO_CONTENT)

    @staticmethod
    def __get_rule_id(resp):
        return resp.json['id']

    def __do_get_one(self, rule_id):
        return self.app.get('/rules/%s' % rule_id, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_post(self, rule):
        return self.app.post_json('/rules', rule, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_put(self, rule_id, rule):
        return self.app.put_json('/rules/%s' % rule_id, rule, expect_errors=True)

    @mock.patch.object(PoolPublisher, 'publish', mock.MagicMock())
    def __do_delete(self, rule_id):
        return self.app.delete('/rules/%s' % rule_id)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import netifaces
from i3pystatus import IntervalModule
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_graph, round_dict, make_bar

def count_bits(integer):
    bits = 0
    while (integer):
        integer &= integer - 1
        bits += 1
    return bits

def v6_to_int(v6):
    return int(v6.replace(":", ""), 16)

def prefix6(mask):
    return count_bits(v6_to_int(mask))

def cidr6(addr, mask):
    return "{addr}/{bits}".format(addr=addr, bits=prefix6(mask))

def v4_to_int(v4):
    sum = 0
    mul = 1
    for part in reversed(v4.split(".")):
        sum += int(part) * mul
        mul *= 2 ** 8
    return sum

def prefix4(mask):
    return count_bits(v4_to_int(mask))

def cidr4(addr, mask):
    return "{addr}/{bits}".format(addr=addr, bits=prefix4(mask))

def get_bonded_slaves():
    try:
        with open("/sys/class/net/bonding_masters") as f:
            masters = f.read().split()
    except FileNotFoundError:
        return {}
    slaves = {}
    for master in masters:
        with open("/sys/class/net/{}/bonding/slaves".format(master)) as f:
            for slave in f.read().split():
                slaves[slave] = master
    return slaves

def sysfs_interface_up(interface, unknown_up=False):
    try:
        with open("/sys/class/net/{}/operstate".format(interface)) as f:
            status = f.read().strip()
    except FileNotFoundError:
        # Interface doesn't exist
        return False

    return status == "up" or unknown_up and status == "unknown"

class NetworkInfo():
    """
    Retrieve network information.
    """

    def __init__(self, interface, ignore_interfaces, detached_down, unknown_up, get_wifi_info=False):
        if interface not in netifaces.interfaces() and not detached_down:
            raise RuntimeError(
                "Unknown interface {iface}!".format(iface=interface))

        self.ignore_interfaces = ignore_interfaces
        self.detached_down = detached_down
        self.unknown_up = unknown_up
        self.get_wifi_info = get_wifi_info

    def get_info(self, interface):
        format_dict = dict(v4="", v4mask="", v4cidr="", v6="", v6mask="", v6cidr="")
        iface_up = sysfs_interface_up(interface, self.unknown_up)
        if not iface_up:
            return format_dict

        network_info = netifaces.ifaddresses(interface)
        slaves = get_bonded_slaves()
        try:
            master = slaves[interface]
        except KeyError:
            pass
        else:
            if sysfs_interface_up(interface, self.unknown_up):
                master_info = netifaces.ifaddresses(master)
                for af in (netifaces.AF_INET, netifaces.AF_INET6):
                    try:
                        network_info[af] = master_info[af]
                    except KeyError:
                        pass

        try:
            mac = network_info[netifaces.AF_PACKET][0]["addr"]
        except KeyError:
            mac = "NONE"
        format_dict['mac'] = mac

        if iface_up:
            format_dict.update(self.extract_network_info(network_info))
            format_dict.update(self.extract_wireless_info(interface))

        return format_dict

    @staticmethod
    def extract_network_info(network_info):
        info = dict()
        if netifaces.AF_INET in network_info:
            v4 = network_info[netifaces.AF_INET][0]
            info["v4"] = v4["addr"]
            info["v4mask"] = v4["netmask"]
            info["v4cidr"] = cidr4(v4["addr"], v4["netmask"])
        if netifaces.AF_INET6 in network_info:
            for v6 in network_info[netifaces.AF_INET6]:
                info["v6"] = v6["addr"]
                info["v6mask"] = v6["netmask"]
                info["v6cidr"] = cidr6(v6["addr"], v6["netmask"])
                if not v6["addr"].startswith("fe80::"):  # prefer non link-local addresses
                    break
        return info

    def extract_wireless_info(self, interface):
        info = dict(essid="", freq="", quality=0.0, quality_bar="")

        # Just return empty values if we're not using any Wifi functionality
        if not self.get_wifi_info:
            return info

        import basiciw

        try:
            iwi = basiciw.iwinfo(interface)
        except Exception:
            # Not a wireless interface
            return info

        info["essid"] = iwi["essid"]
        info["freq"] = iwi["freq"]
        quality = iwi["quality"]
        if quality["quality_max"] > 0:
            info["quality"] = quality["quality"] / quality["quality_max"]
        else:
            info["quality"] = quality["quality"]
        info["quality"] *= 100
        info["quality_bar"] = make_bar(info["quality"])
        info["quality"]='{0:.{1}f}'.format(info["quality"],0)

        return info

class NetworkTraffic():
    """
    Retrieve network traffic information
    """

    pnic = None
    pnic_before = None

    def __init__(self, unknown_up, divisor, round_size):
        self.unknown_up = unknown_up
        self.divisor = divisor
        self.round_size = round_size

    def update_counters(self, interface):
        import psutil

        self.pnic_before = self.pnic
        counters = psutil.net_io_counters(pernic=True)
        self.pnic = counters[interface] if interface in counters else None

    def clear_counters(self):
        self.pnic_before = None
        self.pnic = None

    def get_bytes_sent(self):
        return (self.pnic.bytes_sent - self.pnic_before.bytes_sent) / self.divisor

    def get_bytes_received(self):
        return (self.pnic.bytes_recv - self.pnic_before.bytes_recv) / self.divisor

    def get_packets_sent(self):
        return self.pnic.packets_sent - self.pnic_before.packets_sent

    def get_packets_received(self):
        return self.pnic.packets_recv - self.pnic_before.packets_recv

    def get_usage(self, interface):
        self.update_counters(interface)
        usage = dict(bytes_sent=0, bytes_recv=0, packets_sent=0, packets_recv=0)

        if not sysfs_interface_up(interface, self.unknown_up) or not self.pnic_before:
            return usage
        else:
            usage["bytes_sent"] = self.get_bytes_sent()
            usage["bytes_recv"] = self.get_bytes_received()
            usage["packets_sent"] = self.get_packets_sent()
            usage["packets_recv"] = self.get_packets_received()
            round_dict(usage, self.round_size)
        return usage

class Network(IntervalModule, ColorRangeModule):
    """
    Displays network information for an interface.

    Requires the PyPI packages `psutil`, `colour`, `netifaces` and `basiciw`

    .. rubric:: Available formatters

    Network Traffic Formatters:
    * `{interface}` — the configured network interface
    * `{kbs}` – Float representing kb\s
    * `{network_graph}` – Unicode graph representing network usage
    * `{bytes_sent}` — bytes sent per second (divided by divisor)
    * `{bytes_recv}` — bytes received per second (divided by divisor)
    * `{packets_sent}` — bytes sent per second (divided by divisor)
    * `{packets_recv}` — bytes received per second (divided by divisor)

    Network Information Formatters:
    * `{interface}` — same as setting
    * `{v4}` — IPv4 address
    * `{v4mask}` — subnet mask
    * `{v4cidr}` — IPv4 address in cidr notation (i.e. 192.168.2.204/24)
    * `{v6}` — IPv6 address
    * `{v6mask}` — subnet mask
    * `{v6cidr}` — IPv6 address in cidr notation
    * `{mac}` — MAC of interface

    Wireless Information Formatters:
    * `{essid}` — ESSID of currently connected wifi
    * `{freq}` — Current frequency
    * `{quality}` — Link quality in percent
    * `{quality_bar}` —Bar graphically representing link quality
    """
    settings = (
        ("format_up", "format string"),
        ("format_down", "format string"),
        "color_up",
        "color_down",
        ("interface", "Interface to watch, eg 'eth0'"),
        ("dynamic_color", "Set color dynamically based on network traffic. Note: this overrides color_up"),
        ("start_color", "Hex or English name for start of color range, eg '#00FF00' or 'green'"),
        ("end_color", "Hex or English name for end of color range, eg '#FF0000' or 'red'"),
        ("graph_width", "Width of the network traffic graph"),
        ("graph_style", "Graph style ('blocks', 'braille-fill', 'braille-peak', or 'braille-snake')"),
        ("upper_limit",
         "Expected max kb/s. This value controls how the network traffic graph is drawn and in what color"),
        ("graph_type", "Whether to draw the network traffic graph for input or output. "
                       "Allowed values 'input' or 'output'"),
        ("divisor", "divide all byte values by this value"),
        ("ignore_interfaces", "Array of interfaces to ignore when cycling through "
                              "on click, eg, ['lo']"),
        ("round_size", "defines number of digits in round"),
        ("detached_down", "If the interface doesn't exist, display it as if it were down"),
        ("unknown_up", "If the interface is in unknown state, display it as if it were up"),
    )

    interval = 1
    interface = 'eth0'

    format_up = "{interface} {network_graph}{kbs}KB/s"
    format_down = "{interface}: DOWN"
    color_up = "#00FF00"
    color_down = "#FF0000"
    dynamic_color = True
    graph_type = 'input'
    graph_width = 15
    graph_style = 'blocks'
    upper_limit = 150.0

    # Network traffic settings
    divisor = 1024
    round_size = None

    # Network info settings
    detached_down = True
    unknown_up = False
    ignore_interfaces = ["lo"]

    on_leftclick = "nm-connection-editor"
    on_rightclick = "cycle_interface"
    on_upscroll = ['cycle_interface', 1]
    on_downscroll = ['cycle_interface', -1]

    def init(self):
        # Don't require importing basiciw unless using the functionality it offers.
        if any(s in self.format_up or s in self.format_up for s in
               ['essid', 'freq', 'quality', 'quality_bar']):
            get_wifi_info = True
        else:
            get_wifi_info = False

        self.network_info = NetworkInfo(self.interface, self.ignore_interfaces, self.detached_down, self.unknown_up,
                                        get_wifi_info)

        # Don't require importing psutil unless using the functionality it offers.
        if any(s in self.format_up or s in self.format_down for s in
               ['bytes_sent', 'bytes_recv', 'packets_sent', 'packets_recv', 'network_graph']):
            self.network_traffic = NetworkTraffic(self.unknown_up, self.divisor, self.round_size)
        else:
            self.network_traffic = None

        self.colors = self.get_hex_color_range(self.start_color, self.end_color, int(self.upper_limit))
        self.kbs_arr = [0.0] * self.graph_width

    def cycle_interface(self, increment=1):
        interfaces = [i for i in netifaces.interfaces() if i not in self.ignore_interfaces]
        if self.interface in interfaces:
            next_index = (interfaces.index(self.interface) + increment) % len(interfaces)
            self.interface = interfaces[next_index]
        elif len(interfaces) > 0:
            self.interface = interfaces[0]

        if self.network_traffic:
            self.network_traffic.clear_counters()
            self.kbs_arr = [0.0] * self.graph_width

    def get_network_graph(self, kbs):
        # Cycle array by inserting at the start and chopping off the last element
        self.kbs_arr.insert(0, kbs)
        self.kbs_arr = self.kbs_arr[:self.graph_width]
        return make_graph(self.kbs_arr, 0.0, self.upper_limit, self.graph_style)

    def run(self):
        format_values = dict(kbs="", network_graph="", bytes_sent="", bytes_recv="", packets_sent="", packets_recv="",
                             interface="", v4="", v4mask="", v4cidr="", v6="", v6mask="", v6cidr="", mac="",
                             essid="", freq="", quality="", quality_bar="")
        color = None
        if self.network_traffic:
            network_usage = self.network_traffic.get_usage(self.interface)
            format_values.update(network_usage)
            if self.graph_type == 'input':
                kbs = network_usage['bytes_recv']
            elif self.graph_type == 'output':
                kbs = network_usage['bytes_sent']
            else:
                raise Exception("graph_type must be either 'input' or 'output'!")

            format_values['network_graph'] = self.get_network_graph(kbs)
            format_values['kbs'] = "{0:.1f}".format(round(kbs, 2)).rjust(6)
            color = self.get_gradient(kbs, self.colors, self.upper_limit)

        network_info = self.network_info.get_info(self.interface)
        format_values.update(network_info)

        format_values['interface'] = self.interface
        if sysfs_interface_up(self.interface, self.unknown_up):
            if not self.dynamic_color:
                color = self.color_up

            self.output = {
                "full_text": self.format_up.format(**format_values),
                'color': color,
            }
        else:
            color = self.color_down
            self.output = {
                "full_text": self.format_down.format(**format_values),
                'color': color,
            }
<commit_msg>format wifi percentage<commit_after># -*- coding: utf-8 -*-
import netifaces
from i3pystatus import IntervalModule
from i3pystatus.core.color import ColorRangeModule
from i3pystatus.core.util import make_graph, round_dict, make_bar

def count_bits(integer):
    bits = 0
    while (integer):
        integer &= integer - 1
        bits += 1
    return bits

def v6_to_int(v6):
    return int(v6.replace(":", ""), 16)

def prefix6(mask):
    return count_bits(v6_to_int(mask))

def cidr6(addr, mask):
    return "{addr}/{bits}".format(addr=addr, bits=prefix6(mask))

def v4_to_int(v4):
    sum = 0
    mul = 1
    for part in reversed(v4.split(".")):
        sum += int(part) * mul
        mul *= 2 ** 8
    return sum

def prefix4(mask):
    return count_bits(v4_to_int(mask))

def cidr4(addr, mask):
    return "{addr}/{bits}".format(addr=addr, bits=prefix4(mask))

def get_bonded_slaves():
    try:
        with open("/sys/class/net/bonding_masters") as f:
            masters = f.read().split()
    except FileNotFoundError:
        return {}
    slaves = {}
    for master in masters:
        with open("/sys/class/net/{}/bonding/slaves".format(master)) as f:
            for slave in f.read().split():
                slaves[slave] = master
    return slaves

def sysfs_interface_up(interface, unknown_up=False):
    try:
        with open("/sys/class/net/{}/operstate".format(interface)) as f:
            status = f.read().strip()
    except FileNotFoundError:
        # Interface doesn't exist
        return False

    return status == "up" or unknown_up and status == "unknown"

class NetworkInfo():
    """
    Retrieve network information.
    """

    def __init__(self, interface, ignore_interfaces, detached_down, unknown_up, get_wifi_info=False):
        if interface not in netifaces.interfaces() and not detached_down:
            raise RuntimeError(
                "Unknown interface {iface}!".format(iface=interface))

        self.ignore_interfaces = ignore_interfaces
        self.detached_down = detached_down
        self.unknown_up = unknown_up
        self.get_wifi_info = get_wifi_info

    def get_info(self, interface):
        format_dict = dict(v4="", v4mask="", v4cidr="", v6="", v6mask="", v6cidr="")
        iface_up = sysfs_interface_up(interface, self.unknown_up)
        if not iface_up:
            return format_dict

        network_info = netifaces.ifaddresses(interface)
        slaves = get_bonded_slaves()
        try:
            master = slaves[interface]
        except KeyError:
            pass
        else:
            if sysfs_interface_up(interface, self.unknown_up):
                master_info = netifaces.ifaddresses(master)
                for af in (netifaces.AF_INET, netifaces.AF_INET6):
                    try:
                        network_info[af] = master_info[af]
                    except KeyError:
                        pass

        try:
            mac = network_info[netifaces.AF_PACKET][0]["addr"]
        except KeyError:
            mac = "NONE"
        format_dict['mac'] = mac

        if iface_up:
            format_dict.update(self.extract_network_info(network_info))
            format_dict.update(self.extract_wireless_info(interface))

        return format_dict

    @staticmethod
    def extract_network_info(network_info):
        info = dict()
        if netifaces.AF_INET in network_info:
            v4 = network_info[netifaces.AF_INET][0]
            info["v4"] = v4["addr"]
            info["v4mask"] = v4["netmask"]
            info["v4cidr"] = cidr4(v4["addr"], v4["netmask"])
        if netifaces.AF_INET6 in network_info:
            for v6 in network_info[netifaces.AF_INET6]:
                info["v6"] = v6["addr"]
                info["v6mask"] = v6["netmask"]
                info["v6cidr"] = cidr6(v6["addr"], v6["netmask"])
                if not v6["addr"].startswith("fe80::"):  # prefer non link-local addresses
                    break
        return info

    def extract_wireless_info(self, interface):
        info = dict(essid="", freq="", quality=0.0, quality_bar="")

        # Just return empty values if we're not using any Wifi functionality
        if not self.get_wifi_info:
            return info

        import basiciw

        try:
            iwi = basiciw.iwinfo(interface)
        except Exception:
            # Not a wireless interface
            return info

        info["essid"] = iwi["essid"]
        info["freq"] = iwi["freq"]
        quality = iwi["quality"]
        if quality["quality_max"] > 0:
            info["quality"] = quality["quality"] / quality["quality_max"]
        else:
            info["quality"] = quality["quality"]
        info["quality"] *= 100
        info["quality_bar"] = make_bar(info["quality"])
        info["quality"]='{:.0f}'.format(info["quality"])

        return info

class NetworkTraffic():
    """
    Retrieve network traffic information
    """

    pnic = None
    pnic_before = None

    def __init__(self, unknown_up, divisor, round_size):
        self.unknown_up = unknown_up
        self.divisor = divisor
        self.round_size = round_size

    def update_counters(self, interface):
        import psutil

        self.pnic_before = self.pnic
        counters = psutil.net_io_counters(pernic=True)
        self.pnic = counters[interface] if interface in counters else None

    def clear_counters(self):
        self.pnic_before = None
        self.pnic = None

    def get_bytes_sent(self):
        return (self.pnic.bytes_sent - self.pnic_before.bytes_sent) / self.divisor

    def get_bytes_received(self):
        return (self.pnic.bytes_recv - self.pnic_before.bytes_recv) / self.divisor

    def get_packets_sent(self):
        return self.pnic.packets_sent - self.pnic_before.packets_sent

    def get_packets_received(self):
        return self.pnic.packets_recv - self.pnic_before.packets_recv

    def get_usage(self, interface):
        self.update_counters(interface)
        usage = dict(bytes_sent=0, bytes_recv=0, packets_sent=0, packets_recv=0)

        if not sysfs_interface_up(interface, self.unknown_up) or not self.pnic_before:
            return usage
        else:
            usage["bytes_sent"] = self.get_bytes_sent()
            usage["bytes_recv"] = self.get_bytes_received()
            usage["packets_sent"] = self.get_packets_sent()
            usage["packets_recv"] = self.get_packets_received()
            round_dict(usage, self.round_size)
        return usage

class Network(IntervalModule, ColorRangeModule):
    """
    Displays network information for an interface.

    Requires the PyPI packages `psutil`, `colour`, `netifaces` and `basiciw`

    .. rubric:: Available formatters

    Network Traffic Formatters:
    * `{interface}` — the configured network interface
    * `{kbs}` – Float representing kb\s
    * `{network_graph}` – Unicode graph representing network usage
    * `{bytes_sent}` — bytes sent per second (divided by divisor)
    * `{bytes_recv}` — bytes received per second (divided by divisor)
    * `{packets_sent}` — bytes sent per second (divided by divisor)
    * `{packets_recv}` — bytes received per second (divided by divisor)

    Network Information Formatters:
    * `{interface}` — same as setting
    * `{v4}` — IPv4 address
    * `{v4mask}` — subnet mask
    * `{v4cidr}` — IPv4 address in cidr notation (i.e. 192.168.2.204/24)
    * `{v6}` — IPv6 address
    * `{v6mask}` — subnet mask
    * `{v6cidr}` — IPv6 address in cidr notation
    * `{mac}` — MAC of interface

    Wireless Information Formatters:
    * `{essid}` — ESSID of currently connected wifi
    * `{freq}` — Current frequency
    * `{quality}` — Link quality in percent
    * `{quality_bar}` —Bar graphically representing link quality
    """
    settings = (
        ("format_up", "format string"),
        ("format_down", "format string"),
        "color_up",
        "color_down",
        ("interface", "Interface to watch, eg 'eth0'"),
        ("dynamic_color", "Set color dynamically based on network traffic. Note: this overrides color_up"),
        ("start_color", "Hex or English name for start of color range, eg '#00FF00' or 'green'"),
        ("end_color", "Hex or English name for end of color range, eg '#FF0000' or 'red'"),
        ("graph_width", "Width of the network traffic graph"),
        ("graph_style", "Graph style ('blocks', 'braille-fill', 'braille-peak', or 'braille-snake')"),
        ("upper_limit",
         "Expected max kb/s. This value controls how the network traffic graph is drawn and in what color"),
        ("graph_type", "Whether to draw the network traffic graph for input or output. "
                       "Allowed values 'input' or 'output'"),
        ("divisor", "divide all byte values by this value"),
        ("ignore_interfaces", "Array of interfaces to ignore when cycling through "
                              "on click, eg, ['lo']"),
        ("round_size", "defines number of digits in round"),
        ("detached_down", "If the interface doesn't exist, display it as if it were down"),
        ("unknown_up", "If the interface is in unknown state, display it as if it were up"),
    )

    interval = 1
    interface = 'eth0'

    format_up = "{interface} {network_graph}{kbs}KB/s"
    format_down = "{interface}: DOWN"
    color_up = "#00FF00"
    color_down = "#FF0000"
    dynamic_color = True
    graph_type = 'input'
    graph_width = 15
    graph_style = 'blocks'
    upper_limit = 150.0

    # Network traffic settings
    divisor = 1024
    round_size = None

    # Network info settings
    detached_down = True
    unknown_up = False
    ignore_interfaces = ["lo"]

    on_leftclick = "nm-connection-editor"
    on_rightclick = "cycle_interface"
    on_upscroll = ['cycle_interface', 1]
    on_downscroll = ['cycle_interface', -1]

    def init(self):
        # Don't require importing basiciw unless using the functionality it offers.
        if any(s in self.format_up or s in self.format_up for s in
               ['essid', 'freq', 'quality', 'quality_bar']):
            get_wifi_info = True
        else:
            get_wifi_info = False

        self.network_info = NetworkInfo(self.interface, self.ignore_interfaces, self.detached_down, self.unknown_up,
                                        get_wifi_info)

        # Don't require importing psutil unless using the functionality it offers.
        if any(s in self.format_up or s in self.format_down for s in
               ['bytes_sent', 'bytes_recv', 'packets_sent', 'packets_recv', 'network_graph']):
            self.network_traffic = NetworkTraffic(self.unknown_up, self.divisor, self.round_size)
        else:
            self.network_traffic = None

        self.colors = self.get_hex_color_range(self.start_color, self.end_color, int(self.upper_limit))
        self.kbs_arr = [0.0] * self.graph_width

    def cycle_interface(self, increment=1):
        interfaces = [i for i in netifaces.interfaces() if i not in self.ignore_interfaces]
        if self.interface in interfaces:
            next_index = (interfaces.index(self.interface) + increment) % len(interfaces)
            self.interface = interfaces[next_index]
        elif len(interfaces) > 0:
            self.interface = interfaces[0]

        if self.network_traffic:
            self.network_traffic.clear_counters()
            self.kbs_arr = [0.0] * self.graph_width

    def get_network_graph(self, kbs):
        # Cycle array by inserting at the start and chopping off the last element
        self.kbs_arr.insert(0, kbs)
        self.kbs_arr = self.kbs_arr[:self.graph_width]
        return make_graph(self.kbs_arr, 0.0, self.upper_limit, self.graph_style)

    def run(self):
        format_values = dict(kbs="", network_graph="", bytes_sent="", bytes_recv="", packets_sent="", packets_recv="",
                             interface="", v4="", v4mask="", v4cidr="", v6="", v6mask="", v6cidr="", mac="",
                             essid="", freq="", quality="", quality_bar="")
        color = None
        if self.network_traffic:
            network_usage = self.network_traffic.get_usage(self.interface)
            format_values.update(network_usage)
            if self.graph_type == 'input':
                kbs = network_usage['bytes_recv']
            elif self.graph_type == 'output':
                kbs = network_usage['bytes_sent']
            else:
                raise Exception("graph_type must be either 'input' or 'output'!")

            format_values['network_graph'] = self.get_network_graph(kbs)
            format_values['kbs'] = "{0:.1f}".format(round(kbs, 2)).rjust(6)
            color = self.get_gradient(kbs, self.colors, self.upper_limit)

        network_info = self.network_info.get_info(self.interface)
        format_values.update(network_info)

        format_values['interface'] = self.interface
        if sysfs_interface_up(self.interface, self.unknown_up):
            if not self.dynamic_color:
                color = self.color_up

            self.output = {
                "full_text": self.format_up.format(**format_values),
                'color': color,
            }
        else:
            color = self.color_down
            self.output = {
                "full_text": self.format_down.format(**format_values),
                'color': color,
            }
<|endoftext|>