<commit_before># -*- coding: utf-8 -*-
# vim: sw=4:ts=4:expandtab

"""
swutils
~~~~~~~

Provides methods for interacting with ScraperWiki boxes

Examples:
    Schedule a job::

        job = lambda x: 'hello %s' % x
        exception_handler = ExceptionHandler('reubano@gmail.com').handler
        run_or_schedule(job, True, exception_handler)

Attributes:
    SCHEDULE_TIME (str): Time of the day to run the scheduled job.
"""

from __future__ import (
    absolute_import, division, print_function, with_statement,
    unicode_literals)

import time
import smtplib
import logging
import itertools as it

import schedule as sch
import scraperwiki

from os import environ, path as p
from email.mime.text import MIMEText
from pprint import pprint
from operator import itemgetter
from functools import partial

from testfixtures import LogCapture
from tabutils import process as pr, fntools as ft
from sqlalchemy import Column, Integer
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.schema import MetaData
from sqlalchemy.orm import sessionmaker

__version__ = '0.8.0'

__title__ = 'swutils'
__author__ = 'Reuben Cummings'
__description__ = 'ScraperWiki box utility library'
__email__ = 'reubano@gmail.com'
__license__ = 'MIT'
__copyright__ = 'Copyright 2015 Reuben Cummings'

SCHEDULE_TIME = '10:30'
meta = MetaData()
logger = logging.getLogger('populate')

class ExceptionHandler(object):
    """Creates a logging exception handler with email notifications

    Note: the following doctests assume you have a running Postfix server
    https://www.garron.me/en/mac/postfix-relay-gmail-mac-os-x-local-smtp.html
    `sudo postfix start`
    """
    def __init__(self, to, logfile='log.txt', logname=''):
        """ExceptionHandler constructor

        Args:
            to (str): The email recipient
            logfile (str): The logfile (default: log.txt)
            logname (str): The logger name (default: '')

        Examples:
            >>> ExceptionHandler('reubano@gmail.com')  # doctest: +ELLIPSIS
            <swutils.ExceptionHandler object at 0x...>
        """
        logging.basicConfig(filename=logfile, level=logging.DEBUG)
        self.to = to
        self.logfile = logfile
        self.logger = logging.getLogger(logname)
        open(logfile, 'w').close() if not p.exists(logfile) else None

    def email(self, subject=None, text=None, host='localhost'):
        """Sends the email notification

        Args:
            subject (str): The email subject (default: localhost).
            text (str): The email content (default: None).
            host (str): The email host server (default: localhost).

        Examples:
            >>> to = 'reubano@gmail.com'
            >>> ExceptionHandler(to).email('hello world')  # doctest: +ELLIPSIS
            <smtplib.SMTP instance at 0x...>
        """
        user = environ.get('USER')
        body = 'https://scraperwiki.com/dataset/%s\n\n%s' % (user, text)
        msg = MIMEText(body)

        msg['From'] = '%s@scraperwiki.com' % user
        msg['Subject'] = subject or 'scraperwiki box %s failed' % user
        msg['To'] = self.to

        # Send the message via our own SMTP server, but don't include the
        # envelope header.
        s = smtplib.SMTP(host)
        # s.sendmail(msg['From'], [msg['To']], msg.as_string())
        s.quit()
        return s

    def handler(self, func):
        """Creates the exception handler

        Args:
            func (func): The func to catch exceptions from

        Examples:
            >>> import os
            >>> from tempfile import NamedTemporaryFile
            >>> f = NamedTemporaryFile(delete=False)
            >>> to = 'reubano@gmail.com'
            >>> exc_handler = ExceptionHandler(to, f.name).handler
            >>> job = exc_handler(lambda x: x * 2)
            >>> print(job(2))
            4
            >>> with LogCapture() as l:
            ...     job(None)
            ...     print(l)
            root ERROR
              unsupported operand type(s) for *: 'NoneType' and 'int'
            >>> os.unlink(f.name)

        Returns:
            func: the exception handler
        """
        def wrapper(*args, **kwargs):
            try:
                res = func(*args, **kwargs)
            except Exception as e:
                self.logger.exception(str(e))
                scraperwiki.status('error', 'Error collecting data')

                with open(self.logfile, 'rb') as f:
                    self.email(text=f.read())
            else:
                scraperwiki.status('ok')
                return res

        return wrapper

def run_or_schedule(job, schedule=False, exception_handler=None):
    """Runs a job and optionally schedules it to run later

    Args:
        job (func): The func to run
        schedule (bool): Schedule `func` to run in the future (default: False)
        exception_handler (func): The exception handler to wrap the function in
            (default: None)

    Examples:
        >>> job = partial(pprint, 'hello world')
        >>> run_or_schedule(job)
        u'hello world'
        >>> exception_handler = ExceptionHandler('reubano@gmail.com').handler
        >>> run_or_schedule(job, False, exception_handler)
        u'hello world'
    """
    if exception_handler and schedule:
        job = exception_handler(job)

    job()

    if schedule:
        sch.every(1).day.at(SCHEDULE_TIME).do(job)

        while True:
            sch.run_pending()
            time.sleep(1)

def get_message(count, name, deleted=True):
    """Generates a message on the number of records inserted or deleted

    Args:
        count (int): The number of records
        name (str): The table name
        deleted (bool): The exception handler to wrap the function in
            (default: True)

    Examples:
        >>> print(get_message(5, 'table'))
        Deleted 5 records from table `table`.
        >>> print(get_message(5, 'table', False))
        Inserted 5 records into table `table`.

    Returns
        str: The message
    """
    verb, prep = ('Deleted', 'from') if deleted else ('Inserted', 'into')
    return '%s %s records %s table `%s`.' % (verb, count, prep, name)

def execute(records, engine, table, rid=None):
    in_count = len(records)

    if rid:
        # delete records if already in db
        ids = map(itemgetter(rid), records)
        q = table.query.filter(getattr(table, rid).in_(ids))
        del_count = q.delete(synchronize_session=False)
        engine.session.commit()
    else:
        del_count = 0

    engine.execute(table.insert(), records)
    return del_count, in_count

def get_dynamic_res(engine, get_name, t, **kwargs):
    name, data = t
    f = get_name or unicode.lower
    table_name = f(name)
    Base = declarative_base()

    # dynamically create sqlalchemy table
    attrs = {'__tablename__': table_name}

    result = {
        'table': type(str(name), (Base, kwargs['mixin']), attrs).__table__,
        'rid': None,
        'data': data}

    return result

def res_from_models(models, gen_data, t, **kwargs):
    result = {
        'table': getattr(models, t.get('name').title()).__table__,
        'rid': t.get('rid'),
        'data': gen_data(**pr.merge([kwargs, t]))}

    return result

def res_from_meta(engine, gen_data, t, **kwargs):
    meta.reflect(engine)
    result = {
        'table': meta.tables[t.get('name')],
        'rid': t.get('rid'),
        'data': gen_data(**pr.merge([kwargs, t]))}

    return result

def delete_records(table, rid, engine):
    if not rid:
        # delete all records since there is no way to identify them
        table.query = engine.session.query(table)

        try:
            del_count = table.query.delete(synchronize_session=False)
        except OperationalError:
            table.create(engine)
            del_count = 0
        else:
            engine.session.commit()
    else:
        del_count = 0

    return del_count

def get_tables(gen_data, **kwargs):
    tables = kwargs.get('TABLES')

    if tables:
        dynamic = False
    else:
        dynamic = True
        data = gen_data(**kwargs)
        keyfunc = itemgetter(kwargs['KEY'])
        tables = it.groupby(sorted(data, key=keyfunc), keyfunc)

    return dynamic, tables

def populate(gen_data, engine, models=None, get_name=None, **kwargs):
    """Populates a SQLAlchemy db with data. Supports both declarative
    SQLAlchemy and Flask-SQLAlchemy

    Note: Either `TABLES` or `KEY` must be defined.

    Args:
        gen_data (func): A function used to generate the data to be inserted
            into the db. It will receive keywords comprised of combining
            `kwargs` with a table defined in `TABLES`.

        engine (obj): A SQLAlchemy engine.
        models (module): A models module of SQLAlchemy table classes
            (default: None).
        get_name (func): A function used to generate the table name if
            `TABLES` is unset. It will receive the name of each
            each grouped obtained by grouping the data generated from
            `gen_data` (default: None).
        kwargs (dict): Keyword arguments passed to `gen_data`.

    Kwargs:
        mixin (class): Base table that dynamically create tables inherit.
            Required if `TABLES` is unset.
        TABLES (list[dicts]): The table options. Required if `KEY` is unset.
        KEY (str): The field used to group data generated from `gen_data`.
            Required if `TABLES` is unset.
        ROW_LIMIT (int): The max total number of rows to process
        CHUNK_SIZE (int): The max number of rows to process at one time
        DEBUG (bool): Run in debug mode
        TESTING (bool): Run in test mode

    Examples:
        >>> # Test dynamic tables
        >>> from sqlalchemy import create_engine
        >>> class BaseMixin(object):
        ...    id = Column(Integer, primary_key=True)
        ...    value = Column(Integer)
        ...
        >>> meta = MetaData()
        >>> kwargs = {'KEY': 'kind', 'ROW_LIMIT': 4, 'mixin': BaseMixin}
        >>> f = lambda x: {'kind': 'odd' if x % 2 else 'even', 'value': x}
        >>> gen_data = lambda **x: map(f, range(15))
        >>> engine = create_engine('sqlite:///:memory:')
        >>> populate(gen_data, engine, **kwargs)
        >>> session = sessionmaker(engine)()
        >>> meta.reflect(engine)
        >>> tables = meta.sorted_tables
        >>> dict(session.query(tables[0]).all()) == {1: 0, 2: 2, 3: 4, 4: 6}
        True
        >>> dict(session.query(tables[1]).all()) == {1: 1, 2: 3, 3: 5, 4: 7}
        True
        >>> meta.drop_all(engine)
        >>>
        >>> # Test tables without specifying the `rid`
        >>> Base = declarative_base()
        >>> class Single(Base):
        ...     __tablename__ = 'single'
        ...     id = Column(Integer, primary_key=True)
        ...     rid = Column(Integer)
        ...     value = Column(Integer)
        ...
        >>> class Triple(Base):
        ...     __tablename__ = 'triple'
        ...     id = Column(Integer, primary_key=True)
        ...     rid = Column(Integer)
        ...     value = Column(Integer)
        ...
        >>> options = [
        ...     {'mul': 1, 'name': 'single'}, {'mul': 3, 'name': 'triple'}]
        >>> kwargs = {'TABLES': options, 'ROW_LIMIT': 4}
        >>> def gen_data(**x):
        ...     return ({'value': n * x['mul'], 'rid': n} for n in it.count())
        >>> Base.metadata.create_all(engine)
        >>> populate(gen_data, engine, **kwargs)
        >>> Base.metadata.reflect(engine)
        >>> tables = Base.metadata.sorted_tables
        >>> session.query(tables[0]).all()
        [(1, 0, 0), (2, 1, 1), (3, 2, 2), (4, 3, 3)]
        >>> session.query(tables[1]).all()
        [(1, 0, 0), (2, 1, 3), (3, 2, 6), (4, 3, 9)]
        >>>
        >>> # Test tables with a specified `rid`
        >>> populate(gen_data, engine, rid='rid', **kwargs)
        >>> Base.metadata.reflect(engine)
        >>> tables = Base.metadata.sorted_tables
        >>> session.query(tables[0]).all()
        [(1, 0, 0), (2, 1, 1), (3, 2, 2), (4, 3, 3)]
        >>> session.query(tables[1]).all()
        [(1, 0, 0), (2, 1, 3), (3, 2, 6), (4, 3, 9)]

    Returns
        str: The message
    """
    log_level = logging.DEBUG if kwargs.get('DEBUG') else logging.INFO
    logger.setLevel(log_level)
    console_handler = logging.StreamHandler()
    logger.addHandler(console_handler)
    test = kwargs.get('TESTING')
    row_limit = kwargs.get('ROW_LIMIT')
    chunk_size = min(row_limit or 'inf', kwargs.get('CHUNK_SIZE', row_limit))
    engine.session = sessionmaker(engine)()

    if test:
        meta.create_all(engine)

    dynamic, tables = get_tables(gen_data, **kwargs)

    if dynamic:
        result_func = partial(get_dynamic_res, engine, get_name, **kwargs)
    elif models:
        result_func = partial(res_from_models, models, gen_data, **kwargs)
    else:
        result_func = partial(res_from_meta, engine, gen_data, **kwargs)

    for t in tables:
        count = 0
        result = result_func(t)
        table, rid, data = result['table'], result['rid'], result['data']
        del_count = delete_records(table, rid, engine)
        logger.debug(get_message(del_count, table.name))

        for records in ft.chunk(data, chunk_size):
            del_count, in_count = execute(records, engine, table, rid)
            count += in_count
            logger.debug(get_message(del_count, table.name))
            logger.debug(get_message(in_count, table.name, False))

            if test:
                pprint(records)

            if row_limit and count >= row_limit:
                break

        logger.debug('Success! %s' % get_message(count, table.name, False))
