<commit_before>import os

class Corpus(object):

    def __init__(self):
        current_directory = os.path.dirname(__file__)
        self.data_directory = os.path.join(current_directory, 'data')

    def get_file_path(self, dotted_path):
        """
        Reads a dotted file path and returns the file path.
        """
        parts = dotted_path.split(".")
        if parts[0] == 'chatterbot':
            parts.pop(0)
            parts[0] = self.data_directory

        corpus_path = os.path.join(*parts)

        if os.path.exists(corpus_path + ".json"):
            corpus_path += ".json"

        return corpus_path

    def read_corpus(self, file_name):
        """
        Read and return the data from a corpus json file.
        """
        import json

        with open(file_name) as data_file:
            data = json.load(data_file)
        return data

    def load_corpus(self, dotted_path):
        """
        Return the data contained within a specified corpus.
        """

        corpus_path = self.get_file_path(dotted_path)

        corpora = []

        if os.path.isdir(corpus_path):
            for dirname, dirnames, filenames in os.walk(corpus_path):
                for datafile in filenames:
                    if datafile.endswith(".json"):

                        corpus = self.read_corpus(
                            os.path.join(dirname, datafile)
                        )

                        for key in list(corpus.keys()):
                            corpora.append(corpus[key])
        else:
            corpus = self.read_corpus(corpus_path)

            for key in list(corpus.keys()):
                corpora.append(corpus[key])

        return corpora
<commit_msg>Added UTF-8 support in corpus<commit_after>import os

class Corpus(object):

    def __init__(self):
        current_directory = os.path.dirname(__file__)
        self.data_directory = os.path.join(current_directory, 'data')

    def get_file_path(self, dotted_path):
        """
        Reads a dotted file path and returns the file path.
        """
        parts = dotted_path.split(".")
        if parts[0] == 'chatterbot':
            parts.pop(0)
            parts[0] = self.data_directory

        corpus_path = os.path.join(*parts)

        if os.path.exists(corpus_path + ".json"):
            corpus_path += ".json"

        return corpus_path

    def read_corpus(self, file_name):
        """
        Read and return the data from a corpus json file.
        """
        import json

        with open(file_name, encoding='utf-8') as data_file:
            data = json.load(data_file)
        return data

    def load_corpus(self, dotted_path):
        """
        Return the data contained within a specified corpus.
        """

        corpus_path = self.get_file_path(dotted_path)

        corpora = []

        if os.path.isdir(corpus_path):
            for dirname, dirnames, filenames in os.walk(corpus_path):
                for datafile in filenames:
                    if datafile.endswith(".json"):

                        corpus = self.read_corpus(
                            os.path.join(dirname, datafile)
                        )

                        for key in list(corpus.keys()):
                            corpora.append(corpus[key])
        else:
            corpus = self.read_corpus(corpus_path)

            for key in list(corpus.keys()):
                corpora.append(corpus[key])

        return corpora
<|endoftext|>


<commit_before>import asyncio
import types
import shelve
from datetime import datetime
from time import mktime
from logging import getLogger, StreamHandler, Formatter, INFO, DEBUG, WARNING, ERROR
from functools import wraps

import aiohttp
import click
import feedparser

# Default run parameters #######################################
DEFAULT_RUN_PARAMS = {
    'targets': None,
    'max_entries': None,
    'ignore_fetched': True,
    'asyncio_semaphore': 5,
}

# Logger settings        #######################################
def _create_logger(name):
    formatter = Formatter('[%(levelname)s] %(message)s')
    sh = StreamHandler()
    sh.setFormatter(formatter)
    logger = getLogger(name)
    logger.addHandler(sh)
    return logger

logger = _create_logger(__name__)

# Shelve store            ########################################
class ShelveStore:
    def __init__(self, file_path):
        self.file_path = file_path

    def update_or_create(self, key, value):
        with shelve.open(self.file_path) as db:
            db[key] = value

    def load(self, key):
        with shelve.open(self.file_path) as db:
            return db.get(key, None)

# Fetch feed, entry      ########################################
async def _fetch_feed(feed_url):
    async with aiohttp.ClientSession() as session:
        async with session.get(feed_url) as response:
            body = await response.text()
            parsed = feedparser.parse(body)
            feed = parsed.feed
            feed_info = {
                'site_title': feed.title,
                'site_subtitle': feed.subtitle,
                'site_url': feed.link,
                'fetched_at': datetime.now()
            }
            entries = parsed.entries
            return feed_info, entries

def _get_entry_info(entry):
    published = datetime.fromtimestamp(mktime(entry.published_parsed)) if hasattr(entry, 'published_parsed') else None
    updated = datetime.fromtimestamp(mktime(entry.updated_parsed)) if hasattr(entry, 'updated_parsed') else None
    return {
        'article_title': entry.title,
        'article_url': entry.link,
        'published_at': published,
        'updated_at': updated,
    }