<commit_msg>Bump to version 0.8.1<commit_after># -*- coding: utf-8 -*-
# vim: sw=4:ts=4:expandtab

"""
swutils
~~~~~~~

Provides methods for interacting with ScraperWiki boxes

Examples:
    Schedule a job::

        job = lambda x: 'hello %s' % x
        exception_handler = ExceptionHandler('reubano@gmail.com').handler
        run_or_schedule(job, True, exception_handler)

Attributes:
    SCHEDULE_TIME (str): Time of the day to run the scheduled job.
"""

from __future__ import (
    absolute_import, division, print_function, with_statement,
    unicode_literals)

import time
import smtplib
import logging
import itertools as it

import schedule as sch
import scraperwiki

from os import environ, path as p
from email.mime.text import MIMEText
from pprint import pprint
from operator import itemgetter
from functools import partial

from testfixtures import LogCapture
from tabutils import process as pr, fntools as ft
from sqlalchemy import Column, Integer
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.schema import MetaData
from sqlalchemy.orm import sessionmaker

__version__ = '0.8.1'

__title__ = 'swutils'
__author__ = 'Reuben Cummings'
__description__ = 'ScraperWiki box utility library'
__email__ = 'reubano@gmail.com'
__license__ = 'MIT'
__copyright__ = 'Copyright 2015 Reuben Cummings'

SCHEDULE_TIME = '10:30'
meta = MetaData()
logger = logging.getLogger('populate')

class ExceptionHandler(object):
    """Creates a logging exception handler with email notifications

    Note: the following doctests assume you have a running Postfix server
    https://www.garron.me/en/mac/postfix-relay-gmail-mac-os-x-local-smtp.html
    `sudo postfix start`
    """
    def __init__(self, to, logfile='log.txt', logname=''):
        """ExceptionHandler constructor

        Args:
            to (str): The email recipient
            logfile (str): The logfile (default: log.txt)
            logname (str): The logger name (default: '')

        Examples:
            >>> ExceptionHandler('reubano@gmail.com')  # doctest: +ELLIPSIS
            <swutils.ExceptionHandler object at 0x...>
        """
        logging.basicConfig(filename=logfile, level=logging.DEBUG)
        self.to = to
        self.logfile = logfile
        self.logger = logging.getLogger(logname)
        open(logfile, 'w').close() if not p.exists(logfile) else None

    def email(self, subject=None, text=None, host='localhost'):
        """Sends the email notification

        Args:
            subject (str): The email subject (default: localhost).
            text (str): The email content (default: None).
            host (str): The email host server (default: localhost).

        Examples:
            >>> to = 'reubano@gmail.com'
            >>> ExceptionHandler(to).email('hello world')  # doctest: +ELLIPSIS
            <smtplib.SMTP instance at 0x...>
        """
        user = environ.get('USER')
        body = 'https://scraperwiki.com/dataset/%s\n\n%s' % (user, text)
        msg = MIMEText(body)

        msg['From'] = '%s@scraperwiki.com' % user
        msg['Subject'] = subject or 'scraperwiki box %s failed' % user
        msg['To'] = self.to

        # Send the message via our own SMTP server, but don't include the
        # envelope header.
        s = smtplib.SMTP(host)
        # s.sendmail(msg['From'], [msg['To']], msg.as_string())
        s.quit()
        return s

    def handler(self, func):
        """Creates the exception handler

        Args:
            func (func): The func to catch exceptions from

        Examples:
            >>> import os
            >>> from tempfile import NamedTemporaryFile
            >>> f = NamedTemporaryFile(delete=False)
            >>> to = 'reubano@gmail.com'
            >>> exc_handler = ExceptionHandler(to, f.name).handler
            >>> job = exc_handler(lambda x: x * 2)
            >>> print(job(2))
            4
            >>> with LogCapture() as l:
            ...     job(None)
            ...     print(l)
            root ERROR
              unsupported operand type(s) for *: 'NoneType' and 'int'
            >>> os.unlink(f.name)

        Returns:
            func: the exception handler
        """
        def wrapper(*args, **kwargs):
            try:
                res = func(*args, **kwargs)
            except Exception as e:
                self.logger.exception(str(e))
                scraperwiki.status('error', 'Error collecting data')

                with open(self.logfile, 'rb') as f:
                    self.email(text=f.read())
            else:
                scraperwiki.status('ok')
                return res

        return wrapper

def run_or_schedule(job, schedule=False, exception_handler=None):
    """Runs a job and optionally schedules it to run later

    Args:
        job (func): The func to run
        schedule (bool): Schedule `func` to run in the future (default: False)
        exception_handler (func): The exception handler to wrap the function in
            (default: None)

    Examples:
        >>> job = partial(pprint, 'hello world')
        >>> run_or_schedule(job)
        u'hello world'
        >>> exception_handler = ExceptionHandler('reubano@gmail.com').handler
        >>> run_or_schedule(job, False, exception_handler)
        u'hello world'
    """
    if exception_handler and schedule:
        job = exception_handler(job)

    job()

    if schedule:
        sch.every(1).day.at(SCHEDULE_TIME).do(job)

        while True:
            sch.run_pending()
            time.sleep(1)

def get_message(count, name, deleted=True):
    """Generates a message on the number of records inserted or deleted

    Args:
        count (int): The number of records
        name (str): The table name
        deleted (bool): The exception handler to wrap the function in
            (default: True)

    Examples:
        >>> print(get_message(5, 'table'))
        Deleted 5 records from table `table`.
        >>> print(get_message(5, 'table', False))
        Inserted 5 records into table `table`.

    Returns
        str: The message
    """
    verb, prep = ('Deleted', 'from') if deleted else ('Inserted', 'into')
    return '%s %s records %s table `%s`.' % (verb, count, prep, name)

def execute(records, engine, table, rid=None):
    in_count = len(records)

    if rid:
        # delete records if already in db
        ids = map(itemgetter(rid), records)
        q = table.query.filter(getattr(table, rid).in_(ids))
        del_count = q.delete(synchronize_session=False)
        engine.session.commit()
    else:
        del_count = 0

    engine.execute(table.insert(), records)
    return del_count, in_count

def get_dynamic_res(engine, get_name, t, **kwargs):
    name, data = t
    f = get_name or unicode.lower
    table_name = f(name)
    Base = declarative_base()

    # dynamically create sqlalchemy table
    attrs = {'__tablename__': table_name}

    result = {
        'table': type(str(name), (Base, kwargs['mixin']), attrs).__table__,
        'rid': None,
        'data': data}

    return result

def res_from_models(models, gen_data, t, **kwargs):
    result = {
        'table': getattr(models, t.get('name').title()).__table__,
        'rid': t.get('rid'),
        'data': gen_data(**pr.merge([kwargs, t]))}

    return result

def res_from_meta(engine, gen_data, t, **kwargs):
    meta.reflect(engine)
    result = {
        'table': meta.tables[t.get('name')],
        'rid': t.get('rid'),
        'data': gen_data(**pr.merge([kwargs, t]))}

    return result

def delete_records(table, rid, engine):
    if not rid:
        # delete all records since there is no way to identify them
        table.query = engine.session.query(table)

        try:
            del_count = table.query.delete(synchronize_session=False)
        except OperationalError:
            table.create(engine)
            del_count = 0
        else:
            engine.session.commit()
    else:
        del_count = 0

    return del_count

def get_tables(gen_data, **kwargs):
    tables = kwargs.get('TABLES')

    if tables:
        dynamic = False
    else:
        dynamic = True
        data = gen_data(**kwargs)
        keyfunc = itemgetter(kwargs['KEY'])
        tables = it.groupby(sorted(data, key=keyfunc), keyfunc)

    return dynamic, tables