async def _fetch_body(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.text()

async def _bound_fetch_body(url, semaphore=DEFAULT_RUN_PARAMS['asyncio_semaphore']):
    sem = asyncio.Semaphore(semaphore)
    async with sem:
        return await _fetch_body(url)

async def _fetch_bodies(entries):
    tasks = [asyncio.ensure_future(_bound_fetch_body(e.link)) for e in entries]
    bodies = asyncio.gather(*tasks)
    return await bodies

def _should_fetch(entry_updated_parsed, ignore_fetched, last_fetched_at):
    if ignore_fetched or last_fetched_at is None:
        return True
    if datetime.fromtimestamp(mktime(entry_updated_parsed)) > last_fetched_at:
        return True
    return False

class Feedy:
    feeds = {}
    plugins = []
    store = None

    def __init__(self, store=None):
        if store:
            self.set_store(store)

    def set_store(self, store):
        if hasattr(store, 'update_or_create') or hasattr(store, 'load'):
            self.store = store
        elif isinstance(store, str):
            self.store = ShelveStore(store)
        else:
            raise TypeError("Store must be string or implement ``.update_or_create()`` and ``load()``.")

    def install(self, plugin):
        if hasattr(plugin, 'setup'):
            plugin.setup(self)
        self.plugins.append(plugin)
        logger.debug("{} plugin is registered.".format(plugin.__name__))

    def add(self, feed_url, callback=None):
        @wraps(callback)
        def decorator(callback_func):
            self.feeds[callback_func.__name__] = (feed_url, callback_func)
            return callback_func
        return decorator(callback) if callback else decorator

    def entry_handler(self, callback, body, entry, feed_info):
        for plugin in self.plugins:
            callback = plugin(callback) if callable(plugin) else callback

        entry_info = _get_entry_info(entry)
        feed_info.update(entry_info)  # This dict need not copy using `copy.deepcopy()`.
        callback(info=feed_info, body=body)

    def feed_handler(self, callback, loop, feed_info, entries,
                     max_entries=DEFAULT_RUN_PARAMS['max_entries'],
                     ignore_fetched=DEFAULT_RUN_PARAMS['ignore_fetched']):
        last_fetched_at = self.store.load('{}_fetched_at'.format(callback.__name__)) if self.store else None
        entries = [e for e in entries[:max_entries] if _should_fetch(e.updated_parsed, ignore_fetched, last_fetched_at)]
        future = asyncio.ensure_future(_fetch_bodies(entries))
        bodies = loop.run_until_complete(future)

        if not entries:
            logger.info("No updates in {}.".format(callback.__name__))
        if self.store:
            self.store.update_or_create('{}_fetched_at'.format(callback.__name__), datetime.now())
            logger.debug("A last_fetched_at is just updated.")
        else:
            logger.debug("A last_fetched_at isn't updated.")

        for i, entry in enumerate(entries):
            self.entry_handler(callback, bodies[i], entry, feed_info)

    def target_handler(self, target, loop,
                       max_entries=DEFAULT_RUN_PARAMS['max_entries'],
                       ignore_fetched=DEFAULT_RUN_PARAMS['ignore_fetched']):
        feed_url, callback = self.feeds[target]
        fetched_feed = loop.run_until_complete(_fetch_feed(feed_url))
        feed_info, entries = fetched_feed
        self.feed_handler(callback, loop, feed_info, entries, max_entries=max_entries, ignore_fetched=ignore_fetched)

    def run(self, targets=DEFAULT_RUN_PARAMS['targets'], **kwargs):
        if not targets:
            targets = self.feeds.keys()

        loop = asyncio.get_event_loop()
        for t in targets:
            self.target_handler(t, loop, **kwargs)
        loop.close()

# Command Line Interface ######################################################
def _get_runner(src, obj):
    t = types.ModuleType('runner')
    with src as mod:
        t.__dict__['__file__'] = src.name
        compiled = compile(mod.read(), src.name, 'exec')
        exec(compiled, t.__dict__)
    runner = getattr(t, obj)
    return runner

def set_logger_level(verbose):
    logger_level = (ERROR, WARNING, INFO, DEBUG)
    logger.setLevel(logger_level[verbose])

@click.command()
@click.argument('src', type=click.File('r'), nargs=1)
@click.argument('obj', nargs=1)
@click.option('-v', '--verbose', type=click.IntRange(0, 3), count=True, help='Set log level')
@click.option('-t', '--targets', type=str, multiple=True, default=DEFAULT_RUN_PARAMS['targets'],
              help='The target function names.')
@click.option('-m', '--max-entries', type=int, default=DEFAULT_RUN_PARAMS['max_entries'],
              help="The maximum length for fetching entries every RSS feed")
@click.option('--ignore-fetched/--no-ignore-fetched', default=DEFAULT_RUN_PARAMS['ignore_fetched'],
              help="The maximum length for fetching entries every RSS feed")
def cmd(src, obj, verbose, **kwargs):
    """Run your feedy's project flexibly."""
    set_logger_level(verbose)
    runner = _get_runner(src, obj)
    runner.run(**kwargs)
<commit_msg>Modify defalult run parameters<commit_after>import asyncio
import types
import shelve
from datetime import datetime
from time import mktime
from logging import getLogger, StreamHandler, Formatter, INFO, DEBUG, WARNING, ERROR
from functools import wraps

import aiohttp
import click
import feedparser

# Default run parameters #######################################
DEFAULT_RUN_PARAMS = {
    'targets': None,
    'max_entries': None,
    'ignore_fetched': False,
    'asyncio_semaphore': 5,
}

# Logger settings        #######################################
def _create_logger(name):
    formatter = Formatter('[%(levelname)s] %(message)s')
    sh = StreamHandler()
    sh.setFormatter(formatter)
    logger = getLogger(name)
    logger.addHandler(sh)
    return logger

logger = _create_logger(__name__)

# Shelve store            ########################################
class ShelveStore:
    def __init__(self, file_path):
        self.file_path = file_path

    def update_or_create(self, key, value):
        with shelve.open(self.file_path) as db:
            db[key] = value

    def load(self, key):
        with shelve.open(self.file_path) as db:
            return db.get(key, None)

# Fetch feed, entry      ########################################
async def _fetch_feed(feed_url):
    async with aiohttp.ClientSession() as session:
        async with session.get(feed_url) as response:
            body = await response.text()
            parsed = feedparser.parse(body)
            feed = parsed.feed
            feed_info = {
                'site_title': feed.title,
                'site_subtitle': feed.subtitle,
                'site_url': feed.link,
                'fetched_at': datetime.now()
            }
            entries = parsed.entries
            return feed_info, entries

def _get_entry_info(entry):
    published = datetime.fromtimestamp(mktime(entry.published_parsed)) if hasattr(entry, 'published_parsed') else None
    updated = datetime.fromtimestamp(mktime(entry.updated_parsed)) if hasattr(entry, 'updated_parsed') else None
    return {
        'article_title': entry.title,
        'article_url': entry.link,
        'published_at': published,
        'updated_at': updated,
    }

async def _fetch_body(url):
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.text()

async def _bound_fetch_body(url, semaphore=DEFAULT_RUN_PARAMS['asyncio_semaphore']):
    sem = asyncio.Semaphore(semaphore)
    async with sem:
        return await _fetch_body(url)

async def _fetch_bodies(entries):
    tasks = [asyncio.ensure_future(_bound_fetch_body(e.link)) for e in entries]
    bodies = asyncio.gather(*tasks)
    return await bodies

def _should_fetch(entry_updated_parsed, ignore_fetched, last_fetched_at):
    if ignore_fetched or last_fetched_at is None:
        return True
    if datetime.fromtimestamp(mktime(entry_updated_parsed)) > last_fetched_at:
        return True
    return False

class Feedy:
    feeds = {}
    plugins = []
    store = None

    def __init__(self, store=None):
        if store:
            self.set_store(store)

    def set_store(self, store):
        if hasattr(store, 'update_or_create') or hasattr(store, 'load'):
            self.store = store
        elif isinstance(store, str):
            self.store = ShelveStore(store)
        else:
            raise TypeError("Store must be string or implement ``.update_or_create()`` and ``load()``.")

    def install(self, plugin):
        if hasattr(plugin, 'setup'):
            plugin.setup(self)
        self.plugins.append(plugin)
        logger.debug("{} plugin is registered.".format(plugin.__name__))

    def add(self, feed_url, callback=None):
        @wraps(callback)
        def decorator(callback_func):
            self.feeds[callback_func.__name__] = (feed_url, callback_func)
            return callback_func
        return decorator(callback) if callback else decorator

    def entry_handler(self, callback, body, entry, feed_info):
        for plugin in self.plugins:
            callback = plugin(callback) if callable(plugin) else callback

        entry_info = _get_entry_info(entry)
        feed_info.update(entry_info)  # This dict need not copy using `copy.deepcopy()`.
        callback(info=feed_info, body=body)

    def feed_handler(self, callback, loop, feed_info, entries,
                     max_entries=DEFAULT_RUN_PARAMS['max_entries'],
                     ignore_fetched=DEFAULT_RUN_PARAMS['ignore_fetched']):
        last_fetched_at = self.store.load('{}_fetched_at'.format(callback.__name__)) if self.store else None
        entries = [e for e in entries[:max_entries] if _should_fetch(e.updated_parsed, ignore_fetched, last_fetched_at)]
        future = asyncio.ensure_future(_fetch_bodies(entries))
        bodies = loop.run_until_complete(future)

        if not entries:
            logger.info("No updates in {}.".format(callback.__name__))
        if self.store:
            self.store.update_or_create('{}_fetched_at'.format(callback.__name__), datetime.now())
            logger.debug("A last_fetched_at is just updated.")
        else:
            logger.debug("A last_fetched_at isn't updated.")

        for i, entry in enumerate(entries):
            self.entry_handler(callback, bodies[i], entry, feed_info)

    def target_handler(self, target, loop,
                       max_entries=DEFAULT_RUN_PARAMS['max_entries'],
                       ignore_fetched=DEFAULT_RUN_PARAMS['ignore_fetched']):
        feed_url, callback = self.feeds[target]
        fetched_feed = loop.run_until_complete(_fetch_feed(feed_url))
        feed_info, entries = fetched_feed
        self.feed_handler(callback, loop, feed_info, entries, max_entries=max_entries, ignore_fetched=ignore_fetched)

    def run(self, targets=DEFAULT_RUN_PARAMS['targets'], **kwargs):
        if not targets:
            targets = self.feeds.keys()

        loop = asyncio.get_event_loop()
        for t in targets:
            self.target_handler(t, loop, **kwargs)
        loop.close()

# Command Line Interface ######################################################
def _get_runner(src, obj):
    t = types.ModuleType('runner')
    with src as mod:
        t.__dict__['__file__'] = src.name
        compiled = compile(mod.read(), src.name, 'exec')
        exec(compiled, t.__dict__)
    runner = getattr(t, obj)
    return runner

def set_logger_level(verbose):
    logger_level = (ERROR, WARNING, INFO, DEBUG)
    logger.setLevel(logger_level[verbose])

@click.command()
@click.argument('src', type=click.File('r'), nargs=1)
@click.argument('obj', nargs=1)
@click.option('-v', '--verbose', type=click.IntRange(0, 3), count=True, help='Set log level')
@click.option('-t', '--targets', type=str, multiple=True, default=DEFAULT_RUN_PARAMS['targets'],
              help='The target function names.')
@click.option('-m', '--max-entries', type=int, default=DEFAULT_RUN_PARAMS['max_entries'],
              help="The maximum length for fetching entries every RSS feed")
@click.option('--ignore-fetched/--no-ignore-fetched', default=DEFAULT_RUN_PARAMS['ignore_fetched'],
              help="The maximum length for fetching entries every RSS feed")
def cmd(src, obj, verbose, **kwargs):
    """Run your feedy's project flexibly."""
    set_logger_level(verbose)
    runner = _get_runner(src, obj)
    runner.run(**kwargs)
<|endoftext|>


<commit_before>import re

class FindSpam:
  rules = [
   {'regex': "\\b(baba(ji)?|vashikaran|fashion|here is|porn)\\b", 'all': True,
    'sites': [], 'reason': "Bad keyword detected"},
   {'regex': "\\+\\d{10}|\\+?\\d{2}\\s?\\d{8}", 'all': True,
    'sites': ["patents.stackexchange.com"], 'reason': "Phone number detected"},
   {'regex': "\\b(asshole|crap|fag|fuck|idiot|shit|whore)s?\\b", 'all': True,
    'sites': [], 'reason': "Offensive title detected"},
   {'regex': "^[A-Z0-9\.\?\\s]*$", 'all': True, 'sites': [], 'reason': "All-caps title"}
  ]

  @staticmethod
  def testpost(title, site):
    result = [];
    for rule in FindSpam.rules:
      if rule['all'] != (site in rule['sites']):
        if re.compile(rule['regex']).search(title):
          result.append(rule['reason'])
    return result

  @staticmethod
  def testtitle(title):
    regexes=["\\b(baba(ji)?|vashikaran|fashion|here is|porn)\\b","\\+\\d{10}","\\+?\\d{2}\\s?\\d{8}","\\b(asshole|crap|fag|fuck|idiot|shit|whore)s?\\b"]
    result = []
    p = [not not re.compile(s).search(title) for s in regexes]
    if 'vashikaran' in title or 'baba' in title or True in p:
      result.append('Possible spam')
      # magic if matches word
    elif title.upper() == title:
      result.append('All caps title')
      # magic if all caps
    return result
<commit_msg>Better phone detection. Die, vashikaran, die.<commit_after>import re

class FindSpam:
  rules = [
   {'regex': "\\b(baba(ji)?|vashikaran|fashion|here is|porn)\\b", 'all': True,
    'sites': [], 'reason': "Bad keyword detected"},
   {'regex': "\\+\\d{10}|\\+?\\d{2}[\\s\\-]?\\d{8,11}", 'all': True,
    'sites': ["patents.stackexchange.com"], 'reason': "Phone number detected"},
   {'regex': "\\b(asshole|crap|fag|fuck|idiot|shit|whore)s?\\b", 'all': True,
    'sites': [], 'reason': "Offensive title detected"},
   {'regex': "^[A-Z0-9\\.\\?\\s]*$", 'all': True, 'sites': [], 'reason': "All-caps title"}
  ]

  @staticmethod
  def testpost(title, site):
    result = [];
    for rule in FindSpam.rules:
      if rule['all'] != (site in rule['sites']):
        if re.compile(rule['regex']).search(title):
          result.append(rule['reason'])
    return result

  @staticmethod
  def testtitle(title):
    regexes=["\\b(baba(ji)?|vashikaran|fashion|here is|porn)\\b","\\+\\d{10}","\\+?\\d{2}\\s?\\d{8}","\\b(asshole|crap|fag|fuck|idiot|shit|whore)s?\\b"]
    result = []
    p = [not not re.compile(s).search(title) for s in regexes]
    if 'vashikaran' in title or 'baba' in title or True in p:
      result.append('Possible spam')
      # magic if matches word
    elif title.upper() == title:
      result.append('All caps title')
      # magic if all caps
    return result
<|endoftext|>


<commit_before>"""
This plugin prepares the given servers for RKE and creates an RKE config file
"""
import os
import time

import yaml

from common.methods import set_progress
from infrastructure.models import Server

def create_ssh_keypair(size=2048):
    """
    Make a new ssh keypair of a given size
    :param: size (optional, defaults to 2048). How many bits large should the key be?
    :return: UTF-8 strings representing the public key and private key in that order
    """
    from cryptography.hazmat.primitives import serialization as crypto_serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.hazmat.backends import default_backend as crypto_default_backend

    key = rsa.generate_private_key(
        backend=crypto_default_backend(),
        public_exponent=65537,
        key_size=size,
    )
    private_key = key.private_bytes(
        crypto_serialization.Encoding.PEM,
        crypto_serialization.PrivateFormat.PKCS8,
        crypto_serialization.NoEncryption())
    public_key = key.public_key().public_bytes(
        crypto_serialization.Encoding.OpenSSH,
        crypto_serialization.PublicFormat.OpenSSH
    )
    return public_key.decode('utf-8'), private_key.decode('utf-8')

def generate_rke_yaml(ips, user, ssh_private_key):
    """
    Make the YAML file that RKE is going to use to create the kubernetes cluster
    :param ips: What IP addresses are we working with? Array of strings
    :param user: What user should we create?
    :param ssh_private_key: Private key text (utf-8) to include in yaml
    :return: string with the formatted yaml in it
    """
    nodes = []
    for ip in ips:
        nodes.append(
            {
                'address': ip,
                'user': user,
                'ssh_key': ssh_private_key,
                'role': [],
            },
        )
    for i, role in enumerate(['controlplane', 'etcd']):
        node_idx = i % len(nodes)
        nodes[node_idx]['role'].append(role)

    for node in nodes:
        if len(node['role']) == 0:
            node['role'].append('worker')

    worker_count = 0
    for node in nodes:
        if 'worker' in node['role']:
            worker_count += 1

    if worker_count == 0:
        for node in nodes:
            if 'controlplane' not in node['role']:
                node['role'].append('worker')
                worker_count += 1
    if worker_count == 0:
        nodes[0]['role'].append('worker')

    etcd_count = 0
    for node in nodes:
        if 'etcd' in node['role']:
            etcd_count += 1
    if etcd_count < 3:
        for node in nodes:
            if 'etcd' not in node['role'] and 'controlplane' not in node['role']:
                node['role'].append('etcd')
                etcd_count += 1
                if etcd_count >= 3:
                    break

    services = {
        'etcd': {'image': 'quay.io/coreos/etcd:latest'},
        'kube-api': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kube-controller': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'scheduler': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kubelet': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kubeproxy': {'image': 'rancher/k8s:v1.11.6-rancher2'},
    }
    addons = [
        {
            "apiVersion": "v1",
            "kind": "ServiceAccount",
            "metadata": {
                "name": "cloudbolt-admin",
                "namespace": "kube-system"
            }
        },
        {
            "apiVersion": "rbac.authorization.k8s.io/v1",
            "kind": "ClusterRoleBinding",
            "metadata": {
                "name": "cloudbolt-admin"
            },
            "roleRef": {
                "apiGroup": "rbac.authorization.k8s.io",
                "kind": "ClusterRole",
                "name": "cluster-admin"
            },
            "subjects": [
                {
                    "kind": "ServiceAccount",
                    "name": "cloudbolt-admin",
                    "namespace": "kube-system"
                }
            ]
        },
    ]

    document = {
        'nodes': nodes,
        'services': services,
        'addons': '---\n' + yaml.dump_all(addons),
    }

    return yaml.dump(document)

def find_all_server_ips(blueprint_context):
    """
    Given the blueprint context, find the IP addresses for each server
    :param blueprint_context:
    :return: list of ip addresses represented as strings
    """
    ips = []
    for server_obj in find_all_servers(blueprint_context):
        if hasattr(server_obj, 'ip') and server_obj.ip is not None and server_obj.ip != '':
            ips.append(server_obj.ip)
    return ips

def find_all_servers(blueprint_context):
    """
    Given the blueprint context, find all servers
    :param blueprint_context:
    :return: iterable (yield) over Server objects
    """
    for key, value in blueprint_context.items():
        if isinstance(value, dict) and 'servers' in value:
            server_query_set = value.get('servers')
            for server_obj in server_query_set:
                yield server_obj

def prepare_server_hosts(user, blueprint_context, ssh_public_key):
    docker_script = 'yum -y install docker || exit 1;\n' + \
                    'systemctl enable docker || exit 1;\n' + \
                    'useradd {};\n'.format(user) + \
                    'groupadd docker\n' + \
                    'usermod -aG docker {} || exit 1\n'.format(user) + \
                    'mkdir -p /home/{}/.ssh || exit 1;\n'.format(user) + \
                    'echo \'{}\' >> /home/{}/.ssh/authorized_keys || exit 1;\n'.format(ssh_public_key, user) + \
                    'chown -R {} /home/{}/.ssh || exit 1;\n'.format(user, user) + \
                    'chmod 755 /home/{}/.ssh || exit 1;\n'.format(user) + \
                    'chmod 644 /home/{}/.ssh/authorized_keys || exit 1;\n'.format(user) + \
                    'echo \'net.ipv6.conf.all.forwarding=1\' >> /etc/sysctl.conf || exit 1\n' + \
                    'sysctl -p /etc/sysctl.conf || exit 1'
    for tcp_port_num in [443, 10250, 2379, 2380, 6443]:
        docker_script += '\nfirewall-offline-cmd --add-port={}/tcp || exit 1;'.format(tcp_port_num)
    for udp_port_num in [8285, 8472]:
        docker_script += '\nfirewall-offline-cmd --add-port={}/udp || exit 1;'.format(udp_port_num)
    docker_script += '\nsystemctl restart firewalld;\n'
    print("docker script:\n{}")
    for server in find_all_servers(blueprint_context=blueprint_context):
        set_progress("starting script execution on server {}".format(server.ip))
        server.execute_script(script_contents=docker_script, timeout=700)
        server.reboot()
    time.sleep(10)
    for server in find_all_servers(blueprint_context=blueprint_context):
        server.wait_for_os_readiness()

def run(_job, *_args, **kwargs):
    """
    main entry point for the plugin
    """
    user = 'cbrke'
    blueprint_context = kwargs.get('blueprint_context', {})
    ssh_public_key, ssh_private_key = create_ssh_keypair()
    prepare_server_hosts(user, blueprint_context, ssh_public_key)

    os.makedirs('/var/opt/cloudbolt/rke', exist_ok=True)
    with open("/var/opt/cloudbolt/rke/rke_private_key.pem", 'w') as fl:
        fl.write(ssh_private_key)
    with open("/var/opt/cloudbolt/rke/rke_public_key.pem", 'w') as fl:
        fl.write(ssh_public_key)

    ips = find_all_server_ips(kwargs.get('blueprint_context', {}))
    rke_yaml_text = generate_rke_yaml(ips, user, ssh_private_key)
    with open("/var/opt/cloudbolt/rke/cluster.yml", 'w') as fl:
        fl.write(rke_yaml_text)

    set_progress("Your ssh public and private keys have been generated. Please find them in /var/opt/coudbolt/rke")
    set_progress("Your RKE cluster.yml file has been generated. Please find it in /var/opt/cloudbolt/rke/cluster.yml")

    return "SUCCESS", "./rke up --config=/var/opt/cloudbolt/rke/cluster.yml", ""<commit_msg>Switch from hard-coded /var to the VARDIR in the settings file<commit_after>"""
This plugin prepares the given servers for RKE and creates an RKE config file
"""
import os
import time

import yaml

from common.methods import set_progress
from utilities.logger import ThreadLogger

import settings

CBVARDIR = os.path.join(settings.VARDIR, "opt", "cloudbolt")
logger = ThreadLogger(__name__)

def create_ssh_keypair(size=2048):
    """
    Make a new ssh keypair of a given size
    :param: size (optional, defaults to 2048). How many bits large should the key be?
    :return: UTF-8 strings representing the public key and private key in that order
    """
    from cryptography.hazmat.primitives import serialization as crypto_serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.hazmat.backends import default_backend as crypto_default_backend

    key = rsa.generate_private_key(
        backend=crypto_default_backend(),
        public_exponent=65537,
        key_size=size,
    )
    private_key = key.private_bytes(
        crypto_serialization.Encoding.PEM,
        crypto_serialization.PrivateFormat.PKCS8,
        crypto_serialization.NoEncryption())
    public_key = key.public_key().public_bytes(
        crypto_serialization.Encoding.OpenSSH,
        crypto_serialization.PublicFormat.OpenSSH
    )
    return public_key.decode('utf-8'), private_key.decode('utf-8')

def generate_rke_yaml(ips, user, ssh_private_key):
    """
    Make the YAML file that RKE is going to use to create the kubernetes cluster
    :param ips: What IP addresses are we working with? Array of strings
    :param user: What user should we create?
    :param ssh_private_key: Private key text (utf-8) to include in yaml
    :return: string with the formatted yaml in it
    """
    nodes = []
    for ip in ips:
        nodes.append(
            {
                'address': ip,
                'user': user,
                'ssh_key': ssh_private_key,
                'role': [],
            },
        )
    for i, role in enumerate(['controlplane', 'etcd']):
        node_idx = i % len(nodes)
        nodes[node_idx]['role'].append(role)

    for node in nodes:
        if len(node['role']) == 0:
            node['role'].append('worker')

    worker_count = 0
    for node in nodes:
        if 'worker' in node['role']:
            worker_count += 1

    if worker_count == 0:
        for node in nodes:
            if 'controlplane' not in node['role']:
                node['role'].append('worker')
                worker_count += 1
    if worker_count == 0:
        nodes[0]['role'].append('worker')

    etcd_count = 0
    for node in nodes:
        if 'etcd' in node['role']:
            etcd_count += 1
    if etcd_count < 3:
        for node in nodes:
            if 'etcd' not in node['role'] and 'controlplane' not in node['role']:
                node['role'].append('etcd')
                etcd_count += 1
                if etcd_count >= 3:
                    break

    services = {
        'etcd': {'image': 'quay.io/coreos/etcd:latest'},
        'kube-api': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kube-controller': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'scheduler': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kubelet': {'image': 'rancher/k8s:v1.11.6-rancher2'},
        'kubeproxy': {'image': 'rancher/k8s:v1.11.6-rancher2'},
    }
    addons = [
        {
            "apiVersion": "v1",
            "kind": "ServiceAccount",
            "metadata": {
                "name": "cloudbolt-admin",
                "namespace": "kube-system"
            }
        },
        {
            "apiVersion": "rbac.authorization.k8s.io/v1",
            "kind": "ClusterRoleBinding",
            "metadata": {
                "name": "cloudbolt-admin"
            },
            "roleRef": {
                "apiGroup": "rbac.authorization.k8s.io",
                "kind": "ClusterRole",
                "name": "cluster-admin"
            },
            "subjects": [
                {
                    "kind": "ServiceAccount",
                    "name": "cloudbolt-admin",
                    "namespace": "kube-system"
                }
            ]
        },
    ]

    document = {
        'nodes': nodes,
        'services': services,
        'addons': '---\n' + yaml.dump_all(addons),
    }

    return yaml.dump(document)

def find_all_server_ips(blueprint_context):
    """
    Given the blueprint context, find the IP addresses for each server
    :param blueprint_context:
    :return: list of ip addresses represented as strings
    """
    ips = []
    for server_obj in find_all_servers(blueprint_context):
        if hasattr(server_obj, 'ip') and server_obj.ip is not None and server_obj.ip != '':
            ips.append(server_obj.ip)
    return ips

def find_all_servers(blueprint_context):
    """
    Given the blueprint context, find all servers
    :param blueprint_context:
    :return: iterable (yield) over Server objects
    """
    for key, value in blueprint_context.items():
        if isinstance(value, dict) and 'servers' in value:
            server_query_set = value.get('servers')
            for server_obj in server_query_set:
                yield server_obj

def prepare_server_hosts(user, blueprint_context, ssh_public_key):
    docker_script = 'yum -y install docker || exit 1;\n' + \
                    'systemctl enable docker || exit 1;\n' + \
                    'useradd {};\n'.format(user) + \
                    'groupadd docker\n' + \
                    'usermod -aG docker {} || exit 1\n'.format(user) + \
                    'mkdir -p /home/{}/.ssh || exit 1;\n'.format(user) + \
                    'echo \'{}\' >> /home/{}/.ssh/authorized_keys || exit 1;\n'.format(ssh_public_key, user) + \
                    'chown -R {} /home/{}/.ssh || exit 1;\n'.format(user, user) + \
                    'chmod 755 /home/{}/.ssh || exit 1;\n'.format(user) + \
                    'chmod 644 /home/{}/.ssh/authorized_keys || exit 1;\n'.format(user) + \
                    'echo \'net.ipv6.conf.all.forwarding=1\' >> /etc/sysctl.conf || exit 1\n' + \
                    'sysctl -p /etc/sysctl.conf || exit 1'
    for tcp_port_num in [443, 10250, 2379, 2380, 6443]:
        docker_script += '\nfirewall-offline-cmd --add-port={}/tcp || exit 1;'.format(tcp_port_num)
    for udp_port_num in [8285, 8472]:
        docker_script += '\nfirewall-offline-cmd --add-port={}/udp || exit 1;'.format(udp_port_num)
    docker_script += '\nsystemctl restart firewalld;\n'
    logger.info(f"docker script:\n{docker_script}")
    for server in find_all_servers(blueprint_context=blueprint_context):
        set_progress("starting script execution on server {}".format(server.ip))
        server.execute_script(script_contents=docker_script, timeout=700)
        server.reboot()
    time.sleep(10)
    for server in find_all_servers(blueprint_context=blueprint_context):
        server.wait_for_os_readiness()

def run(_job, *_args, **kwargs):
    """
    main entry point for the plugin
    """
    user = 'cbrke'
    blueprint_context = kwargs.get('blueprint_context', {})
    ssh_public_key, ssh_private_key = create_ssh_keypair()
    prepare_server_hosts(user, blueprint_context, ssh_public_key)

    os.makedirs('/var/opt/cloudbolt/rke', exist_ok=True)
    with open(f"{CBVARDIR}rke/rke_private_key.pem", 'w') as fl:
        fl.write(ssh_private_key)
    with open(f"{CBVARDIR}rke/rke_public_key.pem", 'w') as fl:
        fl.write(ssh_public_key)

    ips = find_all_server_ips(kwargs.get('blueprint_context', {}))
    rke_yaml_text = generate_rke_yaml(ips, user, ssh_private_key)
    with open(f"{CBVARDIR}rke/cluster.yml", 'w') as fl:
        fl.write(rke_yaml_text)

    set_progress(f"Your ssh public and private keys have been generated. Please find them in {CBVARDIR}rke/")
    set_progress(f"Your RKE cluster.yml file has been generated. Please find it in {CBVARDIR}rke/cluster.yml")

    return "SUCCESS", "./rke up --config=/var/opt/cloudbolt/rke/cluster.yml", ""
<|endoftext|>


<commit_before>import os

BASEDIR = os.path.dirname(os.getcwd())
HOME = "file://" + BASEDIR + "/src/"
<commit_msg>correct src path in settings<commit_after>from os.path import dirname, realpath

BASEDIR = dirname(dirname(dirname(realpath(__file__))))
HOME = "file://" + BASEDIR + "/src/"
<|endoftext|>


<commit_before># rain_notifier.py
#
# This is a project that uses the Grove LED attached to an umbrella to remind
# people to take their umbrella if it's raining or going to rain that day.
#
# You will need to modify the first few lines to make this project work for you.
# Fee free to modify the rest of the program to make it behave the way you want it to.
#
# Enter your zip code (or Postal Code) on the following line
# (leave the name 'zip' as is, even if you have a postal code. )
zipcode = 'YOUR_ZIP_CODE'

# once you have gotten an API Key from Wunderground (see tutorial), pleave provide it here
# Wunderground API Key
api_key = 'YOUR_WUNDERGROUND_API_KEY'

# if you wish to use two LEDs, provide the port number of the second LED
# The second LED is the one that will show there's no rain in the weather prediction
# With two LEDs, one will always be on.
# if you do not want a second LED, set this to -1
clearled = 8

# in case of rain, do you want the LED to blink in order to catch your attention more?
# set to True for blinking, set to False for a steady warning
BLINK = True

############################################################################
# no modifications are required after this
# however the next few lines can be tweaked easily
############################################################################

# this is the delay between each polling of Wunderground. It's set to poll once every 5 minutes
# you may prefer a faster polling.
# the Wunderground free account does set a limit

DELAY = 5*60 # in seconds

# Pin for the rain LED on the umbrella
rainled = 7

############################################################################
# nothing needs to be modified after this
# if you are familiar with Python, feel free to change anything to improve this
# we welcome additions and suggestions
############################################################################

import urllib2
import json
from grovepi import *
import time
import sys

def clear_led(status):
    """ change clear sky LED status """
    if clearled > -1:
        digitalWrite(clearled,status)

def assign_rain(status, blink):
    """ controls the LEDs status

    Arguments:
    status: True if it will be raining
            False if there's no rain in the forecast
    blink:  True if the rain led should blink when on
            False if the rain led does not need blinking
    no return value
    """
    if status == True: # there will be rain
        clear_led(0) # Turn off the CLEAR LED on the umbrella

        # Light the RAIN LED on the umbrella
        # check if it needs to be blinking
        if blink == True:
            led_blink(DELAY) # no sleep required here
        else:
            digitalWrite(rainled,1)
            time.sleep(DELAY)

    else: # no rain in forecast
        # turn off the RAIN LED on the umbrella
        digitalWrite(rainled,0)
        # Turn on the CLEAR LED on the umbrella
        clear_led(1)
        time.sleep(DELAY)

def led_blink(blinkdelay):
    """ forces the rain LED to blink

    rain LED will blink at a rate of 0.2 secs on, and 0.2 secs off
    it will blink for a duration set by blinkdelay
    no additional sleep() is required as it's incorporated in the blinking
    """
    count = 0.0
    blinkrate = 0.4
    while count < float(blinkdelay):
        digitalWrite(rainled,1)
        time.sleep(blinkrate)
        digitalWrite(rainled,0)
        time.sleep(blinkrate)
        count += (2*blinkrate)

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

# Zip code of location
url='http://api.wunderground.com/api/'+api_key+'/geolookup/conditions/q/'+zip+'.json'

pinMode(rainled,"OUTPUT")

if len(sys.argv) > 1 and sys.argv[1] == 'test':
    try:
        clear_led(1)
        led_blink(5)
        clear_led(0)
    except KeyboardInterrupt:
        digitalWrite(clearled,0)
        digitalWrite(rainled,0)
    quit()

if clearled > -1:
    pinMode(clearled,"OUTPUT")

try:
    while True:
        f = urllib2.urlopen(url)
        json_string = f.read()
        f.close()
        parsed_json = json.loads(json_string)
        location = parsed_json['location']['city']
        precip_today_in = parsed_json['current_observation']['precip_today_in']

        print "Current precipitation in %s is: %s" % (location, precip_today_in)

        if float(precip_today_in) > 0 :
            print "Rain today, take the umbrella"
            assign_rain(True, BLINK)

        else:
            print "No Rain today"
            assign_rain(False, BLINK)

except KeyboardInterrupt:
    clear_led(0)
    digitalWrite(rainled,0)
<commit_msg>add test mode; PEP8 compliant; supports metric; lets user set up comfort level<commit_after># rain_notifier.py
#
# This is a project that uses the Grove LED attached to an umbrella to remind
# people to take their umbrella if it's raining or going to rain that day.
#
# You will need to modify the first few lines for this project work for you.
# Fee free to modify the rest of the program to make it
# behave the way you want it to.
#
# Enter your zip code (or Postal Code) on the following line
# (leave the name 'zip' as is, even if you have a postal code. )
zipcode = 'YOUR_ZIP_CODE'

# once you have gotten an API Key from Wunderground (see tutorial),
# pleave provide it here
# Wunderground API Key
api_key = 'YOUR_WUNDERGROUND_API_KEY'

# if you wish to use two LEDs, provide the port number of the second LED
# The second LED is the one that will indicate
# if there's no rain in the weather prediction
# With two LEDs, one will always be on.
# if you do not want a second LED, set this to -1
clearled = 8

############################################################################
# no modifications are required after this
# however the next few lines can be tweaked easily
############################################################################

# in case of rain, do you want the LED to blink
# to catch your attention more?
# set to True for blinking, set to False for a steady warning
BLINK = True

# this is the delay between each polling of Wunderground.
# It's set to poll once every 5 minutes
# you may prefer a faster polling.
# the Wunderground free account does set a limit
DELAY = 5*60  # in seconds

# Do you prefer to work in metric or imperial?
# set METRIC to 1 to use the metric system
# set METRIC to 0 to use imperial system
METRIC = 1

# rain threshold: what's your safety zone?
# before you need an umbrella
# if METRIC is set to 1, this will be in millimeters
# if METRIC is set to 0, this will be in inches
RAIN_THRESHOLD = 1

# Pin for the rain LED on the umbrella
rainled = 7

############################################################################
# nothing needs to be modified after this
# if you are familiar with Python, feel free to change anything to improve this
# we welcome additions and suggestions
############################################################################

import urllib2
import json
from grovepi import *
import time
import sys
from my_rain_notifier import *

def clear_led(status):
    """ change clear sky LED status """
    if clearled > -1:
        digitalWrite(clearled, status)

def assign_rain(status, blink):
    """ controls the LEDs status

    Arguments:
    status: True if it will be raining
            False if there's no rain in the forecast
    blink:  True if the rain led should blink when on
            False if the rain led does not need blinking
    no return value
    """

    if status is True:  # there will be rain
        clear_led(0)    # Turn off the CLEAR LED on the umbrella

        # Light the RAIN LED on the umbrella
        # first check if it needs to be blinking
        if blink is True:
            led_blink(DELAY)  # no sleep required here
        else:
            digitalWrite(rainled, 1)
            time.sleep(DELAY)

    else:  # no rain in forecast
        # turn off the RAIN LED on the umbrella
        digitalWrite(rainled, 0)
        # Turn on the CLEAR LED on the umbrella
        clear_led(1)
        time.sleep(DELAY)

def led_blink(blinkdelay):
    """ forces the rain LED to blink
    rain LED will blink at a rate of 0.2 secs on, and 0.2 secs off
    it will blink for a duration set by blinkdelay
    no additional sleep() is required as it's incorporated in the blinking
    """
    count = 0.0
    blinkrate = 0.4
    while count < float(blinkdelay):
        digitalWrite(rainled, 1)
        time.sleep(blinkrate)
        digitalWrite(rainled, 0)
        time.sleep(blinkrate)
        count += (2*blinkrate)

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

# Zip code of location
url = 'http://api.wunderground.com/api/' + api_key
url = url + '/geolookup/conditions/q/' + zipcode + '.json'

pinMode(rainled, "OUTPUT")

if len(sys.argv) > 1 and sys.argv[1] == 'test':
    try:
        clear_led(1)
        led_blink(5)
        clear_led(0)
    except KeyboardInterrupt:
        digitalWrite(clearled, 0)
        digitalWrite(rainled, 0)
    quit()

if clearled > -1:
    pinMode(clearled, "OUTPUT")

try:
    while True:
        f = urllib2.urlopen(url)
        json_string = f.read()
        f.close()
        parsed_json = json.loads(json_string)
        location = parsed_json['location']['city']
        if METRIC is 0:
            precip_today = parsed_json['current_observation']['precip_today_in']
        else:
            precip_today = parsed_json['current_observation']['precip_today_metric']
        print "Current precipitation in %s is: %s" % (location, precip_today)

        if float(precip_today) > float(RAIN_THRESHOLD):
            print "Rain today, take the umbrella"
            assign_rain(True, BLINK)

        else:
            print "No Rain today"
            assign_rain(False, BLINK)

except KeyboardInterrupt:
    clear_led(0)
    digitalWrite(rainled, 0)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# fight.py - fight another user, if they're available
#
# Copyright (c) 2020 Casey Bartlett <caseytb@bu.edu>
#
# See LICENSE for terms of usage, modification and redistribution.

from sopel.module import commands, example, NOLIMIT
import random

# Each user has the following associated state information (<<key: default>>)
# | <<player: nic>>
#	+ <<hitPoints: 100>>
#	+ <<xp: 0>>
#	+ <<xl: 0>>

# Xp levels
# start = 1
# stop  = log10(1000)
# nsteps = 5
# spacing = floor(logspace(start,stop,nsteps))
xlMap = {0:10,
	1:31
	2:100
	3:316
	4:1000}

fightStrings = [ "Bam boom pow! {source} stabs {target} for {damage} damage!",
 "Krakow! {source} filets {target} for {damage} damage!",
 "Reeee {source} uses loud screech on {target} dealing {damage} damage!",
 "Kaboom! {source} uses their dlc for SD to stab {target} for {damage} damage!"]

class fighter:
	def __init__(self, db, nick):
		self.db = db
		self.nick = nick
		self.xl = db.get_nick_value(nick, "xl")
		if not self.xl:
			self.xl = 0
			db.set_nick_value(nick, "xl", self.xl)
		self.xp = db.get_nick_value(nick, "xp")
		if not self.xp:
			self.xp = 0
			db.set_nick_value(nick, "xp", self.xp)
		self.hitPoints = db.get_nick_value(nick, "hitPoints")
		if not self.hitPoints:
			self.hitPoints=100
			db.set_nick_value(nick, "hitPoints", self.hitPoints)
	def receiveDamage(self, damage):
		msg = ""
		if self.hitPoints < damage:
			msg = " {nick} has died !".format(nick=self.nick)
			self.db.set_nick_value(self.nick, "hitPoints", 100)
		else:
			newhp = self.hitPoints - damage
			self.db.set_nick_value(self.nick, "hitPoints", newhp)
		return msg
	def receiveExperience(self, xp):
		self.db.set_nick_value(self.nick, "xp", xp)
		return fightEvents.onXpChange(self, xp)
	def setXl(self, newXl):
		self.db.set_nick_value(self.nick, "xl", newXl)

class fightEvents:
	# Event for "on xp change" - Don't assume only one level can be increased at one
	# time
	# if no xl change:
	# - return
	# if xl changes:
	# - write xl to database
	# - generate message
	def onXpChange(fighter, gainedXp):
		newXl = fighter.xl
		newXp = gainedXp + fighter.xp
		while xlMap[newXl] < newXp:
			newXl += 1
		xlDiff = newXl - fighter.xl
		if xlDiff > 0
			fighter.setXl(newXl)
		else:
			return ""
		msg = "You gained {totalLevels} experience level".format(xlDiff)
		if xlDiff  > 1:
			msg += "s"
		return msg;

def fight(source, target):
	maxIndex = len(fightStrings) - 1
	minDamage = 1
	maxDamage = 10
	minXpGain = 2
	maxXpGain = 15
	index = random.randint(0, maxIndex)
	damage = random.randint(minDamage, maxDamage)
	baseMsg = fightStrings[index].format(source=source, target=target, damage=damage)
	damageMsg = target.receiveDamage(damage)
	xlChangedMessage = ""
	if damageMsg:
		# Generate some random xp:
		xpGained = random.randint(minXpGain, maxXpGain)
		xlChangedMessage = fighter.receiveExperience(xpGained)
	return "{base} {damage} {xl}".format(baseMsg, damageMsg, xlChangedMessage)

@commands('fight')
@commands('.fight pf')
def fight(bot, trigger):
	sourceNick = trigger.nick
	channel = trigger.channel
	targetNick = trigger.group(2)
	if targetNick not in bot.channels[channel].users:
		bot.reply('{target} is not in this channel [[{channel}]]'.format(target=targetNick, channel=channel))
	# load the fighters
	sourceFighter = fighter(bot.db, sourceNick)
	targetFighter = fighter(bot.db, targetNick)
	msg = fight(sourceFighter, targetFighter)
	sopel.say(msg)
<commit_msg>added chance<commit_after>#!/usr/bin/env python
# fight.py - fight another user, if they're available
#
# Copyright (c) 2020 Casey Bartlett <caseytb@bu.edu>
#
# See LICENSE for terms of usage, modification and redistribution.

from sopel.module import commands, example, NOLIMIT
from sopel.tools import Identifier
import random

# Each user has the following associated state information (<<key: default>>)
# | <<player: nic>>
#	+ <<hitPoints: 100>>
#	+ <<xp: 0>>
#	+ <<xl: 0>>

# Xp levels
# start = 1
# stop  = log10(1000)
# nsteps = 5
# spacing = floor(logspace(start,stop,nsteps))
xlMap = {0:10,
	1:31,
	2:100,
	3:316,
	4:1000}

fightStrings = [ "Bam boom pow! {source} stabs {target} for {damage} damage!",
 "Krakow! {source} filets {target} for {damage} damage!",
 "Reeee {source} uses loud screech on {target} dealing {damage} damage!",
 "Kaboom! {source} uses their dlc for SD to stab {target} for {damage} damage!"]

MissStrings = [ "Oh no {source} misses {target} and deals no damage"]

CritMissStrings = [ "Oh no {source} is confused and attacked themself for {damage} damage!"]

class fighter:
	def __init__(self, db, nick):
		self.db = db
		self.nick = nick
		self.xl = db.get_nick_value(nick, "xl")
		if not self.xl:
			self.xl = 0
			db.set_nick_value(nick, "xl", self.xl)
		self.xp = db.get_nick_value(nick, "xp")
		if not self.xp:
			self.xp = 0
			db.set_nick_value(nick, "xp", self.xp)
		self.hitPoints = db.get_nick_value(nick, "hitPoints")
		if not self.hitPoints:
			self.hitPoints=100
			db.set_nick_value(nick, "hitPoints", self.hitPoints)
	def receiveDamage(self, damage):
		msg = ""
		if self.hitPoints < damage:
			msg = " {nick} has died !".format(nick=self.nick)
			self.db.set_nick_value(self.nick, "hitPoints", 100)
		else:
			newhp = self.hitPoints - damage
			self.db.set_nick_value(self.nick, "hitPoints", newhp)
		return msg
	def receiveExperience(self, xp):
		self.db.set_nick_value(self.nick, "xp", xp)
		return fightEvents.onXpChange(self, xp)
	def setXl(self, newXl):
		self.db.set_nick_value(self.nick, "xl", newXl)

class fightEvents:
	# Event for "on xp change" - Don't assume only one level can be increased at one
	# time
	# if no xl change:
	# - return
	# if xl changes:
	# - write xl to database
	# - generate message
	def onXpChange(aFighter, gainedXp):
		newXl = aFighter.xl
		newXp = gainedXp + aFighter.xp
		while xlMap[newXl] < newXp:
			newXl += 1
		xlDiff = newXl - aFighter.xl
		if xlDiff > 0:
			aFighter.setXl(newXl)
		else:
			return ""
		msg = "You gained {totalLevels} experience level".format(totalLevels=xlDiff)
		if xlDiff  > 1:
			msg += "s"
		return msg;
#todo add benifits to levels aka attack chance and armor?
def fightImpl(source, target):
	maxIndex = len(fightStrings) - 1
	minDamage = 1
	maxDamage = 25
	minXpGain = 2
	maxXpGain = 30
	minAttack = 1
	maxAttack = 100
	attack = random.randint(minAttack, maxAttack)
	index = random.randint(0, maxIndex)
	damage = random.randint(minDamage, maxDamage)
	damageMsg =""
	if attack >= 50 and attack < 95:
		baseMsg = fightStrings[index].format(source=source.nick, target=target.nick, damage=damage)
		damageMsg = target.receiveDamage(damage)
	elif attack >= 95:
		damage = damage * 2
		baseMsg = fightStrings[index].format(source=source.nick, target=target.nick, damage=damage)
		damageMsg = target.receiveDamage(damage)
	elif attack < 50 and attack > 5:
		baseMsg = fightStrings[index].format(source=source.nick, target=target.nick)
		damageMsg = target.receiveDamage(damage)
	elif attack <= 5:
		damage = damage / 2
		baseMsg = fightStrings[index].format(source=source.nick, damage=damage)
		damageMsg = source.receiveDamage(damage)
	else:
		baseMsg = fightStrings[index].format(source=source.nick, target=target.nick, damage=damage)
		damageMsg = target.receiveDamage(damage)
	xlChangedMessage = ""

	if damageMsg:
		# Generate some random xp:
		xpGained = random.randint(minXpGain, maxXpGain)
		xlChangedMessage = source.receiveExperience(xpGained)
	return "{base} {damage} {xl}".format(base=baseMsg, damage=damageMsg, xl=xlChangedMessage)

@commands('fight')
@example('.fight pf')
def fight(bot, trigger):
	sourceNick = trigger.nick
	channel = trigger.sender
	if not trigger.group(2):
		bot.say('fight whosits?')
		return
	targetNick = Identifier(trigger.group(2).strip())
	if targetNick == sourceNick:
		bot.say('Are u slow son?')
		return
	if targetNick == "knifebot":
		bot.say('knifebot dodges with robotic skill')
		return
	if channel not in bot.channels:
		bot.reply('You can''t fight here!')
		return
	if targetNick not in bot.channels[channel].users:
		bot.reply('{target} is not in this channel [[{channel}]]'.format(target=targetNick, channel=channel))
		return
	# load the fighters
	sourceFighter = fighter(bot.db, sourceNick)
	targetFighter = fighter(bot.db, targetNick)
	msg = fightImpl(sourceFighter, targetFighter)
	bot.say(msg)

@commands('fstat')
@example('.fstat fooobarrr')
def fighterStatus(bot, trigger):
	if not trigger.group(2):
		bot.say('fighter status for who?')
		return
	targetNick = Identifier(trigger.group(2).strip())
	hitpoints = bot.db.get_nick_value(targetNick,'hitPoints')
	xl = bot.db.get_nick_value(targetNick,'xl')
	if not hitpoints:
		bot.say('I can''t find stats for {nick}'.format(nick=targetNick))
		return
	else:
		bot.say('{nick} has {hp} hit points / 100 @ XL {xl}'.format(nick=targetNick, hp=hitpoints, xl=xl))

<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
# Licensed to the Apache Software Foundation (ASF) 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.

from flask import request
from gcecloudstack import app
from gcecloudstack import authentication
from gcecloudstack.services import requester
from gcecloudstack.controllers import zones, helper, operations, images
import json
import urllib

def _get_instances(authorization, args=None):
    command = 'listVirtualMachines'
    if not args:
        args = {}
    cloudstack_response = requester.make_request(
        command,
        args,
        authorization.client_id,
        authorization.client_secret
    )

    return cloudstack_response

def _deploy_virtual_machine(authorization, args):
    command = 'deployVirtualMachine'
    # do something here to convert the given args.
    convertedargs = {}
    convertedargs['templateid'] = images.get_template_id(
        args['template'], authorization
    )
    convertedargs['zoneid'] = zones.get_zone_id(
        args['zone'], authorization
    )
    convertedargs['serviceofferingid'] = args['serviceoffering']
    cloudstack_response = requester.make_request(
        command,
        convertedargs,
        authorization.client_id,
        authorization.client_secret
    )

    return cloudstack_response

def _cloudstack_instance_to_gce(cloudstack_response, selfLink=None, zone=None):
    response = {}
    response['kind'] = 'compute#instance'
    response['id'] = cloudstack_response['id']
    response['creationTimestamp'] = cloudstack_response['created']
    response['status'] = cloudstack_response['state'].upper()
    response['name'] = cloudstack_response['instancename']
    response['description'] = cloudstack_response['instancename']
    response['machineType'] = cloudstack_response['serviceofferingname']
    response['image'] = cloudstack_response['templatename']
    response['canIpForward'] = 'true'
    response['networkInterfaces'] = []

    networking = {}
    networking['network'] = 'tobereviewed'
    networking['networkIP'] = cloudstack_response['nic'][0]['ipaddress']
    networking['name'] = cloudstack_response['nic'][0]['id']

    response['networkInterfaces'].append(networking)

    if not selfLink:
        response['selfLink'] = request.base_url
    else:
        response['selfLink'] = selfLink

    if not zone:
        response['zone'] = cloudstack_response['zonename']
    else:
        response['zone'] = zone

    return response

@app.route('/' + app.config['PATH'] + '<projectid>/aggregated/instances', methods=['GET'])
@authentication.required
def aggregatedlistinstances(authorization, projectid):
    zonelist = zones.get_zone_names(authorization)
    instancesList = _get_instances(authorization)

    items = {}

    for zone in zonelist:
        zones_instances = []
        if instancesList['listvirtualmachinesresponse']:
            for instance in instancesList['listvirtualmachinesresponse']['virtualmachine']:
                zones_instances.append(
                    _cloudstack_instance_to_gce(
                        cloudstack_response=instance,
                        zone=zone,
                        selfLink=request.base_url + '/' + instance['instancename']
                    )
                )
        items['zone/' + zone] = {}
        items['zone/' + zone]['zone'] = zone
        items['zone/' + zone]['instances'] = zones_instances

    populated_response = {
        'kind': 'compute#instanceAggregatedList',
        'id': 'projects/' + projectid + '/instances',
        'selfLink': request.base_url,
        'items': items
    }

    return helper.createsuccessfulresponse(data=populated_response)

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances', methods=['GET'])
@authentication.required
def listinstances(authorization, projectid, zone):
    instance = None

    if 'filter' in request.args:
        filter = urllib.unquote_plus(request.args['filter'])
        filter = dict(filter.split(' eq ') for values in filter)
        if filter['name']:
            instance = filter['name']

    if instance:
        instanceList = _get_instances(
            authorization,
            args={'keyword': instance}
        )
    else:
        instanceList = _get_instances(authorization)

    items = []
    if instanceList['listvirtualmachinesresponse']:
        for instance in instanceList['listvirtualmachinesresponse']['virtualmachine']:
            items.append(_cloudstack_instance_to_gce(instance))

    populated_response = {
        'kind': 'compute#instanceList',
        'id': 'projects/' + projectid + '/instances',
        'selfLink': request.base_url,
        'items': items
    }

    return helper.createsuccessfulresponse(data=populated_response)

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances/<instance>', methods=['GET'])
@authentication.required
def getinstance(projectid, authorization, zone, instance):
    cloudstack_response = _get_instances(
        authorization,
        args={'keyword': instance}
    )

    return helper.createsuccessfulresponse(
        data=_cloudstack_instance_to_gce(cloudstack_response['listvirtualmachinesresponse']['virtualmachine'][0])
    )

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances', methods=['POST'])
@authentication.required
def addinstance(authorization, projectid, zone):
    data = json.loads(request.data)
    args = {}
    args['name'] = data['name']
    args['serviceoffering'] = data['machineType'].rsplit('/', 1)[1]
    args['template'] = data['image'].rsplit('/', 1)[1]
    args['zone'] = zone

    deploymentResult = _deploy_virtual_machine(authorization, args)

    populated_response = operations.create_response(
        projectid=projectid,
        operationid=deploymentResult['deployvirtualmachineresponse']['jobid'],
        authorization=authorization
    )

    app.logger.debug(str(populated_response))

    return helper.createsuccessfulresponse(data=populated_response)
<commit_msg>Fix issues around instance naming<commit_after>#!/usr/bin/env python
# encoding: utf-8
# Licensed to the Apache Software Foundation (ASF) 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.

from flask import request
from gcecloudstack import app
from gcecloudstack import authentication
from gcecloudstack.services import requester
from gcecloudstack.controllers import zones, helper, operations, images
import json
import urllib

def _get_instances(authorization, args=None):
    command = 'listVirtualMachines'
    if not args:
        args = {}
    cloudstack_response = requester.make_request(
        command,
        args,
        authorization.client_id,
        authorization.client_secret
    )

    return cloudstack_response

def _deploy_virtual_machine(authorization, args):
    command = 'deployVirtualMachine'
    # do something here to convert the given args.
    convertedargs = {}
    convertedargs['templateid'] = images.get_template_id(
        args['template'], authorization
    )
    convertedargs['zoneid'] = zones.get_zone_id(
        args['zone'], authorization
    )
    convertedargs['serviceofferingid'] = args['serviceoffering']
    convertedargs['displayname'] = args['name']
    convertedargs['name'] = args['name']
    cloudstack_response = requester.make_request(
        command,
        convertedargs,
        authorization.client_id,
        authorization.client_secret
    )

    return cloudstack_response

def _cloudstack_instance_to_gce(cloudstack_response, selfLink=None, zone=None):
    response = {}
    response['kind'] = 'compute#instance'
    response['id'] = cloudstack_response['id']
    response['creationTimestamp'] = cloudstack_response['created']
    response['status'] = cloudstack_response['state'].upper()
    response['name'] = cloudstack_response['name']
    response['description'] = cloudstack_response['name']
    response['machineType'] = cloudstack_response['serviceofferingname']
    response['image'] = cloudstack_response['templatename']
    response['canIpForward'] = 'true'
    response['networkInterfaces'] = []

    networking = {}
    networking['network'] = 'tobereviewed'
    networking['networkIP'] = cloudstack_response['nic'][0]['ipaddress']
    networking['name'] = cloudstack_response['nic'][0]['id']

    response['networkInterfaces'].append(networking)

    if not selfLink:
        response['selfLink'] = request.base_url
    else:
        response['selfLink'] = selfLink

    if not zone:
        response['zone'] = cloudstack_response['zonename']
    else:
        response['zone'] = zone

    return response

@app.route('/' + app.config['PATH'] + '<projectid>/aggregated/instances', methods=['GET'])
@authentication.required
def aggregatedlistinstances(authorization, projectid):
    zonelist = zones.get_zone_names(authorization)
    instancesList = _get_instances(authorization)

    items = {}

    for zone in zonelist:
        zones_instances = []
        if instancesList['listvirtualmachinesresponse']:
            for instance in instancesList['listvirtualmachinesresponse']['virtualmachine']:
                zones_instances.append(
                    _cloudstack_instance_to_gce(
                        cloudstack_response=instance,
                        zone=zone,
                        selfLink=request.base_url + '/' + instance['name']
                    )
                )
        items['zone/' + zone] = {}
        items['zone/' + zone]['zone'] = zone
        items['zone/' + zone]['instances'] = zones_instances

    populated_response = {
        'kind': 'compute#instanceAggregatedList',
        'id': 'projects/' + projectid + '/instances',
        'selfLink': request.base_url,
        'items': items
    }

    return helper.createsuccessfulresponse(data=populated_response)

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances', methods=['GET'])
@authentication.required
def listinstances(authorization, projectid, zone):
    instance = None

    if 'filter' in request.args:
        filter = urllib.unquote_plus(request.args['filter'])
        filter = dict(filter.split(' eq ') for values in filter)
        if filter['name']:
            instance = filter['name']

    if instance:
        instanceList = _get_instances(
            authorization,
            args={'keyword': instance}
        )
    else:
        instanceList = _get_instances(authorization)

    items = []
    if instanceList['listvirtualmachinesresponse']:
        for instance in instanceList['listvirtualmachinesresponse']['virtualmachine']:
            items.append(_cloudstack_instance_to_gce(instance))

    populated_response = {
        'kind': 'compute#instanceList',
        'id': 'projects/' + projectid + '/instances',
        'selfLink': request.base_url,
        'items': items
    }

    return helper.createsuccessfulresponse(data=populated_response)

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances/<instance>', methods=['GET'])
@authentication.required
def getinstance(projectid, authorization, zone, instance):
    cloudstack_response = _get_instances(
        authorization,
        args={'keyword': instance}
    )

    return helper.createsuccessfulresponse(
        data=_cloudstack_instance_to_gce(cloudstack_response['listvirtualmachinesresponse']['virtualmachine'][0])
    )

@app.route('/' + app.config['PATH'] + '<projectid>/zones/<zone>/instances', methods=['POST'])
@authentication.required
def addinstance(authorization, projectid, zone):
    data = json.loads(request.data)
    args = {}
    args['name'] = data['name']
    args['serviceoffering'] = data['machineType'].rsplit('/', 1)[1]
    args['template'] = data['image'].rsplit('/', 1)[1]
    args['zone'] = zone

    deploymentResult = _deploy_virtual_machine(authorization, args)

    populated_response = operations.create_response(
        projectid=projectid,
        operationid=deploymentResult['deployvirtualmachineresponse']['jobid'],
        authorization=authorization
    )

    app.logger.debug(str(populated_response))

    return helper.createsuccessfulresponse(data=populated_response)
<|endoftext|>


<commit_before>"""
The ``vtkInterface`` module provides an way for ``discretize`` meshes to be
converted to VTK data objects (and back when possible).

"""
import os
import numpy as np
import six

import vtk
import vtk.util.numpy_support as nps
from vtk import VTK_VERSION
from vtk import vtkXMLRectilinearGridWriter
from vtk import vtkXMLUnstructuredGridWriter
from vtk import vtkXMLStructuredGridWriter
from vtk import vtkXMLRectilinearGridReader

class vtkInterface(object):
    """This class is full of static methods that enable ``discretize`` meshes to
    be converted to VTK data objects (and back when possible).
    """

    def __treeMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkUnstructuredGrid`` object of this tree mesh and the
        given models as ``CellData`` of that VTK dataset.

        Input:
        :param mesh, discretize.TreeMesh - The tree mesh to convert to a ``vtkUnstructuredGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        # Make the data parts for the vtu object
        # Points
        ptsMat = np.vstack((mesh.gridN, mesh.gridhN))

        # Adjust if result was 2D (voxels are pixels in 2D):
        VTK_CELL_TYPE = vtk.VTK_VOXEL
        if ptsMat.shape[1] == 2:
            # TODO: assumes any 2D tree meshes are on the XY Plane
            # Add Z values of 0.0 if 2D
            ptsMat = np.c_[ptsMat, np.zeros(ptsMat.shape[0])]
            VTK_CELL_TYPE = vtk.VTK_PIXEL
        if ptsMat.shape[1] != 3:
            raise RuntimeError('Points of the mesh are improperly defined.')

        vtkPts = vtk.vtkPoints()
        vtkPts.SetData(nps.numpy_to_vtk(ptsMat, deep=True))

        # Cells
        cellArray = [c for c in mesh]
        cellConn = np.array([cell.nodes for cell in cellArray])

        cellsMat = np.concatenate((np.ones((cellConn.shape[0], 1), dtype=np.int64)*cellConn.shape[1], cellConn), axis=1).ravel()
        cellsArr = vtk.vtkCellArray()
        cellsArr.SetNumberOfCells(cellConn.shape[0])
        cellsArr.SetCells(cellConn.shape[0], nps.numpy_to_vtkIdTypeArray(cellsMat, deep=True))

        # Make the object
        vtuObj = vtk.vtkUnstructuredGrid()
        vtuObj.SetPoints(vtkPts)
        vtuObj.SetCells(VTK_CELL_TYPE, cellsArr)
        # Add the level of refinement as a cell array
        cell_levels = np.array([cell._level for cell in cellArray])
        refineLevelArr = nps.numpy_to_vtk(cell_levels, deep=1)
        refineLevelArr.SetName('octreeLevel')
        vtuObj.GetCellData().AddArray(refineLevelArr)
        # Assign the model('s) to the object
        if models is not None:
            for name, mod in models.items():
                # Convert numpy array
                vtkDoubleArr = nps.numpy_to_vtk(mod, deep=1)
                vtkDoubleArr.SetName(name)
                vtuObj.GetCellData().AddArray(vtkDoubleArr)

        return vtuObj

    def __tensorMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkRectilinearGrid`` object of this tensor mesh and the
        given models as ``CellData`` of that grid.

        Input:
        :param mesh, discretize.TensorMesh - The tensor mesh to convert to a ``vtkRectilinearGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        # Deal with dimensionalities
        if mesh.dim >= 1:
            vX = mesh.vectorNx
            xD = mesh.nNx
            yD, zD = 1, 1
            vY, vZ = np.array([0, 0])
        if mesh.dim >= 2:
            vY = mesh.vectorNy
            yD = mesh.nNy
        if mesh.dim == 3:
            vZ = mesh.vectorNz
            zD = mesh.nNz
        # Use rectilinear VTK grid.
        # Assign the spatial information.
        vtkObj = vtk.vtkRectilinearGrid()
        vtkObj.SetDimensions(xD, yD, zD)
        vtkObj.SetXCoordinates(nps.numpy_to_vtk(vX, deep=1))
        vtkObj.SetYCoordinates(nps.numpy_to_vtk(vY, deep=1))
        vtkObj.SetZCoordinates(nps.numpy_to_vtk(vZ, deep=1))

        # Assign the model('s) to the object
        if models is not None:
            for item in models.items():
                # Convert numpy array
                vtkDoubleArr = nps.numpy_to_vtk(item[1], deep=1)
                vtkDoubleArr.SetName(item[0])
                vtkObj.GetCellData().AddArray(vtkDoubleArr)
            # Set the active scalar
            vtkObj.GetCellData().SetActiveScalars(list(models.keys())[0])
        return vtkObj

    def __curvilinearMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkStructuredGrid`` of this mesh and the given
        models as ``CellData`` of that object.

        Input:
        :param mesh, discretize.CurvilinearMesh - The curvilinear mesh to convert to a ``vtkStructuredGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        # Make the data parts for the vtu object
        # Points
        ptsMat = mesh.gridN

        dims = [mesh.nCx, mesh.nCy, mesh.nCz]
        # Adjust if result was 2D:
        if ptsMat.shape[1] == 2:
            # Figure out which dim is null
            nullDim = dims.index(None)
            ptsMat = np.insert(ptsMat, nullDim, np.zeros(ptsMat.shape[0]), axis=1)
        if ptsMat.shape[1] != 3:
            raise RuntimeError('Points of the mesh are improperly defined.')

        vtkPts = vtk.vtkPoints()
        vtkPts.SetData(nps.numpy_to_vtk(ptsMat, deep=True))

        dims = [mesh.nCx, mesh.nCy, mesh.nCz]
        for i, d in enumerate(dims):
            if d is None:
                dims[i] = 0
            dims[i] = dims[i] + 1

        output = vtk.vtkStructuredGrid()
        output.SetDimensions(dims[0], dims[1], dims[2]) # note this subtracts 1
        output.SetPoints(vtkPts)

        # Assign the model('s) to the object
        if models is not None:
            for name, mod in models.items():
                # Convert numpy array
                vtkDoubleArr = nps.numpy_to_vtk(mod, deep=1)
                vtkDoubleArr.SetName(name)
                output.GetCellData().AddArray(vtkDoubleArr)

        return output

    def __cylMeshToVTK(mesh, models=None):
        # TODO: implement this!
        pass

    def toVTK(mesh, models=None):
        """Convert any mesh object to it's proper VTK data object."""
        converters = {
            'TreeMesh' : vtkInterface.__treeMeshToVTK,
            'TensorMesh' : vtkInterface.__tensorMeshToVTK,
            'CurvilinearMesh' : vtkInterface.__curvilinearMeshToVTK,
            #TODO: 'CylMesh' : vtkInterface._cylMeshToVTK,
            }
        key = type(mesh).__name__
        try:
            convert = converters[key]
        except:
            raise RuntimeError('Mesh type `%s` is not currently supported for VTK conversion.' % key)
        return convert(mesh, models=models)

    @staticmethod
    def _saveUnstructuredGrid(fileName, vtkUnstructGrid, directory=''):
        """Saves a VTK unstructured grid file (vtu) for an already generated
        ``vtkUnstructuredGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkUnstructGrid, vtk.vtkUnstructuredGrid):
            raise RuntimeError('`_saveUnstructuredGrid` can only handle `vtkUnstructuredGrid` objects. `%s` is not supported.' % vtkUnstructGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vtu'
        elif ext not in '.vtu':
            raise IOError('{:s} is an incorrect extension, has to be .vtu'.format(ext))
        # Make the writer
        vtuWriteFilter = vtkXMLUnstructuredGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            vtuWriteFilter.SetInputDataObject(vtkUnstructGrid)
        else:
            vtuWriteFilter.SetInput(vtkUnstructGrid)
        vtuWriteFilter.SetFileName(fileName)
        # Write the file
        vtuWriteFilter.Update()

    @staticmethod
    def _saveStructuredGrid(fileName, vtkStructGrid, directory=''):
        """Saves a VTK structured grid file (vtk) for an already generated
        ``vtkStructuredGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkStructGrid, vtk.vtkStructuredGrid):
            raise RuntimeError('`_saveStructuredGrid` can only handle `vtkStructuredGrid` objects. `%s` is not supported.' % vtkStructGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vts'
        elif ext not in '.vts':
            raise IOError('{:s} is an incorrect extension, has to be .vts'.format(ext))
        # Make the writer
        writer = vtkXMLStructuredGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            writer.SetInputDataObject(vtkStructGrid)
        else:
            writer.SetInput(vtkStructGrid)
        writer.SetFileName(fileName)
        # Write the file
        writer.Update()

    @staticmethod
    def _saveRectilinearGrid(fileName, vtkRectGrid, directory=''):
        """Saves a VTK rectilinear file (vtr) ffor an already generated
        ``vtkRectilinearGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkRectGrid, vtk.vtkRectilinearGrid):
            raise RuntimeError('`_saveRectilinearGrid` can only handle `vtkRectilinearGrid` objects. `%s` is not supported.' % vtkRectGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vtr'
        elif ext not in '.vtr':
            raise IOError('{:s} is an incorrect extension, has to be .vtr'.format(ext))
        # Write the file.
        vtrWriteFilter = vtkXMLRectilinearGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            vtrWriteFilter.SetInputDataObject(vtkRectGrid)
        else:
            vtuWriteFilter.SetInput(vtuObj)
        vtrWriteFilter.SetFileName(fname)
        vtrWriteFilter.Update()

    def writeVTK(mesh, fileName, models=None, directory=''):
        """Makes and saves a VTK object from this mesh and models

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        :param dict models: dictionary of numpy.array - Name('s) and array('s).
        Match number of cells
        """
        vtkObj = vtkInterface.toVTK(mesh, models=models)
        writers = {
            'vtkUnstructuredGrid' : vtkInterface._saveUnstructuredGrid,
            'vtkRectilinearGrid' : vtkInterface._saveRectilinearGrid,
            'vtkStructuredGrid' : vtkInterface._saveStructuredGrid,
            #TODO: 'CylMesh' : vtkInterface.???,
            }
        key = type(vtkObj).__name__
        try:
            write = writers[key]
        except:
            raise RuntimeError('VTK data type `%s` is not currently supported.' % key)
        return write(fileName, vtkObj, directory=directory)

class vtkTensorRead(object):
    """Provides a convienance method for reading VTK Rectilinear Grid files
    as TensorMesh objects."""

    @classmethod
    def readVTK(TensorMesh, fileName, directory=''):
        """Read VTK Rectilinear (vtr xml file) and return Tensor mesh and model

        Input:
        :param str fileName: path to the vtr model file to read or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives

        Output:
        :rtype: tuple
        :return: (TensorMesh, modelDictionary)
        """
        fname = os.path.join(directory, fileName)
        # Read the file
        vtrReader = vtkXMLRectilinearGridReader()
        vtrReader.SetFileName(fname)
        vtrReader.Update()
        vtrGrid = vtrReader.GetOutput()
        # Sort information
        hx = np.abs(np.diff(nps.vtk_to_numpy(vtrGrid.GetXCoordinates())))
        xR = nps.vtk_to_numpy(vtrGrid.GetXCoordinates())[0]
        hy = np.abs(np.diff(nps.vtk_to_numpy(vtrGrid.GetYCoordinates())))
        yR = nps.vtk_to_numpy(vtrGrid.GetYCoordinates())[0]
        zD = np.diff(nps.vtk_to_numpy(vtrGrid.GetZCoordinates()))
        # Check the direction of hz
        if np.all(zD < 0):
            hz = np.abs(zD[::-1])
            zR = nps.vtk_to_numpy(vtrGrid.GetZCoordinates())[-1]
        else:
            hz = np.abs(zD)
            zR = nps.vtk_to_numpy(vtrGrid.GetZCoordinates())[0]
        x0 = np.array([xR, yR, zR])

        # Make the object
        tensMsh = TensorMesh([hx, hy, hz], x0=x0)

        # Grap the models
        models = {}
        for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()):
            modelName = vtrGrid.GetCellData().GetArrayName(i)
            if np.all(zD < 0):
                modFlip = nps.vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
                tM = tensMsh.r(modFlip, 'CC', 'CC', 'M')
                modArr = tensMsh.r(tM[:, :, ::-1], 'CC', 'CC', 'V')
            else:
                modArr = nps.vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
            models[modelName] = modArr

        # Return the data
        return tensMsh, models
<commit_msg>Cleanup the vtkInterface<commit_after>"""
The ``vtkInterface`` module provides an way for ``discretize`` meshes to be
converted to VTK data objects (and back when possible).

"""
import os
import numpy as np
import six

from ..utils import cyl2cart

import vtk
import vtk.util.numpy_support as nps
from vtk import VTK_VERSION
from vtk import vtkXMLRectilinearGridWriter
from vtk import vtkXMLUnstructuredGridWriter
from vtk import vtkXMLStructuredGridWriter
from vtk import vtkXMLRectilinearGridReader

class vtkInterface(object):
    """This class is full of methods that enable ``discretize`` meshes to
    be converted to VTK data objects (and back when possible).
    """
    # NOTE: I name wrangle the class specific VTK conversions to force the user
    #       to use the ``toVTK()`` method.

    @staticmethod
    def __assignCellData(vtkDS, models=None):
        """Assign the model(s) to the VTK dataset as CellData

        Input:
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells
        """
        nc = vtkDS.GetNumberOfCells()
        if models is not None:
            for name, mod in models.items():
                # Convert numpy array
                if mod.size != nc:
                    raise RuntimeError('Number of model cells ({}) does not match number of mesh cells ({}).'.format(mod.size, nc))
                vtkDoubleArr = nps.numpy_to_vtk(mod, deep=1)
                vtkDoubleArr.SetName(name)
                vtkDS.GetCellData().AddArray(vtkDoubleArr)
        return vtkDS

    def __treeMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkUnstructuredGrid`` object of this tree mesh and the
        given models as ``CellData`` of that VTK dataset.

        Input:
        :param mesh, discretize.TreeMesh - The tree mesh to convert to a ``vtkUnstructuredGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        # Make the data parts for the vtu object
        # Points
        ptsMat = np.vstack((mesh.gridN, mesh.gridhN))

        # Adjust if result was 2D (voxels are pixels in 2D):
        VTK_CELL_TYPE = vtk.VTK_VOXEL
        if ptsMat.shape[1] == 2:
            # TODO: assumes any 2D tree meshes are on the XY Plane
            # Add Z values of 0.0 if 2D
            ptsMat = np.c_[ptsMat, np.zeros(ptsMat.shape[0])]
            VTK_CELL_TYPE = vtk.VTK_PIXEL
        if ptsMat.shape[1] != 3:
            raise RuntimeError('Points of the mesh are improperly defined.')
        # Grab the points
        vtkPts = vtk.vtkPoints()
        vtkPts.SetData(nps.numpy_to_vtk(ptsMat, deep=True))
        # Cells
        cellArray = [c for c in mesh]
        cellConn = np.array([cell.nodes for cell in cellArray])
        cellsMat = np.concatenate((np.ones((cellConn.shape[0], 1), dtype=np.int64)*cellConn.shape[1], cellConn), axis=1).ravel()
        cellsArr = vtk.vtkCellArray()
        cellsArr.SetNumberOfCells(cellConn.shape[0])
        cellsArr.SetCells(cellConn.shape[0], nps.numpy_to_vtkIdTypeArray(cellsMat, deep=True))
        # Make the object
        output = vtk.vtkUnstructuredGrid()
        output.SetPoints(vtkPts)
        output.SetCells(VTK_CELL_TYPE, cellsArr)
        # Add the level of refinement as a cell array
        cell_levels = np.array([cell._level for cell in cellArray])
        refineLevelArr = nps.numpy_to_vtk(cell_levels, deep=1)
        refineLevelArr.SetName('octreeLevel')
        output.GetCellData().AddArray(refineLevelArr)
        # Assign the model('s) to the object
        return vtkInterface.__assignCellData(output, models=models)

    def __tensorMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkRectilinearGrid`` object of this tensor mesh and the
        given models as ``CellData`` of that grid.

        Input:
        :param mesh, discretize.TensorMesh - The tensor mesh to convert to a ``vtkRectilinearGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        # Deal with dimensionalities
        if mesh.dim >= 1:
            vX = mesh.vectorNx
            xD = mesh.nNx
            yD, zD = 1, 1
            vY, vZ = np.array([0, 0])
        if mesh.dim >= 2:
            vY = mesh.vectorNy
            yD = mesh.nNy
        if mesh.dim == 3:
            vZ = mesh.vectorNz
            zD = mesh.nNz
        # Use rectilinear VTK grid.
        # Assign the spatial information.
        output = vtk.vtkRectilinearGrid()
        output.SetDimensions(xD, yD, zD)
        output.SetXCoordinates(nps.numpy_to_vtk(vX, deep=1))
        output.SetYCoordinates(nps.numpy_to_vtk(vY, deep=1))
        output.SetZCoordinates(nps.numpy_to_vtk(vZ, deep=1))
        # Assign the model('s) to the object
        return vtkInterface.__assignCellData(output, models=models)

    @staticmethod
    def __createStructGrid(ptsMat, dims, models=None):
        # Adjust if result was 2D:
        if ptsMat.shape[1] == 2:
            # Figure out which dim is null
            nullDim = dims.index(None)
            ptsMat = np.insert(ptsMat, nullDim, np.zeros(ptsMat.shape[0]), axis=1)
        if ptsMat.shape[1] != 3:
            raise RuntimeError('Points of the mesh are improperly defined.')
        # Convert the points
        vtkPts = vtk.vtkPoints()
        vtkPts.SetData(nps.numpy_to_vtk(ptsMat, deep=True))
        # Uncover hidden dimension
        for i, d in enumerate(dims):
            if d is None:
                dims[i] = 0
            dims[i] = dims[i] + 1
        output = vtk.vtkStructuredGrid()
        output.SetDimensions(dims[0], dims[1], dims[2]) # note this subtracts 1
        output.SetPoints(vtkPts)
        # Assign the model('s) to the object
        return vtkInterface.__assignCellData(output, models=models)

    def __curvilinearMeshToVTK(mesh, models=None):
        """
        Constructs a ``vtkStructuredGrid`` of this mesh and the given
        models as ``CellData`` of that object.

        Input:
        :param mesh, discretize.CurvilinearMesh - The curvilinear mesh to convert to a ``vtkStructuredGrid``
        :param models, dictionary of numpy.array - Name('s) and array('s). Match number of cells

        """
        ptsMat = mesh.gridN
        dims = [mesh.nCx, mesh.nCy, mesh.nCz]
        return vtkInterface.__createStructGrid(ptsMat, dims, models=models)

    def __cylMeshToVTK(mesh, models=None):
        """This treats the CylindricalMesh defined in cylindrical coordinates
        :math:`(r, \theta, z)` and transforms it to cartesian coordinates.
        """
        # # Points
        # ptsMat = cyl2cart(mesh.gridN)
        # dims = [mesh.nCx, mesh.nCy, mesh.nCz]
        # return vtkInterface.__createStructGrid(ptsMat, dims, models=models)
        pass

    def toVTK(mesh, models=None):
        """Convert any mesh object to it's proper VTK data object."""
        converters = {
            'TreeMesh' : vtkInterface.__treeMeshToVTK,
            'TensorMesh' : vtkInterface.__tensorMeshToVTK,
            'CurvilinearMesh' : vtkInterface.__curvilinearMeshToVTK,
            #TODO: 'CylMesh' : vtkInterface.__cylMeshToVTK,
            }
        key = type(mesh).__name__
        try:
            convert = converters[key]
        except:
            raise RuntimeError('Mesh type `%s` is not currently supported for VTK conversion.' % key)
        return convert(mesh, models=models)

    @staticmethod
    def _saveUnstructuredGrid(fileName, vtkUnstructGrid, directory=''):
        """Saves a VTK unstructured grid file (vtu) for an already generated
        ``vtkUnstructuredGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkUnstructGrid, vtk.vtkUnstructuredGrid):
            raise RuntimeError('`_saveUnstructuredGrid` can only handle `vtkUnstructuredGrid` objects. `%s` is not supported.' % vtkUnstructGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vtu'
        elif ext not in '.vtu':
            raise IOError('{:s} is an incorrect extension, has to be .vtu'.format(ext))
        # Make the writer
        vtuWriteFilter = vtkXMLUnstructuredGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            vtuWriteFilter.SetInputDataObject(vtkUnstructGrid)
        else:
            vtuWriteFilter.SetInput(vtkUnstructGrid)
        vtuWriteFilter.SetFileName(fileName)
        # Write the file
        vtuWriteFilter.Update()

    @staticmethod
    def _saveStructuredGrid(fileName, vtkStructGrid, directory=''):
        """Saves a VTK structured grid file (vtk) for an already generated
        ``vtkStructuredGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkStructGrid, vtk.vtkStructuredGrid):
            raise RuntimeError('`_saveStructuredGrid` can only handle `vtkStructuredGrid` objects. `%s` is not supported.' % vtkStructGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vts'
        elif ext not in '.vts':
            raise IOError('{:s} is an incorrect extension, has to be .vts'.format(ext))
        # Make the writer
        writer = vtkXMLStructuredGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            writer.SetInputDataObject(vtkStructGrid)
        else:
            writer.SetInput(vtkStructGrid)
        writer.SetFileName(fileName)
        # Write the file
        writer.Update()

    @staticmethod
    def _saveRectilinearGrid(fileName, vtkRectGrid, directory=''):
        """Saves a VTK rectilinear file (vtr) ffor an already generated
        ``vtkRectilinearGrid`` object.

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        """
        if not isinstance(vtkRectGrid, vtk.vtkRectilinearGrid):
            raise RuntimeError('`_saveRectilinearGrid` can only handle `vtkRectilinearGrid` objects. `%s` is not supported.' % vtkRectGrid.__class__)
        # Check the extension of the fileName
        fname = os.path.join(directory, fileName)
        ext = os.path.splitext(fname)[1]
        if ext is '':
            fname = fname + '.vtr'
        elif ext not in '.vtr':
            raise IOError('{:s} is an incorrect extension, has to be .vtr'.format(ext))
        # Write the file.
        vtrWriteFilter = vtkXMLRectilinearGridWriter()
        if float(VTK_VERSION.split('.')[0]) >= 6:
            vtrWriteFilter.SetInputDataObject(vtkRectGrid)
        else:
            vtuWriteFilter.SetInput(vtuObj)
        vtrWriteFilter.SetFileName(fname)
        vtrWriteFilter.Update()

    def writeVTK(mesh, fileName, models=None, directory=''):
        """Makes and saves a VTK object from this mesh and models

        Input:
        :param str fileName:  path to the output vtk file or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives
        :param dict models: dictionary of numpy.array - Name('s) and array('s).
        Match number of cells
        """
        vtkObj = vtkInterface.toVTK(mesh, models=models)
        writers = {
            'vtkUnstructuredGrid' : vtkInterface._saveUnstructuredGrid,
            'vtkRectilinearGrid' : vtkInterface._saveRectilinearGrid,
            'vtkStructuredGrid' : vtkInterface._saveStructuredGrid,
            }
        key = type(vtkObj).__name__
        try:
            write = writers[key]
        except:
            raise RuntimeError('VTK data type `%s` is not currently supported.' % key)
        return write(fileName, vtkObj, directory=directory)

class vtkTensorRead(object):
    """Provides a convienance method for reading VTK Rectilinear Grid files
    as TensorMesh objects."""

    @classmethod
    def readVTK(TensorMesh, fileName, directory=''):
        """Read VTK Rectilinear (vtr xml file) and return Tensor mesh and model

        Input:
        :param str fileName: path to the vtr model file to read or just its name if directory is specified
        :param str directory: directory where the UBC GIF file lives

        Output:
        :rtype: tuple
        :return: (TensorMesh, modelDictionary)
        """
        fname = os.path.join(directory, fileName)
        # Read the file
        vtrReader = vtkXMLRectilinearGridReader()
        vtrReader.SetFileName(fname)
        vtrReader.Update()
        vtrGrid = vtrReader.GetOutput()
        # Sort information
        hx = np.abs(np.diff(nps.vtk_to_numpy(vtrGrid.GetXCoordinates())))
        xR = nps.vtk_to_numpy(vtrGrid.GetXCoordinates())[0]
        hy = np.abs(np.diff(nps.vtk_to_numpy(vtrGrid.GetYCoordinates())))
        yR = nps.vtk_to_numpy(vtrGrid.GetYCoordinates())[0]
        zD = np.diff(nps.vtk_to_numpy(vtrGrid.GetZCoordinates()))
        # Check the direction of hz
        if np.all(zD < 0):
            hz = np.abs(zD[::-1])
            zR = nps.vtk_to_numpy(vtrGrid.GetZCoordinates())[-1]
        else:
            hz = np.abs(zD)
            zR = nps.vtk_to_numpy(vtrGrid.GetZCoordinates())[0]
        x0 = np.array([xR, yR, zR])

        # Make the object
        tensMsh = TensorMesh([hx, hy, hz], x0=x0)

        # Grap the models
        models = {}
        for i in np.arange(vtrGrid.GetCellData().GetNumberOfArrays()):
            modelName = vtrGrid.GetCellData().GetArrayName(i)
            if np.all(zD < 0):
                modFlip = nps.vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
                tM = tensMsh.r(modFlip, 'CC', 'CC', 'M')
                modArr = tensMsh.r(tM[:, :, ::-1], 'CC', 'CC', 'V')
            else:
                modArr = nps.vtk_to_numpy(vtrGrid.GetCellData().GetArray(i))
            models[modelName] = modArr

        # Return the data
        return tensMsh, models
<|endoftext|>