def populate(gen_data, engine, models=None, get_name=None, **kwargs):
    """Populates a SQLAlchemy db with data. Supports both declarative
    SQLAlchemy and Flask-SQLAlchemy

    Note: Either `TABLES` or `KEY` must be defined.

    Args:
        gen_data (func): A function used to generate the data to be inserted
            into the db. It will receive keywords comprised of combining
            `kwargs` with a table defined in `TABLES`.

        engine (obj): A SQLAlchemy engine.
        models (module): A models module of SQLAlchemy table classes
            (default: None).
        get_name (func): A function used to generate the table name if
            `TABLES` is unset. It will receive the name of each
            each grouped obtained by grouping the data generated from
            `gen_data` (default: None).
        kwargs (dict): Keyword arguments passed to `gen_data`.

    Kwargs:
        mixin (class): Base table that dynamically create tables inherit.
            Required if `TABLES` is unset.
        TABLES (list[dicts]): The table options. Required if `KEY` is unset.
        KEY (str): The field used to group data generated from `gen_data`.
            Required if `TABLES` is unset.
        ROW_LIMIT (int): The max total number of rows to process
        CHUNK_SIZE (int): The max number of rows to process at one time
        DEBUG (bool): Run in debug mode
        TESTING (bool): Run in test mode

    Examples:
        >>> # Test dynamic tables
        >>> from sqlalchemy import create_engine
        >>> class BaseMixin(object):
        ...    id = Column(Integer, primary_key=True)
        ...    value = Column(Integer)
        ...
        >>> meta = MetaData()
        >>> kwargs = {'KEY': 'kind', 'ROW_LIMIT': 4, 'mixin': BaseMixin}
        >>> f = lambda x: {'kind': 'odd' if x % 2 else 'even', 'value': x}
        >>> gen_data = lambda **x: map(f, range(15))
        >>> engine = create_engine('sqlite:///:memory:')
        >>> populate(gen_data, engine, **kwargs)
        >>> session = sessionmaker(engine)()
        >>> meta.reflect(engine)
        >>> tables = meta.sorted_tables
        >>> dict(session.query(tables[0]).all()) == {1: 0, 2: 2, 3: 4, 4: 6}
        True
        >>> dict(session.query(tables[1]).all()) == {1: 1, 2: 3, 3: 5, 4: 7}
        True
        >>> meta.drop_all(engine)
        >>>
        >>> # Test tables without specifying the `rid`
        >>> Base = declarative_base()
        >>> class Single(Base):
        ...     __tablename__ = 'single'
        ...     id = Column(Integer, primary_key=True)
        ...     rid = Column(Integer)
        ...     value = Column(Integer)
        ...
        >>> class Triple(Base):
        ...     __tablename__ = 'triple'
        ...     id = Column(Integer, primary_key=True)
        ...     rid = Column(Integer)
        ...     value = Column(Integer)
        ...
        >>> options = [
        ...     {'mul': 1, 'name': 'single'}, {'mul': 3, 'name': 'triple'}]
        >>> kwargs = {'TABLES': options, 'ROW_LIMIT': 4}
        >>> def gen_data(**x):
        ...     return ({'value': n * x['mul'], 'rid': n} for n in it.count())
        >>> Base.metadata.create_all(engine)
        >>> populate(gen_data, engine, **kwargs)
        >>> Base.metadata.reflect(engine)
        >>> tables = Base.metadata.sorted_tables
        >>> session.query(tables[0]).all()
        [(1, 0, 0), (2, 1, 1), (3, 2, 2), (4, 3, 3)]
        >>> session.query(tables[1]).all()
        [(1, 0, 0), (2, 1, 3), (3, 2, 6), (4, 3, 9)]
        >>>
        >>> # Test tables with a specified `rid`
        >>> populate(gen_data, engine, rid='rid', **kwargs)
        >>> Base.metadata.reflect(engine)
        >>> tables = Base.metadata.sorted_tables
        >>> session.query(tables[0]).all()
        [(1, 0, 0), (2, 1, 1), (3, 2, 2), (4, 3, 3)]
        >>> session.query(tables[1]).all()
        [(1, 0, 0), (2, 1, 3), (3, 2, 6), (4, 3, 9)]

    Returns
        str: The message
    """
    log_level = logging.DEBUG if kwargs.get('DEBUG') else logging.INFO
    logger.setLevel(log_level)
    console_handler = logging.StreamHandler()
    logger.addHandler(console_handler)
    test = kwargs.get('TESTING')
    row_limit = kwargs.get('ROW_LIMIT')
    chunk_size = min(row_limit or 'inf', kwargs.get('CHUNK_SIZE', row_limit))
    engine.session = sessionmaker(engine)()

    if test:
        meta.create_all(engine)

    dynamic, tables = get_tables(gen_data, **kwargs)

    if dynamic:
        result_func = partial(get_dynamic_res, engine, get_name, **kwargs)
    elif models:
        result_func = partial(res_from_models, models, gen_data, **kwargs)
    else:
        result_func = partial(res_from_meta, engine, gen_data, **kwargs)

    for t in tables:
        count = 0
        result = result_func(t)
        table, rid, data = result['table'], result['rid'], result['data']
        del_count = delete_records(table, rid, engine)
        logger.debug(get_message(del_count, table.name))

        for records in ft.chunk(data, chunk_size):
            del_count, in_count = execute(records, engine, table, rid)
            count += in_count
            logger.debug(get_message(del_count, table.name))
            logger.debug(get_message(in_count, table.name, False))

            if test:
                pprint(records)

            if row_limit and count >= row_limit:
                break

        logger.debug('Success! %s' % get_message(count, table.name, False))
<|endoftext|>


<commit_before>__module_name__ = "Discord Bot Bridge"
__module_author__ = "FichteFoll"
__module_version__ = "0.2.0"
__module_description__ = "Translates messages bridged from Discord into the native IRC protocol"

import re

import hexchat

# associates channels where this functionality is active with the nickname of the bot
BOT_MAP = {
    "#nanaone": "_dc_"
}

MODE_CHAR = "⇔"
REVERSE_COLOR = "\026"
ITALICS = "\035"

def iter_fill(iterable, n, fillvalue=None):
    """Iterate over iterable and append fillvalue until n values have been yielded."""
    for x in iterable:
        yield x
        n -= 1
    for _ in range(n):
        yield fillvalue

def is_user_in_channel(nick, context=hexchat):
    for user in context.get_list('users'):
        if hexchat.nickcmp(nick, user.nick) == 0:
            return True
    return False

def msg_cb(word, word_eol, event_name, attrs):
    # Trailing empty words are not included in the word list (here: mode_char, identified_text).
    # For our custom print event, we don't care about the mode_char
    # because we either override it mode char or don't do anything.
    colored_nick, text, _, identified_text = iter_fill(word, 4, "")
    nick = hexchat.strip(colored_nick)

    channel = hexchat.get_info('channel')
    # if channel in BOT_MAP:
    #     print("nick", repr(nick))
    #     print("nickcmp", nick, BOT_MAP[channel], hexchat.nickcmp(nick, BOT_MAP[channel]))
    #     print(".lower()==?", nick.lower() == BOT_MAP[channel].lower())

    if channel not in BOT_MAP:
        return hexchat.EAT_NONE
    elif hexchat.nickcmp(nick, BOT_MAP[channel]) != 0:
        return hexchat.EAT_NONE

    match = re.match(r"^<([^>]+)> (.*)", text)
    if not match:
        return hexchat.EAT_NONE
    original_nick, message = match.groups()

    if original_nick == hexchat.get_info('nick'):
        event_name = 'Your Message'
    else:
        if not is_user_in_channel(original_nick):
            # TODO test if emit_print also works
            spaceless_nick = original_nick.replace(" ", "_")
            print("making", spaceless_nick, "join")
            hexchat.command("RECV :{nick}!someone@discord.server JOIN {channel}"
                            .format(nick=spaceless_nick, channel=channel))

    # try to decect actions
    if message.startswith(REVERSE_COLOR) and message.endswith(REVERSE_COLOR):
        message = message.strip(REVERSE_COLOR)
        event_name = event_name.replace("Message", "Action").replace("Msg", "Action")

    # replace with proper escape code
    message = message.replace(REVERSE_COLOR, ITALICS)

    hexchat.emit_print(event_name, original_nick, message, MODE_CHAR, identified_text,
                       time=attrs.time)

    return hexchat.EAT_ALL

if __name__ == '__main__':
    for event in ('Channel Message', 'Channel Msg Hilight',
                  'Channel Action', 'Channel Action Hilight'):
        hexchat.hook_print_attrs(event, msg_cb, event, hexchat.PRI_HIGHEST)

    print(__module_name__, __module_version__, "loaded")
<commit_msg>Remove debug prints<commit_after>__module_name__ = "Discord Bot Bridge"
__module_author__ = "FichteFoll"
__module_version__ = "0.2.0"
__module_description__ = "Translates messages bridged from Discord into the native IRC protocol"

import re

import hexchat

# associates channels where this functionality is active with the nickname of the bot
BOT_MAP = {
    "#nanaone": "_dc_"
}

MODE_CHAR = "⇔"
REVERSE_COLOR = "\026"
ITALICS = "\035"

def iter_fill(iterable, n, fillvalue=None):
    """Iterate over iterable and append fillvalue until n values have been yielded."""
    for x in iterable:
        yield x
        n -= 1
    for _ in range(n):
        yield fillvalue

def is_user_in_channel(nick, context=hexchat):
    for user in context.get_list('users'):
        if hexchat.nickcmp(nick, user.nick) == 0:
            return True
    return False

def msg_cb(word, word_eol, event_name, attrs):
    # Trailing empty words are not included in the word list (here: mode_char, identified_text).
    # For our custom print event, we don't care about the mode_char
    # because we either override it mode char or don't do anything.
    colored_nick, text, _, identified_text = iter_fill(word, 4, "")
    nick = hexchat.strip(colored_nick)

    channel = hexchat.get_info('channel')
    if channel not in BOT_MAP:
        return hexchat.EAT_NONE
    elif hexchat.nickcmp(nick, BOT_MAP[channel]) != 0:
        return hexchat.EAT_NONE

    match = re.match(r"^<([^>]+)> (.*)", text)
    if not match:
        return hexchat.EAT_NONE
    original_nick, message = match.groups()

    if original_nick == hexchat.get_info('nick'):
        event_name = 'Your Message'
    else:
        if not is_user_in_channel(original_nick):
            # TODO test if emit_print also works
            spaceless_nick = original_nick.replace(" ", "_")
            hexchat.command("RECV :{nick}!someone@discord.server JOIN {channel}"
                            .format(nick=spaceless_nick, channel=channel))

    # try to decect actions
    if message.startswith(REVERSE_COLOR) and message.endswith(REVERSE_COLOR):
        message = message.strip(REVERSE_COLOR)
        event_name = event_name.replace("Message", "Action").replace("Msg", "Action")

    # replace with proper escape code
    message = message.replace(REVERSE_COLOR, ITALICS)

    hexchat.emit_print(event_name, original_nick, message, MODE_CHAR, identified_text,
                       time=attrs.time)

    return hexchat.EAT_ALL

if __name__ == '__main__':
    for event in ('Channel Message', 'Channel Msg Hilight',
                  'Channel Action', 'Channel Action Hilight'):
        hexchat.hook_print_attrs(event, msg_cb, event, hexchat.PRI_HIGHEST)

    print(__module_name__, __module_version__, "loaded")
<|endoftext|>


<commit_before>"""
Statespace Tools

Author: Chad Fulton
License: Simplified-BSD
"""
from __future__ import division, absolute_import, print_function

import numpy as np
import pandas as pd
from statsmodels.tools.data import _is_using_pandas
from . import (
    _statespace, _kalman_filter, _kalman_smoother, _simulation_smoother
)

try:
    from scipy.linalg.blas import find_best_blas_type
except ImportError:  # pragma: no cover
    # Shim for SciPy 0.11, derived from tag=0.11 scipy.linalg.blas
    _type_conv = {'f': 's', 'd': 'd', 'F': 'c', 'D': 'z', 'G': 'z'}

    def find_best_blas_type(arrays):
        dtype, index = max(
            [(ar.dtype, i) for i, ar in enumerate(arrays)])
        prefix = _type_conv.get(dtype.char, 'd')
        return (prefix, dtype, None)

prefix_dtype_map = {
    's': np.float32, 'd': np.float64, 'c': np.complex64, 'z': np.complex128
}
prefix_statespace_map = {
    's': _statespace.sStatespace, 'd': _statespace.dStatespace,
    'c': _statespace.cStatespace, 'z': _statespace.zStatespace
}
prefix_kalman_filter_map = {
    's': _kalman_filter.sKalmanFilter, 'd': _kalman_filter.dKalmanFilter,
    'c': _kalman_filter.cKalmanFilter, 'z': _kalman_filter.zKalmanFilter
}
prefix_kalman_smoother_map = {
    's': _kalman_smoother.sKalmanSmoother,
    'd': _kalman_smoother.dKalmanSmoother,
    'c': _kalman_smoother.cKalmanSmoother,
    'z': _kalman_smoother.zKalmanSmoother
}
prefix_simulation_smoother_map = {
    's': _simulation_smoother.sSimulationSmoother,
    'd': _simulation_smoother.dSimulationSmoother,
    'c': _simulation_smoother.cSimulationSmoother,
    'z': _simulation_smoother.zSimulationSmoother
}

def companion_matrix(polynomial):
    r"""
    Create a companion matrix

    Parameters
    ----------
    polynomial : array_like, optional.
        If an iterable, interpreted as the coefficients of the polynomial from
        which to form the companion matrix. Polynomial coefficients are in
        order of increasing degree. If an integer, the size of the companion
        matrix (the polynomial coefficients are then set to zeros).

    Returns
    -------
    companion_matrix : array

    Notes
    -----

    Returns a matrix of the form

    .. math::
        \begin{bmatrix}
            \phi_1 & 1      & 0 & \cdots & 0 \\
            \phi_2 & 0      & 1 &        & 0 \\
            \vdots &        &   & \ddots & 0 \\
                   &        &   &        & 1 \\
            \phi_n & 0      & 0 & \cdots & 0 \\
        \end{bmatrix}

    where some or all of the :math:`\phi_i` may be non-zero (if `polynomial` is
    None, then all are equal to zero).

    If the coefficients provided are :math:`(c_0, c_1, \dots, c_{n})`,
    then the companion matrix is an :math:`n \times n` matrix formed with the
    elements in the first column defined as
    :math:`\phi_i = -\frac{c_i}{c_0}, i \in 1, \dots, n`.
    """
    if isinstance(polynomial, int):
        n = polynomial
        polynomial = None
    else:
        n = len(polynomial)-1
        polynomial = np.asanyarray(polynomial)

    matrix = np.zeros((n, n))
    idx = np.diag_indices(n - 1)
    idx = (idx[0], idx[1] + 1)
    matrix[idx] = 1
    if polynomial is not None and n > 0:
        matrix[:, 0] = -polynomial[1:] / polynomial[0]
    return matrix

def diff(series, k_diff=1, k_seasonal_diff=None, k_seasons=1):
    """
    Difference a series simply and/or seasonally along the zero-th axis.

    Given a series (denoted :math:`y_t`), performs the differencing operation

    .. math::

        \Delta^d \Delta_s^D y_t

    where :math:`d = diff`, :math:`s = k\_seasons`, :math:`D = seasonal\_diff`,
    and :math:`\Delta` is the difference operator.

    Parameters
    ----------
    series : array_like
        The series to be differenced.
    diff : int, optional
        The number of simple differences to perform. Default is 1.
    seasonal_diff : int or None, optional
        The number of seasonal differences to perform. Default is no seasonal
        differencing.
    k_seasons : int, optional
        The seasonal lag. Default is 1. Unused if there is no seasonal
        differencing.

    Returns
    -------
    differenced : array
        The differenced array.
    """
    pandas = _is_using_pandas(series, None)
    differenced = np.asanyarray(series) if not pandas else series

    # Seasonal differencing
    if k_seasonal_diff is not None:
        while k_seasonal_diff > 0:
            if not pandas:
                differenced = (
                    differenced[k_seasons:] - differenced[:-k_seasons]
                )
            else:
                differenced = differenced.diff(k_seasons)[k_seasons:]
            k_seasonal_diff -= 1

    # Simple differencing
    if not pandas:
        differenced = np.diff(differenced, k_diff, axis=0)
    else:
        while k_diff > 0:
            differenced = differenced.diff()[1:]
            k_diff -= 1
    return differenced

def concat(series, axis=0, allow_mix=False):
    """
    Concatenate a set of series.

    Parameters
    ----------
    series : iterable
        An iterable of series to be concatenated
    axis : int, optional
        The axis along which to concatenate. Default is 1 (columns).
    allow_mix : bool
        Whether or not to allow a mix of pandas and non-pandas objects. Default
        is False. If true, the returned object is an ndarray, and additional
        pandas metadata (e.g. column names, indices, etc) is lost.

    Returns
    -------
    concatenated : array or pd.DataFrame
        The concatenated array. Will be a DataFrame if series are pandas
        objects.
    """
    is_pandas = np.r_[[_is_using_pandas(s, None) for s in series]]

    if np.all(is_pandas):
        concatenated = pd.concat(series, axis=axis)
    elif np.all(~is_pandas) or allow_mix:
        concatenated = np.concatenate(series, axis=axis)
    else:
        raise ValueError('Attempted to concatenate Pandas objects with'
                         ' non-Pandas objects with `allow_mix=False`.')

    return concatenated

def is_invertible(polynomial, threshold=1.):
    """
    Determine if a polynomial is invertible.

    Requires all roots of the polynomial lie inside the unit circle.

    Parameters
    ----------
    polynomial : array_like
        Coefficients of a polynomial, in order of increasing degree.
        For example, `polynomial=[1, -0.5]` corresponds to the polynomial
        :math:`1 - 0.5x` which has root :math:`2`.
    threshold : number
        Allowed threshold for `is_invertible` to return True. Default is 1.

    Notes
    -----

    If the coefficients provided are :math:`(c_0, c_1, \dots, c_n)`, then
    the corresponding polynomial is :math:`c_0 + c_1 L + \dots + c_n L^n`.

    There are three equivalent methods of determining if the polynomial
    represented by the coefficients is invertible:

    The first method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (1 - \lambda_1 L)
                 (1 - \lambda_2 L) \dots (1 - \lambda_n L)

    In order for :math:`C(L)` to be invertible, it must be that each factor
    :math:`(1 - \lambda_i L)` is invertible; the condition is then that
    :math:`|\lambda_i| < 1`, where :math:`\lambda_i` is a root of the
    polynomial.

    The second method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (L - \zeta_1 L) (L - \zeta_2) \dots (L - \zeta_3)

    The condition is now :math:`|\zeta_i| > 1`, where :math:`\zeta_i` is a root
    of the polynomial with reversed coefficients and
    :math:`\lambda_i = \frac{1}{\zeta_i}`.

    Finally, a companion matrix can be formed using the coefficients of the
    polynomial. Then the eigenvalues of that matrix give the roots of the
    polynomial. This last method is the one actually used.

    See Also
    --------
    dismalpy.ssm.tools.companion_matrix

    Examples
    --------
    >>> dp.ssm.is_invertible([0.5])
    True
    >>> dp.ssm.is_invertible([1])
    False
    """
    # First method:
    # np.all(np.abs(np.roots(np.r_[1, params])) < 1)
    # Second method:
    # np.all(np.abs(np.roots(np.r_[1, params][::-1])) > 1)
    # Final method:
    eigvals = np.linalg.eigvals(companion_matrix(polynomial))
    return np.all(np.abs(eigvals) < threshold)

def constrain_stationary_univariate(unconstrained):
    """
    Transform unconstrained parameters used by the optimizer to constrained
    parameters used in likelihood evaluation

    Parameters
    ----------
    unconstrained : array
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    Returns
    -------
    constrained : array
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    References
    ----------

    Monahan, John F. 1984.
    "A Note on Enforcing Stationarity in
    Autoregressive-moving Average Models."
    Biometrika 71 (2) (August 1): 403-404.
    """

    n = unconstrained.shape[0]
    y = np.zeros((n, n), dtype=unconstrained.dtype)
    r = unconstrained/((1 + unconstrained**2)**0.5)
    for k in range(n):
        for i in range(k):
            y[k, i] = y[k - 1, i] + r[k] * y[k - 1, k - i - 1]
        y[k, k] = r[k]
    return -y[n - 1, :]

def unconstrain_stationary_univariate(constrained):
    """
    Transform constrained parameters used in likelihood evaluation
    to unconstrained parameters used by the optimizer

    Parameters
    ----------
    constrained : array
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    Returns
    -------
    unconstrained : array
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    References
    ----------

    Monahan, John F. 1984.
    "A Note on Enforcing Stationarity in
    Autoregressive-moving Average Models."
    Biometrika 71 (2) (August 1): 403-404.
    """
    n = constrained.shape[0]
    y = np.zeros((n, n), dtype=constrained.dtype)
    y[n-1:] = -constrained
    for k in range(n-1, 0, -1):
        for i in range(k):
            y[k-1, i] = (y[k, i] - y[k, k]*y[k, k-i-1]) / (1 - y[k, k]**2)
    r = y.diagonal()
    x = r / ((1 - r**2)**0.5)
    return x

def validate_matrix_shape(name, shape, nrows, ncols, nobs):
    """
    Validate the shape of a possibly time-varing matrix, or raise an exception

    Parameters
    ----------
    name : str
        The name of the matrix being validated (used in exception messages)
    shape : array_like
        The shape of the matrix to be validated. May be of size 2 or (if
        the matrix is time-varying) 3.
    nrows : int
        The expected number of rows.
    ncols : int
        The expected number of columns.
    nobs : int
        The number of observations (used to validate the last dimension of a
        time-varying matrix)

    Raises
    ------
    ValueError
        If the matrix is not of the desired shape.
    """
    ndim = len(shape)

    # Enforce dimension
    if ndim not in [2, 3]:
        raise ValueError('Invalid value for %s matrix. Requires a'
                         ' 2- or 3-dimensional array, got %d dimensions' %
                         (name, ndim))
    # Enforce the shape of the matrix
    if not shape[0] == nrows:
        raise ValueError('Invalid dimensions for %s matrix: requires %d'
                         ' rows, got %d' % (name, nrows, shape[0]))
    if not shape[1] == ncols:
        raise ValueError('Invalid dimensions for %s matrix: requires %d'
                         ' columns, got %d' % (name, ncols, shape[1]))

    # If we don't yet know `nobs`, don't allow time-varying arrays
    if nobs is None and not (ndim == 2 or shape[-1] == 1):
        raise ValueError('Invalid dimensions for %s matrix: time-varying'
                         ' matrices cannot be given unless `nobs` is specified'
                         ' (implicitly when a dataset is bound or else set'
                         ' explicity)' % name)

    # Enforce time-varying array size
    if ndim == 3 and nobs is not None and not shape[-1] in [1, nobs]:
        raise ValueError('Invalid dimensions for time-varying %s'
                         ' matrix. Requires shape (*,*,%d), got %s' %
                         (name, nobs, str(shape)))

def validate_vector_shape(name, shape, nrows, nobs):
    """
    Validate the shape of a possibly time-varing vector, or raise an exception

    Parameters
    ----------
    name : str
        The name of the vector being validated (used in exception messages)
    shape : array_like
        The shape of the vector to be validated. May be of size 1 or (if
        the vector is time-varying) 2.
    nrows : int
        The expected number of rows (elements of the vector).
    nobs : int
        The number of observations (used to validate the last dimension of a
        time-varying vector)

    Raises
    ------
    ValueError
        If the vector is not of the desired shape.
    """
    ndim = len(shape)
    # Enforce dimension
    if ndim not in [1, 2]:
        raise ValueError('Invalid value for %s vector. Requires a'
                         ' 1- or 2-dimensional array, got %d dimensions' %
                         (name, ndim))
    # Enforce the shape of the vector
    if not shape[0] == nrows:
        raise ValueError('Invalid dimensions for %s vector: requires %d'
                         ' rows, got %d' % (name, nrows, shape[0]))

    # If we don't yet know `nobs`, don't allow time-varying arrays
    if nobs is None and not (ndim == 1 or shape[-1] == 1):
        raise ValueError('Invalid dimensions for %s vector: time-varying'
                         ' vectors cannot be given unless `nobs` is specified'
                         ' (implicitly when a dataset is bound or else set'
                         ' explicity)' % name)

    # Enforce time-varying array size
    if ndim == 2 and not shape[1] in [1, nobs]:
        raise ValueError('Invalid dimensions for time-varying %s'
                         ' vector. Requires shape (*,%d), got %s' %
                         (name, nobs, str(shape)))
<commit_msg>CLEAN: Typos, misc.<commit_after>"""
Statespace Tools

Author: Chad Fulton
License: Simplified-BSD
"""
from __future__ import division, absolute_import, print_function

import numpy as np
import pandas as pd
from statsmodels.tools.data import _is_using_pandas
from . import (
    _statespace, _kalman_filter, _kalman_smoother, _simulation_smoother
)

try:
    from scipy.linalg.blas import find_best_blas_type
except ImportError:  # pragma: no cover
    # Shim for SciPy 0.11, derived from tag=0.11 scipy.linalg.blas
    _type_conv = {'f': 's', 'd': 'd', 'F': 'c', 'D': 'z', 'G': 'z'}

    def find_best_blas_type(arrays):
        dtype, index = max(
            [(ar.dtype, i) for i, ar in enumerate(arrays)])
        prefix = _type_conv.get(dtype.char, 'd')
        return prefix, dtype, None

prefix_dtype_map = {
    's': np.float32, 'd': np.float64, 'c': np.complex64, 'z': np.complex128
}
prefix_statespace_map = {
    's': _statespace.sStatespace, 'd': _statespace.dStatespace,
    'c': _statespace.cStatespace, 'z': _statespace.zStatespace
}
prefix_kalman_filter_map = {
    's': _kalman_filter.sKalmanFilter, 'd': _kalman_filter.dKalmanFilter,
    'c': _kalman_filter.cKalmanFilter, 'z': _kalman_filter.zKalmanFilter
}
prefix_kalman_smoother_map = {
    's': _kalman_smoother.sKalmanSmoother,
    'd': _kalman_smoother.dKalmanSmoother,
    'c': _kalman_smoother.cKalmanSmoother,
    'z': _kalman_smoother.zKalmanSmoother
}
prefix_simulation_smoother_map = {
    's': _simulation_smoother.sSimulationSmoother,
    'd': _simulation_smoother.dSimulationSmoother,
    'c': _simulation_smoother.cSimulationSmoother,
    'z': _simulation_smoother.zSimulationSmoother
}

def companion_matrix(polynomial):
    r"""
    Create a companion matrix

    Parameters
    ----------
    polynomial : array_like, optional.
        If an iterable, interpreted as the coefficients of the polynomial from
        which to form the companion matrix. Polynomial coefficients are in
        order of increasing degree. If an integer, the size of the companion
        matrix (the polynomial coefficients are then set to zeros).

    Returns
    -------
    companion_matrix : array

    Notes
    -----

    Returns a matrix of the form

    .. math::
        \begin{bmatrix}
            \phi_1 & 1      & 0 & \cdots & 0 \\
            \phi_2 & 0      & 1 &        & 0 \\
            \vdots &        &   & \ddots & 0 \\
                   &        &   &        & 1 \\
            \phi_n & 0      & 0 & \cdots & 0 \\
        \end{bmatrix}

    where some or all of the :math:`\phi_i` may be non-zero (if `polynomial` is
    None, then all are equal to zero).

    If the coefficients provided are :math:`(c_0, c_1, \dots, c_{n})`,
    then the companion matrix is an :math:`n \times n` matrix formed with the
    elements in the first column defined as
    :math:`\phi_i = -\frac{c_i}{c_0}, i \in 1, \dots, n`.
    """
    if isinstance(polynomial, int):
        n = polynomial
        polynomial = None
    else:
        n = len(polynomial)-1
        polynomial = np.asanyarray(polynomial)

    matrix = np.zeros((n, n))
    idx = np.diag_indices(n - 1)
    idx = (idx[0], idx[1] + 1)
    matrix[idx] = 1
    if polynomial is not None and n > 0:
        matrix[:, 0] = -polynomial[1:] / polynomial[0]
    return matrix

def diff(series, k_diff=1, k_seasonal_diff=None, k_seasons=1):
    """
    Difference a series simply and/or seasonally along the zero-th axis.

    Given a series (denoted :math:`y_t`), performs the differencing operation

    .. math::

        \Delta^d \Delta_s^D y_t

    where :math:`d = diff`, :math:`s = k\_seasons`, :math:`D = seasonal\_diff`,
    and :math:`\Delta` is the difference operator.

    Parameters
    ----------
    series : array_like
        The series to be differenced.
    diff : int, optional
        The number of simple differences to perform. Default is 1.
    seasonal_diff : int or None, optional
        The number of seasonal differences to perform. Default is no seasonal
        differencing.
    k_seasons : int, optional
        The seasonal lag. Default is 1. Unused if there is no seasonal
        differencing.

    Returns
    -------
    differenced : array
        The differenced array.
    """
    pandas = _is_using_pandas(series, None)
    differenced = np.asanyarray(series) if not pandas else series

    # Seasonal differencing
    if k_seasonal_diff is not None:
        while k_seasonal_diff > 0:
            if not pandas:
                differenced = (
                    differenced[k_seasons:] - differenced[:-k_seasons]
                )
            else:
                differenced = differenced.diff(k_seasons)[k_seasons:]
            k_seasonal_diff -= 1

    # Simple differencing
    if not pandas:
        differenced = np.diff(differenced, k_diff, axis=0)
    else:
        while k_diff > 0:
            differenced = differenced.diff()[1:]
            k_diff -= 1
    return differenced

def concat(series, axis=0, allow_mix=False):
    """
    Concatenate a set of series.

    Parameters
    ----------
    series : iterable
        An iterable of series to be concatenated
    axis : int, optional
        The axis along which to concatenate. Default is 1 (columns).
    allow_mix : bool
        Whether or not to allow a mix of pandas and non-pandas objects. Default
        is False. If true, the returned object is an ndarray, and additional
        pandas metadata (e.g. column names, indices, etc) is lost.

    Returns
    -------
    concatenated : array or pd.DataFrame
        The concatenated array. Will be a DataFrame if series are pandas
        objects.
    """
    is_pandas = np.r_[[_is_using_pandas(s, None) for s in series]]

    if np.all(is_pandas):
        concatenated = pd.concat(series, axis=axis)
    elif np.all(~is_pandas) or allow_mix:
        concatenated = np.concatenate(series, axis=axis)
    else:
        raise ValueError('Attempted to concatenate Pandas objects with'
                         ' non-Pandas objects with `allow_mix=False`.')

    return concatenated

def is_invertible(polynomial, threshold=1.):
    """
    Determine if a polynomial is invertible.

    Requires all roots of the polynomial lie inside the unit circle.

    Parameters
    ----------
    polynomial : array_like
        Coefficients of a polynomial, in order of increasing degree.
        For example, `polynomial=[1, -0.5]` corresponds to the polynomial
        :math:`1 - 0.5x` which has root :math:`2`.
    threshold : number
        Allowed threshold for `is_invertible` to return True. Default is 1.

    Notes
    -----

    If the coefficients provided are :math:`(c_0, c_1, \dots, c_n)`, then
    the corresponding polynomial is :math:`c_0 + c_1 L + \dots + c_n L^n`.

    There are three equivalent methods of determining if the polynomial
    represented by the coefficients is invertible:

    The first method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (1 - \lambda_1 L)
                 (1 - \lambda_2 L) \dots (1 - \lambda_n L)

    In order for :math:`C(L)` to be invertible, it must be that each factor
    :math:`(1 - \lambda_i L)` is invertible; the condition is then that
    :math:`|\lambda_i| < 1`, where :math:`\lambda_i` is a root of the
    polynomial.

    The second method factorizes the polynomial into:

    .. math::

        C(L) & = c_0 + c_1 L + \dots + c_n L^n \\
             & = constant (L - \zeta_1 L) (L - \zeta_2) \dots (L - \zeta_3)

    The condition is now :math:`|\zeta_i| > 1`, where :math:`\zeta_i` is a root
    of the polynomial with reversed coefficients and
    :math:`\lambda_i = \frac{1}{\zeta_i}`.

    Finally, a companion matrix can be formed using the coefficients of the
    polynomial. Then the eigenvalues of that matrix give the roots of the
    polynomial. This last method is the one actually used.

    See Also
    --------
    dismalpy.ssm.tools.companion_matrix

    Examples
    --------
    >>> dp.ssm.is_invertible([0.5])
    True
    >>> dp.ssm.is_invertible([1])
    False
    """
    # First method:
    # np.all(np.abs(np.roots(np.r_[1, params])) < 1)
    # Second method:
    # np.all(np.abs(np.roots(np.r_[1, params][::-1])) > 1)
    # Final method:
    eigvals = np.linalg.eigvals(companion_matrix(polynomial))
    return np.all(np.abs(eigvals) < threshold)

def constrain_stationary_univariate(unconstrained):
    """
    Transform unconstrained parameters used by the optimizer to constrained
    parameters used in likelihood evaluation

    Parameters
    ----------
    unconstrained : array
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    Returns
    -------
    constrained : array
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    References
    ----------

    Monahan, John F. 1984.
    "A Note on Enforcing Stationarity in
    Autoregressive-moving Average Models."
    Biometrika 71 (2) (August 1): 403-404.
    """

    n = unconstrained.shape[0]
    y = np.zeros((n, n), dtype=unconstrained.dtype)
    r = unconstrained/((1 + unconstrained**2)**0.5)
    for k in range(n):
        for i in range(k):
            y[k, i] = y[k - 1, i] + r[k] * y[k - 1, k - i - 1]
        y[k, k] = r[k]
    return -y[n - 1, :]

def unconstrain_stationary_univariate(constrained):
    """
    Transform constrained parameters used in likelihood evaluation
    to unconstrained parameters used by the optimizer

    Parameters
    ----------
    constrained : array
        Constrained parameters of, e.g., an autoregressive or moving average
        component, to be transformed to arbitrary parameters used by the
        optimizer.

    Returns
    -------
    unconstrained : array
        Unconstrained parameters used by the optimizer, to be transformed to
        stationary coefficients of, e.g., an autoregressive or moving average
        component.

    References
    ----------

    Monahan, John F. 1984.
    "A Note on Enforcing Stationarity in
    Autoregressive-moving Average Models."
    Biometrika 71 (2) (August 1): 403-404.
    """
    n = constrained.shape[0]
    y = np.zeros((n, n), dtype=constrained.dtype)
    y[n-1:] = -constrained
    for k in range(n-1, 0, -1):
        for i in range(k):
            y[k-1, i] = (y[k, i] - y[k, k]*y[k, k-i-1]) / (1 - y[k, k]**2)
    r = y.diagonal()
    x = r / ((1 - r**2)**0.5)
    return x

def validate_matrix_shape(name, shape, nrows, ncols, nobs):
    """
    Validate the shape of a possibly time-varying matrix, or raise an exception

    Parameters
    ----------
    name : str
        The name of the matrix being validated (used in exception messages)
    shape : array_like
        The shape of the matrix to be validated. May be of size 2 or (if
        the matrix is time-varying) 3.
    nrows : int
        The expected number of rows.
    ncols : int
        The expected number of columns.
    nobs : int
        The number of observations (used to validate the last dimension of a
        time-varying matrix)

    Raises
    ------
    ValueError
        If the matrix is not of the desired shape.
    """
    ndim = len(shape)

    # Enforce dimension
    if ndim not in [2, 3]:
        raise ValueError('Invalid value for %s matrix. Requires a'
                         ' 2- or 3-dimensional array, got %d dimensions' %
                         (name, ndim))
    # Enforce the shape of the matrix
    if not shape[0] == nrows:
        raise ValueError('Invalid dimensions for %s matrix: requires %d'
                         ' rows, got %d' % (name, nrows, shape[0]))
    if not shape[1] == ncols:
        raise ValueError('Invalid dimensions for %s matrix: requires %d'
                         ' columns, got %d' % (name, ncols, shape[1]))

    # If we don't yet know `nobs`, don't allow time-varying arrays
    if nobs is None and not (ndim == 2 or shape[-1] == 1):
        raise ValueError('Invalid dimensions for %s matrix: time-varying'
                         ' matrices cannot be given unless `nobs` is specified'
                         ' (implicitly when a dataset is bound or else set'
                         ' explicity)' % name)

    # Enforce time-varying array size
    if ndim == 3 and nobs is not None and not shape[-1] in [1, nobs]:
        raise ValueError('Invalid dimensions for time-varying %s'
                         ' matrix. Requires shape (*,*,%d), got %s' %
                         (name, nobs, str(shape)))

def validate_vector_shape(name, shape, nrows, nobs):
    """
    Validate the shape of a possibly time-varying vector, or raise an exception

    Parameters
    ----------
    name : str
        The name of the vector being validated (used in exception messages)
    shape : array_like
        The shape of the vector to be validated. May be of size 1 or (if
        the vector is time-varying) 2.
    nrows : int
        The expected number of rows (elements of the vector).
    nobs : int
        The number of observations (used to validate the last dimension of a
        time-varying vector)

    Raises
    ------
    ValueError
        If the vector is not of the desired shape.
    """
    ndim = len(shape)
    # Enforce dimension
    if ndim not in [1, 2]:
        raise ValueError('Invalid value for %s vector. Requires a'
                         ' 1- or 2-dimensional array, got %d dimensions' %
                         (name, ndim))
    # Enforce the shape of the vector
    if not shape[0] == nrows:
        raise ValueError('Invalid dimensions for %s vector: requires %d'
                         ' rows, got %d' % (name, nrows, shape[0]))

    # If we don't yet know `nobs`, don't allow time-varying arrays
    if nobs is None and not (ndim == 1 or shape[-1] == 1):
        raise ValueError('Invalid dimensions for %s vector: time-varying'
                         ' vectors cannot be given unless `nobs` is specified'
                         ' (implicitly when a dataset is bound or else set'
                         ' explicity)' % name)

    # Enforce time-varying array size
    if ndim == 2 and not shape[1] in [1, nobs]:
        raise ValueError('Invalid dimensions for time-varying %s'
                         ' vector. Requires shape (*,%d), got %s' %
                         (name, nobs, str(shape)))
<|endoftext|>


<commit_before>import os

import enaml
from enaml.qt.qt_application import QtApplication
from PyQt4.QtCore import QFileSystemWatcher, QUrl
from atom.api import Atom, Unicode, observe, Typed, Property

from btsync import BTSync

# Directory where all synced sites will be stored. Each site will be synced to
# a directory whose name is the secret.
STORAGE_PATH = u'/Users/jack/Desktop/synced_secrets'

class SyncNet(Atom):

    # The text currently entered in the secret field.
    entered_secret = Unicode()

    # A list of all locally synced secrets.
    known_secrets = Property()

    # Instance of the BTSync API wrapper.
    btsync = Typed(BTSync, ())

    # The QUrl object referencing the currently displayed resource.
    # FIXME remove this QT dependency.
    html_path = Typed(QUrl, ())

    # Root path where all synced site directoryies are added.
    storage_path = Unicode(STORAGE_PATH)

    # The filesystem watcher that monitors all currently synced site
    # directories.
    _watcher = Typed(QFileSystemWatcher)

    ### Public Interface  #####################################################

    def init_secret(self, secret):
        """ Creates a new directry at `self.storage_path` and adds it to the
        BTSync object to be synced with the given `secret`. The secret is
        assumed valid.

        Parameters
        ----------
        secret : str
            BTSync secret string referencing a directory of html files. This
            secret is assumed to already exist on the network.

        Notes
        -----
        The newly created folder's name will be the given `secret`.

        """
        path = os.path.join(self.storage_path, secret)
        os.mkdir(path)
        self._watcher.addPath(path)
        self.btsync.add_folder(path, secret)

    def load_secret(self, secret):
        """ Display the HTML files referenced by the given secret in the View.
        If the secret is not synced locally, it will be initialized and synced.

        Parameters
        ----------
        secret : str
            BTSync secret string

        Raises
        ------
        RuntimeError if `secret` is invalid

        """
        if not self.is_valid_secret(secret):
            msg = 'Attempted to load invalid secret: {}'.format(secret)
            raise RuntimeError(msg)

        if secret not in self.known_secrets:
            self.init_secret(secret)

        path = os.path.join(self.storage_path, secret)
        path = os.path.join(path, 'index.html')
        self.html_path = QUrl(path)

    def is_valid_secret(self, secret):
        """ True if the given `secret` is a valid btsync secret string.
        """
        return 'error' not in self.btsync.get_secrets(secret)

    ### Observers  ############################################################

    @observe('entered_secret')
    def _secret_changed(self, change):
        """ Load the entered secret into the HTML View if it is valid.
        """
        secret = self.entered_secret.upper()
        if self.is_valid_secret(secret):
            self.load_secret(secret)

    def on_directory_changed(self, dirname):
        """ Slot connected to the `QFileSystemWatcher.directoryChanged` Signal.
        """
        # If the directory containing the currently loaded secret changes, it
        # is reloaded.
        _, secret = os.path.split(os.path.normpath(dirname))
        if secret == self.entered_secret:
            self.load_secret(secret)

    def on_link_clicked(self, link):
        """ Slot connected to the `QWebView.linkClicked` Signal.
        """
        print link.toString()
        if link.scheme() == 'sync':
            secret = link.host().upper()
            if self.is_valid_secret(secret):
                self.load_secret(secret)
            else:
                print 'failed to load: {}'.format(link.toString())
        else:
            self.html_path = link

    ### Default methods  ######################################################

    def _default__watcher(self):
        _watcher = QFileSystemWatcher()
        _watcher.directoryChanged.connect(self.on_directory_changed)
        return _watcher

    ### Property getters  #####################################################

    def _get_known_secrets(self):
        """ List of all locally synced secrets. Getter for known_secrets.
        """
        return os.listdir(self.storage_path)

if __name__ == '__main__':
    with enaml.imports():
        from sync_net_view import SyncNetView
    syncnet = SyncNet()
    app = QtApplication()
    view = SyncNetView(model=syncnet)
    view.show()
    app.start()
<commit_msg>clarify docs<commit_after>import os

import enaml
from enaml.qt.qt_application import QtApplication
from PyQt4.QtCore import QFileSystemWatcher, QUrl
from atom.api import Atom, Unicode, observe, Typed, Property

from btsync import BTSync

# Directory where all synced sites will be stored. Each site will be synced to
# a directory whose name is the secret.
STORAGE_PATH = u'/Users/jack/Desktop/synced_secrets'

class SyncNet(Atom):

    # The text currently entered in the secret field.
    entered_secret = Unicode()

    # A list of all locally synced secrets.
    known_secrets = Property()

    # Instance of the BTSync API wrapper.
    btsync = Typed(BTSync, ())

    # The QUrl object referencing the currently displayed resource. It must be
    # replaced wholesale for the UI to react.
    # FIXME remove this QT dependency.
    html_path = Typed(QUrl, ())

    # Root path where all synced site directoryies are added.
    storage_path = Unicode(STORAGE_PATH)

    # The filesystem watcher that monitors all currently synced site
    # directories.
    _watcher = Typed(QFileSystemWatcher)

    ### Public Interface  #####################################################

    def init_secret(self, secret):
        """ Creates a new directry at `self.storage_path` and adds it to the
        BTSync object to be synced with the given `secret`. The secret is
        assumed valid.

        Parameters
        ----------
        secret : str
            BTSync secret string referencing a directory of html files. This
            secret is assumed to already exist on the network.

        Notes
        -----
        The newly created folder's name will be the given `secret`.

        """
        path = os.path.join(self.storage_path, secret)
        os.mkdir(path)
        self._watcher.addPath(path)
        self.btsync.add_folder(path, secret)

    def load_secret(self, secret):
        """ Display the HTML files referenced by the given secret in the View.
        If the secret is not synced locally, it will be initialized and synced.

        Parameters
        ----------
        secret : str
            BTSync secret string

        Raises
        ------
        RuntimeError if `secret` is invalid

        """
        if not self.is_valid_secret(secret):
            msg = 'Attempted to load invalid secret: {}'.format(secret)
            raise RuntimeError(msg)

        if secret not in self.known_secrets:
            self.init_secret(secret)

        path = os.path.join(self.storage_path, secret)
        path = os.path.join(path, 'index.html')
        self.html_path = QUrl(path)

    def is_valid_secret(self, secret):
        """ True if the given `secret` is a valid btsync secret string.
        """
        return 'error' not in self.btsync.get_secrets(secret)

    ### Observers  ############################################################

    @observe('entered_secret')
    def _secret_changed(self, change):
        """ Load the entered secret into the HTML View if it is valid.
        """
        secret = self.entered_secret.upper()
        if self.is_valid_secret(secret):
            self.load_secret(secret)

    def on_directory_changed(self, dirname):
        """ Slot connected to the `QFileSystemWatcher.directoryChanged` Signal.
        """
        # If the directory containing the currently loaded secret changes, it
        # is reloaded.
        _, secret = os.path.split(os.path.normpath(dirname))
        if secret == self.entered_secret:
            self.load_secret(secret)

    def on_link_clicked(self, link):
        """ Slot connected to the `QWebView.linkClicked` Signal.
        """
        print link.toString()
        if link.scheme() == 'sync':
            secret = link.host().upper()
            if self.is_valid_secret(secret):
                self.load_secret(secret)
            else:
                print 'failed to load: {}'.format(link.toString())
        else:
            self.html_path = link

    ### Default methods  ######################################################

    def _default__watcher(self):
        _watcher = QFileSystemWatcher()
        _watcher.directoryChanged.connect(self.on_directory_changed)
        return _watcher

    ### Property getters  #####################################################

    def _get_known_secrets(self):
        """ List of all locally synced secrets. Getter for known_secrets.
        """
        return os.listdir(self.storage_path)

if __name__ == '__main__':
    with enaml.imports():
        from sync_net_view import SyncNetView
    syncnet = SyncNet()
    app = QtApplication()
    view = SyncNetView(model=syncnet)
    view.show()
    app.start()
<|endoftext|>


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

import os, sys, subprocess, datetime, json

def _get_kernel():
    kernel =  subprocess.check_output(['uname', '-r']).decode('utf-8').strip('\n')
    kernel = kernel.split('-')
    return kernel

def _get_fqdn():
    fqdn = subprocess.check_output(['hostname', '-f']).decode('utf-8').strip('\n')
    fqdn = fqdn.split('.')
    return fqdn

def _get_uptime():
    with open('/proc/uptime', 'r') as upfile:
        uptime = datetime.timedelta(seconds = float(upfile.readline().split()[0]))
    return {'days':uptime.days, 'hours':int(uptime.seconds/3600), 'mins':int(uptime.seconds%3600/60), 'secs':int(uptime.seconds%3600%60)}

def _get_date():
    dt = str(datetime.datetime.now()).split()
    dt[1] = dt[1].split('.')[0]
    return dt

def _get_ip_addr():
    addr = subprocess.check_output(['ip', 'addr']).decode('utf-8').strip('\n')
    addr = addr.split('\n')
    addrs = {}
    for line in addr:
        if line[0].isdigit():
            iface = line.split()[1].strip(':')
            addrs[iface] = {'ipv6':[], 'ipv4':[]}
        if 'inet6' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv6'].append(address)
        elif 'inet' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv4'].append(address)
    return addrs

def _get_ip_route():
    route = subprocess.check_output(['ip', 'route']).decode('utf-8').strip('\n')
    route = route.split('\n')
    routes = {}
    for line in route:
        info = line.split()
        name = info[0]
        routes[name] = {}
        for i in range(1,len(info)):
            if info[i] == 'dev':
                routes[name]['dev'] = info[i+1]
            elif info[i] == 'via':
                routes[name]['via'] = info[i+1]
            elif info[i] == 'src':
                routes[name]['src'] = info[i+1]
    return routes

def _get_users():
    if not os.path.isfile('/etc/passwd'):
        return NULL
    else:
        users = {}
        try:
            with open('/etc/passwd', 'r') as passwd:
                for line in passwd:
                    line = line.strip('\n').split(':')
                    if int(line[2]) >= 1000:
                        users[line[0]] = {'uid':int(line[2]),
                                          'gid':int(line[3]),
                                          'home':line[5],
                                          'shell':line[6]}
        except:
            return NULL
        return users

def _get_disks():
    try:
        diskdict = {}
        disks = subprocess.check_output(['lsblk', '-lnf']).decode('utf-8').strip('\n').split()
        return
    except:
        return

def _detect_distro():
    return

def _get_processes():
    return

def _get_hosts():
    return

def _get_dns():
    return

def _get_mounts():
    return

def _get_cpuinfo():
    if os.path.isfile('/proc/cpuinfo'):
        with open('/proc/cpuinfo', 'r') as cpuinfo:
            for line in cpuinfo:
                if 'model name' in line:
                    return line.strip('\n').split(':')[1].split()

def _get_meminfo():
    meminfo = subprocess.check_output(['free', '-h']).decode('utf-8').strip('\n').split()
    mem = meminfo.index("Mem:")
    swap = meminfo.index("Swap:")
    return {'total':meminfo[int(mem)+1],
            'used':meminfo[int(mem)+2],
            'swap_total':meminfo[int(swap)+1],
            'swap_used':meminfo[int(swap)+2]}

def _get_current_user():
    return os.getenv('USER')

def _get_systemd_services():
    return

def _get_initd_services():
    return

def _get_usb_dev():
    return

def _get_pci_dev():
    return

def _get_listening_ports():
    return

def full_print():
    print('Kernel:\t{0}-{1}'.format(*_get_kernel()))
    print('Host:\t{}'.format('.'.join(_get_fqdn())))
    print('Uptime:\t{days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))
    print('Date:\t{0} {1}'.format(*_get_date()))

    print('IP Addresses: ')
    ipinfo = _get_ip_addr()
    for iface in ipinfo:
        print('{}:'.format(iface))
        if not ipinfo[iface]['ipv4']:
            continue
        else:
            print('\tIPv4:')
            for addr in ipinfo[iface]['ipv4']:
                print('\t\t{}'.format('/'.join(addr)))
        if not ipinfo[iface]['ipv6']:
            continue
        else:
            print('\tIPv6:')
            for addr in ipinfo[iface]['ipv6']:
                print('\t\t{}'.format('/'.join(addr)))

    print('Routes:')
    iproute = _get_ip_route()
    for route in iproute:
        proute = '\t{}'.format(route)
        for key in iproute[route]:
            #proute = proute+' '+key+' '+iproute[route][key]
            proute = '{0} {1} {2}'.format(proute, key, iproute[route][key])
        print(proute)

    print('Users:')
    users = _get_users()
    for user in users:
        print('\t{0}: {1}'.format(user,
            'uid={uid}, gid={gid}, home={home}, shell={shell}'.format(**users[user])))

    print('CPU:\t{}'.format(' '.join(_get_cpuinfo())))
    print('Memory:\t{used}/{total} Swap: {swap_used}/{swap_total}'.format(**_get_meminfo()))

def short_print():
    print(' '.join(_get_date()))
    print('{0}@{1}'.format(_get_current_user(), _get_fqdn()[0]))
    print('-'.join(_get_kernel()))
    print('Up: {days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))

def get_json():
    jdb = {}
    jdb['kernel'] = _get_kernel()
    jdb['hostname'] = _get_fqdn()
    jdb['uptime'] = _get_uptime()
    jdb['date'] = _get_date()
    jdb['ipaddr'] = _get_ip_addr()
    jdb['iproute'] = _get_ip_route()
    jdb['users'] = _get_users()
    jdb['cpu'] = _get_cpuinfo()
    jdb['mem'] = _get_meminfo()
    return jdb

if __name__ == '__main__':
    # full_print()
    # print('==============================')
    # short_print()
    # with open('host.json', 'w') as outfile:
    #     json.dump(get_json(), outfile)
    _get_disks()
<commit_msg>added format comment<commit_after>#!/usr/bin/env python3

import os, sys, subprocess, datetime, json

def _get_kernel():
    kernel =  subprocess.check_output(['uname', '-r']).decode('utf-8').strip('\n')
    kernel = kernel.split('-')
    return kernel

def _get_fqdn():
    fqdn = subprocess.check_output(['hostname', '-f']).decode('utf-8').strip('\n')
    fqdn = fqdn.split('.')
    return fqdn

def _get_uptime():
    with open('/proc/uptime', 'r') as upfile:
        uptime = datetime.timedelta(seconds = float(upfile.readline().split()[0]))
    return {'days':uptime.days, 'hours':int(uptime.seconds/3600), 'mins':int(uptime.seconds%3600/60), 'secs':int(uptime.seconds%3600%60)}

def _get_date():
    dt = str(datetime.datetime.now()).split()
    dt[1] = dt[1].split('.')[0]
    return dt

def _get_ip_addr():
    addr = subprocess.check_output(['ip', 'addr']).decode('utf-8').strip('\n')
    addr = addr.split('\n')
    addrs = {}
    for line in addr:
        if line[0].isdigit():
            iface = line.split()[1].strip(':')
            addrs[iface] = {'ipv6':[], 'ipv4':[]}
        if 'inet6' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv6'].append(address)
        elif 'inet' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv4'].append(address)
    return addrs

def _get_ip_route():
    route = subprocess.check_output(['ip', 'route']).decode('utf-8').strip('\n')
    route = route.split('\n')
    routes = {}
    for line in route:
        info = line.split()
        name = info[0]
        routes[name] = {}
        for i in range(1,len(info)):
            if info[i] == 'dev':
                routes[name]['dev'] = info[i+1]
            elif info[i] == 'via':
                routes[name]['via'] = info[i+1]
            elif info[i] == 'src':
                routes[name]['src'] = info[i+1]
    return routes

def _get_users():
    if not os.path.isfile('/etc/passwd'):
        return NULL
    else:
        users = {}
        try:
            with open('/etc/passwd', 'r') as passwd:
                for line in passwd:
                    line = line.strip('\n').split(':')
                    if int(line[2]) >= 1000:
                        users[line[0]] = {'uid':int(line[2]),
                                          'gid':int(line[3]),
                                          'home':line[5],
                                          'shell':line[6]}
        except:
            return NULL
        return users

def _get_disks():
    try:
        diskdict = {}
        disks = subprocess.check_output(['lsblk', '-lnf']).decode('utf-8').strip('\n').split()
        return
        # Format:
        # {
        # 'sda': {
        #         'sda1': {
        #                 'fs':'vfat',
        #                 'mount':'/boot/efi'
        #                 },
        #         'sda2': {
        #                 'fs':'ext4',
        #                 'mount':'/'
        #                 }
        #         },
        # 'sdb': {
        #         'sdb1': {
        #                 'fs':'ntfs',
        #                 'mount':'/windows'
        #                 }
        #         }
        # }
    except:
        return

def _detect_distro():
    return

def _get_processes():
    return

def _get_hosts():
    return

def _get_dns():
    return

def _get_mounts():
    return

def _get_cpuinfo():
    if os.path.isfile('/proc/cpuinfo'):
        with open('/proc/cpuinfo', 'r') as cpuinfo:
            for line in cpuinfo:
                if 'model name' in line:
                    return line.strip('\n').split(':')[1].split()

def _get_meminfo():
    meminfo = subprocess.check_output(['free', '-h']).decode('utf-8').strip('\n').split()
    mem = meminfo.index("Mem:")
    swap = meminfo.index("Swap:")
    return {'total':meminfo[int(mem)+1],
            'used':meminfo[int(mem)+2],
            'swap_total':meminfo[int(swap)+1],
            'swap_used':meminfo[int(swap)+2]}

def _get_current_user():
    return os.getenv('USER')

def _get_systemd_services():
    return

def _get_initd_services():
    return

def _get_usb_dev():
    return

def _get_pci_dev():
    return

def _get_listening_ports():
    return

def full_print():
    print('Kernel:\t{0}-{1}'.format(*_get_kernel()))
    print('Host:\t{}'.format('.'.join(_get_fqdn())))
    print('Uptime:\t{days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))
    print('Date:\t{0} {1}'.format(*_get_date()))

    print('IP Addresses: ')
    ipinfo = _get_ip_addr()
    for iface in ipinfo:
        print('{}:'.format(iface))
        if not ipinfo[iface]['ipv4']:
            continue
        else:
            print('\tIPv4:')
            for addr in ipinfo[iface]['ipv4']:
                print('\t\t{}'.format('/'.join(addr)))
        if not ipinfo[iface]['ipv6']:
            continue
        else:
            print('\tIPv6:')
            for addr in ipinfo[iface]['ipv6']:
                print('\t\t{}'.format('/'.join(addr)))

    print('Routes:')
    iproute = _get_ip_route()
    for route in iproute:
        proute = '\t{}'.format(route)
        for key in iproute[route]:
            #proute = proute+' '+key+' '+iproute[route][key]
            proute = '{0} {1} {2}'.format(proute, key, iproute[route][key])
        print(proute)

    print('Users:')
    users = _get_users()
    for user in users:
        print('\t{0}: {1}'.format(user,
            'uid={uid}, gid={gid}, home={home}, shell={shell}'.format(**users[user])))

    print('CPU:\t{}'.format(' '.join(_get_cpuinfo())))
    print('Memory:\t{used}/{total} Swap: {swap_used}/{swap_total}'.format(**_get_meminfo()))

def short_print():
    print(' '.join(_get_date()))
    print('{0}@{1}'.format(_get_current_user(), _get_fqdn()[0]))
    print('-'.join(_get_kernel()))
    print('Up: {days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))

def get_json():
    jdb = {}
    jdb['kernel'] = _get_kernel()
    jdb['hostname'] = _get_fqdn()
    jdb['uptime'] = _get_uptime()
    jdb['date'] = _get_date()
    jdb['ipaddr'] = _get_ip_addr()
    jdb['iproute'] = _get_ip_route()
    jdb['users'] = _get_users()
    jdb['cpu'] = _get_cpuinfo()
    jdb['mem'] = _get_meminfo()
    return jdb

if __name__ == '__main__':
    # full_print()
    # print('==============================')
    # short_print()
    # with open('host.json', 'w') as outfile:
    #     json.dump(get_json(), outfile)
    _get_disks()
<|endoftext|>


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

import os, sys, subprocess, datetime, json

def _get_kernel():
    kernel =  subprocess.check_output(['uname', '-r']).decode('utf-8').strip('\n')
    kernel = kernel.split('-')
    return kernel

def _get_fqdn():
    fqdn = subprocess.check_output(['hostname', '-f']).decode('utf-8').strip('\n')
    fqdn = fqdn.split('.')
    return fqdn

def _get_uptime():
    with open('/proc/uptime', 'r') as upfile:
        uptime = datetime.timedelta(seconds = float(upfile.readline().split()[0]))
    return {'days':uptime.days, 'hours':int(uptime.seconds/3600), 'mins':int(uptime.seconds%3600/60), 'secs':int(uptime.seconds%3600%60)}

def _get_date():
    dt = str(datetime.datetime.now()).split()
    dt[1] = dt[1].split('.')[0]
    return dt

def _get_ip_addr():
    addr = subprocess.check_output(['ip', 'addr']).decode('utf-8').strip('\n')
    addr = addr.split('\n')
    addrs = {}
    for line in addr:
        if line[0].isdigit():
            iface = line.split()[1].strip(':')
            addrs[iface] = {'ipv6':[], 'ipv4':[]}
        if 'inet6' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv6'].append(address)
        elif 'inet' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv4'].append(address)
    return addrs

def _get_ip_route():
    route = subprocess.check_output(['ip', 'route']).decode('utf-8').strip('\n')
    route = route.split('\n')
    routes = {}
    for line in route:
        info = line.split()
        name = info[0]
        routes[name] = {}
        for i in range(1,len(info)):
            if info[i] == 'dev':
                routes[name]['dev'] = info[i+1]
            elif info[i] == 'via':
                routes[name]['via'] = info[i+1]
            elif info[i] == 'src':
                routes[name]['src'] = info[i+1]
    return routes

def _get_users():
    if not os.path.isfile('/etc/passwd'):
        return NULL
    else:
        users = {}
        try:
            with open('/etc/passwd', 'r') as passwd:
                for line in passwd:
                    line = line.strip('\n').split(':')
                    if int(line[2]) >= 1000:
                        users[line[0]] = {'uid':int(line[2]),
                                          'gid':int(line[3]),
                                          'home':line[5],
                                          'shell':line[6]}
        except:
            return NULL
        return users

def _get_disks():
    try:
        diskdict = {}
        disks = subprocess.check_output(['lsblk', '-lnf']).decode('utf-8').strip('\n').split()
        return
        # Format:
        # {
        # 'sda': {
        #         'sda1': {
        #                 'fs':'vfat',
        #                 'mount':'/boot/efi'
        #                 },
        #         'sda2': {
        #                 'fs':'ext4',
        #                 'mount':'/'
        #                 }
        #         },
        # 'sdb': {
        #         'sdb1': {
        #                 'fs':'ntfs',
        #                 'mount':'/windows'
        #                 }
        #         }
        # }
    except:
        return

def _detect_distro():
    return

def _get_processes():
    return

def _get_hosts():
    return

def _get_dns():
    return

def _get_mounts():
    return

def _get_cpuinfo():
    if os.path.isfile('/proc/cpuinfo'):
        with open('/proc/cpuinfo', 'r') as cpuinfo:
            for line in cpuinfo:
                if 'model name' in line:
                    return line.strip('\n').split(':')[1].split()

def _get_meminfo():
    meminfo = subprocess.check_output(['free', '-h']).decode('utf-8').strip('\n').split()
    mem = meminfo.index("Mem:")
    swap = meminfo.index("Swap:")
    return {'total':meminfo[int(mem)+1],
            'used':meminfo[int(mem)+2],
            'swap_total':meminfo[int(swap)+1],
            'swap_used':meminfo[int(swap)+2]}

def _get_current_user():
    return os.getenv('USER')

def _get_systemd_services():
    return

def _get_initd_services():
    return

def _get_usb_dev():
    return

def _get_pci_dev():
    return

def _get_listening_ports():
    return

def full_print():
    print('Kernel:\t{0}-{1}'.format(*_get_kernel()))
    print('Host:\t{}'.format('.'.join(_get_fqdn())))
    print('Uptime:\t{days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))
    print('Date:\t{0} {1}'.format(*_get_date()))

    print('IP Addresses: ')
    ipinfo = _get_ip_addr()
    for iface in ipinfo:
        print('{}:'.format(iface))
        if not ipinfo[iface]['ipv4']:
            continue
        else:
            print('\tIPv4:')
            for addr in ipinfo[iface]['ipv4']:
                print('\t\t{}'.format('/'.join(addr)))
        if not ipinfo[iface]['ipv6']:
            continue
        else:
            print('\tIPv6:')
            for addr in ipinfo[iface]['ipv6']:
                print('\t\t{}'.format('/'.join(addr)))

    print('Routes:')
    iproute = _get_ip_route()
    for route in iproute:
        proute = '\t{}'.format(route)
        for key in iproute[route]:
            #proute = proute+' '+key+' '+iproute[route][key]
            proute = '{0} {1} {2}'.format(proute, key, iproute[route][key])
        print(proute)

    print('Users:')
    users = _get_users()
    for user in users:
        print('\t{0}: {1}'.format(user,
            'uid={uid}, gid={gid}, home={home}, shell={shell}'.format(**users[user])))

    print('CPU:\t{}'.format(' '.join(_get_cpuinfo())))
    print('Memory:\t{used}/{total} Swap: {swap_used}/{swap_total}'.format(**_get_meminfo()))

def short_print():
    print(' '.join(_get_date()))
    print('{0}@{1}'.format(_get_current_user(), _get_fqdn()[0]))
    print('-'.join(_get_kernel()))
    print('Up: {days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))

def get_json():
    jdb = {}
    jdb['kernel'] = _get_kernel()
    jdb['hostname'] = _get_fqdn()
    jdb['uptime'] = _get_uptime()
    jdb['date'] = _get_date()
    jdb['ipaddr'] = _get_ip_addr()
    jdb['iproute'] = _get_ip_route()
    jdb['users'] = _get_users()
    jdb['cpu'] = _get_cpuinfo()
    jdb['mem'] = _get_meminfo()
    return jdb

if __name__ == '__main__':
    # full_print()
    # print('==============================')
    # short_print()
    # with open('host.json', 'w') as outfile:
    #     json.dump(get_json(), outfile)
    _get_disks()
<commit_msg>fixed nfsnobody<commit_after>#!/usr/bin/env python3

import os, sys, subprocess, datetime, json

def _get_kernel():
    kernel =  subprocess.check_output(['uname', '-r']).decode('utf-8').strip('\n')
    kernel = kernel.split('-')
    return kernel

def _get_fqdn():
    fqdn = subprocess.check_output(['hostname', '-f']).decode('utf-8').strip('\n')
    fqdn = fqdn.split('.')
    return fqdn

def _get_uptime():
    with open('/proc/uptime', 'r') as upfile:
        uptime = datetime.timedelta(seconds = float(upfile.readline().split()[0]))
    return {'days':uptime.days, 'hours':int(uptime.seconds/3600), 'mins':int(uptime.seconds%3600/60), 'secs':int(uptime.seconds%3600%60)}

def _get_date():
    dt = str(datetime.datetime.now()).split()
    dt[1] = dt[1].split('.')[0]
    return dt

def _get_ip_addr():
    addr = subprocess.check_output(['ip', 'addr']).decode('utf-8').strip('\n')
    addr = addr.split('\n')
    addrs = {}
    for line in addr:
        if line[0].isdigit():
            iface = line.split()[1].strip(':')
            addrs[iface] = {'ipv6':[], 'ipv4':[]}
        if 'inet6' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv6'].append(address)
        elif 'inet' in line:
            address = line.split()[1].split('/')
            addrs[iface]['ipv4'].append(address)
    return addrs

def _get_ip_route():
    route = subprocess.check_output(['ip', 'route']).decode('utf-8').strip('\n')
    route = route.split('\n')
    routes = {}
    for line in route:
        info = line.split()
        name = info[0]
        routes[name] = {}
        for i in range(1,len(info)):
            if info[i] == 'dev':
                routes[name]['dev'] = info[i+1]
            elif info[i] == 'via':
                routes[name]['via'] = info[i+1]
            elif info[i] == 'src':
                routes[name]['src'] = info[i+1]
    return routes

def _get_users():
    if not os.path.isfile('/etc/passwd'):
        return NULL
    else:
        users = {}
        try:
            with open('/etc/passwd', 'r') as passwd:
                for line in passwd:
                    line = line.strip('\n').split(':')
                    if int(line[2]) >= 1000 and int(line[2]) < 65534:
                        users[line[0]] = {'uid':int(line[2]),
                                          'gid':int(line[3]),
                                          'home':line[5],
                                          'shell':line[6]}
        except:
            return NULL
        return users

def _get_disks():
    try:
        diskdict = {}
        disks = subprocess.check_output(['lsblk', '-lnf']).decode('utf-8').strip('\n').split()
        return
        # Format:
        # {
        # 'sda': {
        #         'sda1': {
        #                 'fs':'vfat',
        #                 'mount':'/boot/efi'
        #                 },
        #         'sda2': {
        #                 'fs':'ext4',
        #                 'mount':'/'
        #                 }
        #         },
        # 'sdb': {
        #         'sdb1': {
        #                 'fs':'ntfs',
        #                 'mount':'/windows'
        #                 }
        #         }
        # }
    except:
        return

def _detect_distro():
    return

def _get_processes():
    return

def _get_hosts():
    return

def _get_dns():
    return

def _get_mounts():
    return

def _get_cpuinfo():
    if os.path.isfile('/proc/cpuinfo'):
        with open('/proc/cpuinfo', 'r') as cpuinfo:
            for line in cpuinfo:
                if 'model name' in line:
                    return line.strip('\n').split(':')[1].split()

def _get_meminfo():
    meminfo = subprocess.check_output(['free', '-h']).decode('utf-8').strip('\n').split()
    mem = meminfo.index("Mem:")
    swap = meminfo.index("Swap:")
    return {'total':meminfo[int(mem)+1],
            'used':meminfo[int(mem)+2],
            'swap_total':meminfo[int(swap)+1],
            'swap_used':meminfo[int(swap)+2]}

def _get_current_user():
    return os.getenv('USER')

def _get_systemd_services():
    return

def _get_initd_services():
    return

def _get_usb_dev():
    return

def _get_pci_dev():
    return

def _get_listening_ports():
    return

def full_print():
    print('Kernel:\t{0}-{1}'.format(*_get_kernel()))
    print('Host:\t{}'.format('.'.join(_get_fqdn())))
    print('Uptime:\t{days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))
    print('Date:\t{0} {1}'.format(*_get_date()))

    print('IP Addresses: ')
    ipinfo = _get_ip_addr()
    for iface in ipinfo:
        print('{}:'.format(iface))
        if not ipinfo[iface]['ipv4']:
            continue
        else:
            print('\tIPv4:')
            for addr in ipinfo[iface]['ipv4']:
                print('\t\t{}'.format('/'.join(addr)))
        if not ipinfo[iface]['ipv6']:
            continue
        else:
            print('\tIPv6:')
            for addr in ipinfo[iface]['ipv6']:
                print('\t\t{}'.format('/'.join(addr)))

    print('Routes:')
    iproute = _get_ip_route()
    for route in iproute:
        proute = '\t{}'.format(route)
        for key in iproute[route]:
            #proute = proute+' '+key+' '+iproute[route][key]
            proute = '{0} {1} {2}'.format(proute, key, iproute[route][key])
        print(proute)

    print('Users:')
    users = _get_users()
    for user in users:
        print('\t{0}: {1}'.format(user,
            'uid={uid}, gid={gid}, home={home}, shell={shell}'.format(**users[user])))

    print('CPU:\t{}'.format(' '.join(_get_cpuinfo())))
    print('Memory:\t{used}/{total} Swap: {swap_used}/{swap_total}'.format(**_get_meminfo()))

def short_print():
    print(' '.join(_get_date()))
    print('{0}@{1}'.format(_get_current_user(), _get_fqdn()[0]))
    print('-'.join(_get_kernel()))
    print('Up: {days} Days, {hours} Hours, {mins} Minutes'.format(**_get_uptime()))

def get_json():
    jdb = {}
    jdb['kernel'] = _get_kernel()
    jdb['hostname'] = _get_fqdn()
    jdb['uptime'] = _get_uptime()
    jdb['date'] = _get_date()
    jdb['ipaddr'] = _get_ip_addr()
    jdb['iproute'] = _get_ip_route()
    jdb['users'] = _get_users()
    jdb['cpu'] = _get_cpuinfo()
    jdb['mem'] = _get_meminfo()
    return jdb

if __name__ == '__main__':
     full_print()
     print('==============================')
     short_print()
     with open('host.json', 'w') as outfile:
         json.dump(get_json(), outfile)
<|endoftext|>