<commit_before>"""
Support for recording details.

Component that records all events and state changes. Allows other components
to query this database.

For more details about this component, please refer to the documentation at
https://home-assistant.io/components/recorder/
"""
import asyncio
import concurrent.futures
import logging
from os import path
import queue
import threading
import time
from collections import namedtuple
from datetime import datetime, timedelta
from typing import Optional, Dict

import voluptuous as vol

from homeassistant.core import (
    HomeAssistant, callback, CoreState)
from homeassistant.const import (
    ATTR_ENTITY_ID, CONF_ENTITIES, CONF_EXCLUDE, CONF_DOMAINS,
    CONF_INCLUDE, EVENT_HOMEASSISTANT_STOP, EVENT_HOMEASSISTANT_START,
    EVENT_STATE_CHANGED, EVENT_TIME_CHANGED, MATCH_ALL)
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entityfilter import generate_filter
from homeassistant.helpers.typing import ConfigType
import homeassistant.util.dt as dt_util
from homeassistant import config as conf_util

from . import purge, migration
from .const import DATA_INSTANCE
from .util import session_scope

REQUIREMENTS = ['sqlalchemy==1.1.14']

_LOGGER = logging.getLogger(__name__)

DOMAIN = 'recorder'

SERVICE_PURGE = 'purge'

ATTR_KEEP_DAYS = 'keep_days'

SERVICE_PURGE_SCHEMA = vol.Schema({
    vol.Required(ATTR_KEEP_DAYS):
        vol.All(vol.Coerce(int), vol.Range(min=0))
})

DEFAULT_URL = 'sqlite:///{hass_config_path}'
DEFAULT_DB_FILE = 'home-assistant_v2.db'

CONF_DB_URL = 'db_url'
CONF_PURGE_KEEP_DAYS = 'purge_keep_days'
CONF_PURGE_INTERVAL = 'purge_interval'
CONF_EVENT_TYPES = 'event_types'

CONNECT_RETRY_WAIT = 3

FILTER_SCHEMA = vol.Schema({
    vol.Optional(CONF_EXCLUDE, default={}): vol.Schema({
        vol.Optional(CONF_ENTITIES, default=[]): cv.entity_ids,
        vol.Optional(CONF_DOMAINS, default=[]):
            vol.All(cv.ensure_list, [cv.string]),
        vol.Optional(CONF_EVENT_TYPES, default=[]):
            vol.All(cv.ensure_list, [cv.string])
    }),
    vol.Optional(CONF_INCLUDE, default={}): vol.Schema({
        vol.Optional(CONF_ENTITIES, default=[]): cv.entity_ids,
        vol.Optional(CONF_DOMAINS, default=[]):
            vol.All(cv.ensure_list, [cv.string])
    })
})

CONFIG_SCHEMA = vol.Schema({
    DOMAIN: FILTER_SCHEMA.extend({
        vol.Inclusive(CONF_PURGE_KEEP_DAYS, 'purge'):
            vol.All(vol.Coerce(int), vol.Range(min=1)),
        vol.Inclusive(CONF_PURGE_INTERVAL, 'purge'):
            vol.All(vol.Coerce(int), vol.Range(min=1)),
        vol.Optional(CONF_DB_URL): cv.string,
    })
}, extra=vol.ALLOW_EXTRA)

def wait_connection_ready(hass):
    """
    Wait till the connection is ready.

    Returns a coroutine object.
    """
    return hass.data[DATA_INSTANCE].async_db_ready

def run_information(hass, point_in_time: Optional[datetime]=None):
    """Return information about current run.

    There is also the run that covers point_in_time.
    """
    from . import models
    ins = hass.data[DATA_INSTANCE]

    recorder_runs = models.RecorderRuns
    if point_in_time is None or point_in_time > ins.recording_start:
        return ins.run_info

    with session_scope(hass=hass) as session:
        res = session.query(recorder_runs).filter(
            (recorder_runs.start < point_in_time) &
            (recorder_runs.end > point_in_time)).first()
        if res:
            session.expunge(res)
        return res

@asyncio.coroutine
def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
    """Set up the recorder."""
    conf = config.get(DOMAIN, {})
    keep_days = conf.get(CONF_PURGE_KEEP_DAYS)
    purge_interval = conf.get(CONF_PURGE_INTERVAL)

    db_url = conf.get(CONF_DB_URL, None)
    if not db_url:
        db_url = DEFAULT_URL.format(
            hass_config_path=hass.config.path(DEFAULT_DB_FILE))

    include = conf.get(CONF_INCLUDE, {})
    exclude = conf.get(CONF_EXCLUDE, {})
    instance = hass.data[DATA_INSTANCE] = Recorder(
        hass=hass, keep_days=keep_days, purge_interval=purge_interval,
        uri=db_url, include=include, exclude=exclude)
    instance.async_initialize()
    instance.start()

    @asyncio.coroutine
    def async_handle_purge_service(service):
        """Handle calls to the purge service."""
        instance.do_adhoc_purge(service.data[ATTR_KEEP_DAYS])

    descriptions = yield from hass.async_add_job(
        conf_util.load_yaml_config_file, path.join(
            path.dirname(__file__), 'services.yaml'))

    hass.services.async_register(DOMAIN, SERVICE_PURGE,
                                 async_handle_purge_service,
                                 descriptions.get(SERVICE_PURGE),
                                 schema=SERVICE_PURGE_SCHEMA)

    return (yield from instance.async_db_ready)

PurgeTask = namedtuple('PurgeTask', ['keep_days'])

class Recorder(threading.Thread):
    """A threaded recorder class."""

    def __init__(self, hass: HomeAssistant, keep_days: int,
                 purge_interval: int, uri: str,
                 include: Dict, exclude: Dict) -> None:
        """Initialize the recorder."""
        threading.Thread.__init__(self, name='Recorder')

        self.hass = hass
        self.keep_days = keep_days
        self.purge_interval = purge_interval
        self.queue = queue.Queue()  # type: Any
        self.recording_start = dt_util.utcnow()
        self.db_url = uri
        self.async_db_ready = asyncio.Future(loop=hass.loop)
        self.engine = None  # type: Any
        self.run_info = None  # type: Any

        self.entity_filter = generate_filter(include.get(CONF_DOMAINS, []),
                                             include.get(CONF_ENTITIES, []),
                                             exclude.get(CONF_DOMAINS, []),
                                             exclude.get(CONF_ENTITIES, []))
        self.exclude_t = exclude.get(CONF_EVENT_TYPES, [])

        self.get_session = None

    @callback
    def async_initialize(self):
        """Initialize the recorder."""
        self.hass.bus.async_listen(MATCH_ALL, self.event_listener)

    def do_adhoc_purge(self, keep_days):
        """Trigger an adhoc purge retaining keep_days worth of data."""
        if keep_days is not None:
            self.queue.put(PurgeTask(keep_days))

    def run(self):
        """Start processing events to save."""
        from .models import States, Events
        from homeassistant.components import persistent_notification
        from sqlalchemy import exc

        tries = 1
        connected = False

        while not connected and tries <= 10:
            if tries != 1:
                time.sleep(CONNECT_RETRY_WAIT)
            try:
                self._setup_connection()
                migration.migrate_schema(self)
                self._setup_run()
                connected = True
                _LOGGER.debug("Connected to recorder database")
            except Exception as err:  # pylint: disable=broad-except
                _LOGGER.error("Error during connection setup: %s (retrying "
                              "in %s seconds)", err, CONNECT_RETRY_WAIT)
                tries += 1

        if not connected:
            @callback
            def connection_failed():
                """Connect failed tasks."""
                self.async_db_ready.set_result(False)
                persistent_notification.async_create(
                    self.hass,
                    "The recorder could not start, please check the log",
                    "Recorder")

            self.hass.add_job(connection_failed)
            return

        shutdown_task = object()
        hass_started = concurrent.futures.Future()

        @callback
        def register():
            """Post connection initialize."""
            self.async_db_ready.set_result(True)

            def shutdown(event):
                """Shut down the Recorder."""
                if not hass_started.done():
                    hass_started.set_result(shutdown_task)
                self.queue.put(None)
                self.join()

            self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP,
                                            shutdown)

            if self.hass.state == CoreState.running:
                hass_started.set_result(None)
            else:
                @callback
                def notify_hass_started(event):
                    """Notify that hass has started."""
                    hass_started.set_result(None)

                self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START,
                                                notify_hass_started)

            if self.keep_days and self.purge_interval:
                async_track_point_in_time = \
                    self.hass.helpers.event.async_track_point_in_time

                @callback
                def async_purge(now):
                    """Trigger the purge and schedule the next run."""
                    self.queue.put(PurgeTask(self.keep_days))
                    async_track_point_in_time(async_purge, now + timedelta(
                        days=self.purge_interval))

                earliest = dt_util.utcnow() + timedelta(minutes=30)
                run = latest = dt_util.utcnow() + \
                    timedelta(days=self.purge_interval)
                with session_scope(session=self.get_session()) as session:
                    event = session.query(Events).first()
                    if event is not None:
                        session.expunge(event)
                        run = dt_util.UTC.localize(event.time_fired) + \
                            timedelta(days=self.keep_days+self.purge_interval)
                run = min(latest, max(run, earliest))

                _LOGGER.debug("Scheduling purge run for %s", run)
                async_track_point_in_time(async_purge, run)

        self.hass.add_job(register)
        result = hass_started.result()

        # If shutdown happened before HASS finished starting
        if result is shutdown_task:
            return

        while True:
            event = self.queue.get()

            if event is None:
                self._close_run()
                self._close_connection()
                self.queue.task_done()
                return
            elif isinstance(event, PurgeTask):
                purge.purge_old_data(self, event.keep_days)
                self.queue.task_done()
                continue
            elif event.event_type == EVENT_TIME_CHANGED:
                self.queue.task_done()
                continue
            elif event.event_type in self.exclude_t:
                self.queue.task_done()
                continue

            entity_id = event.data.get(ATTR_ENTITY_ID)
            if entity_id is not None:
                if not self.entity_filter(entity_id):
                    self.queue.task_done()
                    continue

            tries = 1
            updated = False
            while not updated and tries <= 10:
                if tries != 1:
                    time.sleep(CONNECT_RETRY_WAIT)
                try:
                    with session_scope(session=self.get_session()) as session:
                        dbevent = Events.from_event(event)
                        session.add(dbevent)

                        if event.event_type == EVENT_STATE_CHANGED:
                            dbstate = States.from_event(event)
                            dbstate.event_id = dbevent.event_id
                            session.add(dbstate)
                    updated = True

                except exc.OperationalError as err:
                    _LOGGER.error("Error in database connectivity: %s. "
                                  "(retrying in %s seconds)", err,
                                  CONNECT_RETRY_WAIT)
                    tries += 1

            if not updated:
                _LOGGER.error("Error in database update. Could not save "
                              "after %d tries. Giving up", tries)

            self.queue.task_done()

    @callback
    def event_listener(self, event):
        """Listen for new events and put them in the process queue."""
        self.queue.put(event)

    def block_till_done(self):
        """Block till all events processed."""
        self.queue.join()

    def _setup_connection(self):
        """Ensure database is ready to fly."""
        from sqlalchemy import create_engine, event
        from sqlalchemy.engine import Engine
        from sqlalchemy.orm import scoped_session
        from sqlalchemy.orm import sessionmaker
        from sqlite3 import Connection

        from . import models

        kwargs = {}

        # pylint: disable=unused-variable
        @event.listens_for(Engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            """Set sqlite's WAL mode."""
            if isinstance(dbapi_connection, Connection):
                old_isolation = dbapi_connection.isolation_level
                dbapi_connection.isolation_level = None
                cursor = dbapi_connection.cursor()
                cursor.execute("PRAGMA journal_mode=WAL")
                cursor.close()
                dbapi_connection.isolation_level = old_isolation

        if self.db_url == 'sqlite://' or ':memory:' in self.db_url:
            from sqlalchemy.pool import StaticPool

            kwargs['connect_args'] = {'check_same_thread': False}
            kwargs['poolclass'] = StaticPool
            kwargs['pool_reset_on_return'] = None
        else:
            kwargs['echo'] = False

        if self.engine is not None:
            self.engine.dispose()

        self.engine = create_engine(self.db_url, **kwargs)
        models.Base.metadata.create_all(self.engine)
        self.get_session = scoped_session(sessionmaker(bind=self.engine))

    def _close_connection(self):
        """Close the connection."""
        self.engine.dispose()
        self.engine = None
        self.get_session = None

    def _setup_run(self):
        """Log the start of the current run."""
        from .models import RecorderRuns

        with session_scope(session=self.get_session()) as session:
            for run in session.query(RecorderRuns).filter_by(end=None):
                run.closed_incorrect = True
                run.end = self.recording_start
                _LOGGER.warning("Ended unfinished session (id=%s from %s)",
                                run.run_id, run.start)
                session.add(run)

            self.run_info = RecorderRuns(
                start=self.recording_start,
                created=dt_util.utcnow()
            )
            session.add(self.run_info)
            session.flush()
            session.expunge(self.run_info)

    def _close_run(self):
        """Save end time for current run."""
        with session_scope(session=self.get_session()) as session:
            self.run_info.end = dt_util.utcnow()
            session.add(self.run_info)
        self.run_info = None
<commit_msg>Fix recorder purge (#10318)<commit_after>"""
Support for recording details.

Component that records all events and state changes. Allows other components
to query this database.

For more details about this component, please refer to the documentation at
https://home-assistant.io/components/recorder/
"""
import asyncio
import concurrent.futures
import logging
from os import path
import queue
import threading
import time
from collections import namedtuple
from datetime import datetime, timedelta
from typing import Optional, Dict

import voluptuous as vol

from homeassistant.core import (
    HomeAssistant, callback, CoreState)
from homeassistant.const import (
    ATTR_ENTITY_ID, CONF_ENTITIES, CONF_EXCLUDE, CONF_DOMAINS,
    CONF_INCLUDE, EVENT_HOMEASSISTANT_STOP, EVENT_HOMEASSISTANT_START,
    EVENT_STATE_CHANGED, EVENT_TIME_CHANGED, MATCH_ALL)
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entityfilter import generate_filter
from homeassistant.helpers.typing import ConfigType
import homeassistant.util.dt as dt_util
from homeassistant import config as conf_util

from . import purge, migration
from .const import DATA_INSTANCE
from .util import session_scope

REQUIREMENTS = ['sqlalchemy==1.1.14']

_LOGGER = logging.getLogger(__name__)

DOMAIN = 'recorder'

SERVICE_PURGE = 'purge'

ATTR_KEEP_DAYS = 'keep_days'

SERVICE_PURGE_SCHEMA = vol.Schema({
    vol.Required(ATTR_KEEP_DAYS):
        vol.All(vol.Coerce(int), vol.Range(min=0))
})

DEFAULT_URL = 'sqlite:///{hass_config_path}'
DEFAULT_DB_FILE = 'home-assistant_v2.db'

CONF_DB_URL = 'db_url'
CONF_PURGE_KEEP_DAYS = 'purge_keep_days'
CONF_PURGE_INTERVAL = 'purge_interval'
CONF_EVENT_TYPES = 'event_types'

CONNECT_RETRY_WAIT = 3

FILTER_SCHEMA = vol.Schema({
    vol.Optional(CONF_EXCLUDE, default={}): vol.Schema({
        vol.Optional(CONF_ENTITIES, default=[]): cv.entity_ids,
        vol.Optional(CONF_DOMAINS, default=[]):
            vol.All(cv.ensure_list, [cv.string]),
        vol.Optional(CONF_EVENT_TYPES, default=[]):
            vol.All(cv.ensure_list, [cv.string])
    }),
    vol.Optional(CONF_INCLUDE, default={}): vol.Schema({
        vol.Optional(CONF_ENTITIES, default=[]): cv.entity_ids,
        vol.Optional(CONF_DOMAINS, default=[]):
            vol.All(cv.ensure_list, [cv.string])
    })
})

CONFIG_SCHEMA = vol.Schema({
    DOMAIN: FILTER_SCHEMA.extend({
        vol.Inclusive(CONF_PURGE_KEEP_DAYS, 'purge'):
            vol.All(vol.Coerce(int), vol.Range(min=1)),
        vol.Inclusive(CONF_PURGE_INTERVAL, 'purge'):
            vol.All(vol.Coerce(int), vol.Range(min=1)),
        vol.Optional(CONF_DB_URL): cv.string,
    })
}, extra=vol.ALLOW_EXTRA)

def wait_connection_ready(hass):
    """
    Wait till the connection is ready.

    Returns a coroutine object.
    """
    return hass.data[DATA_INSTANCE].async_db_ready

def run_information(hass, point_in_time: Optional[datetime]=None):
    """Return information about current run.

    There is also the run that covers point_in_time.
    """
    from . import models
    ins = hass.data[DATA_INSTANCE]

    recorder_runs = models.RecorderRuns
    if point_in_time is None or point_in_time > ins.recording_start:
        return ins.run_info

    with session_scope(hass=hass) as session:
        res = session.query(recorder_runs).filter(
            (recorder_runs.start < point_in_time) &
            (recorder_runs.end > point_in_time)).first()
        if res:
            session.expunge(res)
        return res

@asyncio.coroutine
def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
    """Set up the recorder."""
    conf = config.get(DOMAIN, {})
    keep_days = conf.get(CONF_PURGE_KEEP_DAYS)
    purge_interval = conf.get(CONF_PURGE_INTERVAL)

    db_url = conf.get(CONF_DB_URL, None)
    if not db_url:
        db_url = DEFAULT_URL.format(
            hass_config_path=hass.config.path(DEFAULT_DB_FILE))

    include = conf.get(CONF_INCLUDE, {})
    exclude = conf.get(CONF_EXCLUDE, {})
    instance = hass.data[DATA_INSTANCE] = Recorder(
        hass=hass, keep_days=keep_days, purge_interval=purge_interval,
        uri=db_url, include=include, exclude=exclude)
    instance.async_initialize()
    instance.start()

    @asyncio.coroutine
    def async_handle_purge_service(service):
        """Handle calls to the purge service."""
        instance.do_adhoc_purge(service.data[ATTR_KEEP_DAYS])

    descriptions = yield from hass.async_add_job(
        conf_util.load_yaml_config_file, path.join(
            path.dirname(__file__), 'services.yaml'))

    hass.services.async_register(DOMAIN, SERVICE_PURGE,
                                 async_handle_purge_service,
                                 descriptions.get(SERVICE_PURGE),
                                 schema=SERVICE_PURGE_SCHEMA)

    return (yield from instance.async_db_ready)

PurgeTask = namedtuple('PurgeTask', ['keep_days'])

class Recorder(threading.Thread):
    """A threaded recorder class."""

    def __init__(self, hass: HomeAssistant, keep_days: int,
                 purge_interval: int, uri: str,
                 include: Dict, exclude: Dict) -> None:
        """Initialize the recorder."""
        threading.Thread.__init__(self, name='Recorder')

        self.hass = hass
        self.keep_days = keep_days
        self.purge_interval = purge_interval
        self.queue = queue.Queue()  # type: Any
        self.recording_start = dt_util.utcnow()
        self.db_url = uri
        self.async_db_ready = asyncio.Future(loop=hass.loop)
        self.engine = None  # type: Any
        self.run_info = None  # type: Any

        self.entity_filter = generate_filter(include.get(CONF_DOMAINS, []),
                                             include.get(CONF_ENTITIES, []),
                                             exclude.get(CONF_DOMAINS, []),
                                             exclude.get(CONF_ENTITIES, []))
        self.exclude_t = exclude.get(CONF_EVENT_TYPES, [])

        self.get_session = None

    @callback
    def async_initialize(self):
        """Initialize the recorder."""
        self.hass.bus.async_listen(MATCH_ALL, self.event_listener)

    def do_adhoc_purge(self, keep_days):
        """Trigger an adhoc purge retaining keep_days worth of data."""
        if keep_days is not None:
            self.queue.put(PurgeTask(keep_days))

    def run(self):
        """Start processing events to save."""
        from .models import States, Events
        from homeassistant.components import persistent_notification
        from sqlalchemy import exc

        tries = 1
        connected = False

        while not connected and tries <= 10:
            if tries != 1:
                time.sleep(CONNECT_RETRY_WAIT)
            try:
                self._setup_connection()
                migration.migrate_schema(self)
                self._setup_run()
                connected = True
                _LOGGER.debug("Connected to recorder database")
            except Exception as err:  # pylint: disable=broad-except
                _LOGGER.error("Error during connection setup: %s (retrying "
                              "in %s seconds)", err, CONNECT_RETRY_WAIT)
                tries += 1

        if not connected:
            @callback
            def connection_failed():
                """Connect failed tasks."""
                self.async_db_ready.set_result(False)
                persistent_notification.async_create(
                    self.hass,
                    "The recorder could not start, please check the log",
                    "Recorder")

            self.hass.add_job(connection_failed)
            return

        shutdown_task = object()
        hass_started = concurrent.futures.Future()

        @callback
        def register():
            """Post connection initialize."""
            self.async_db_ready.set_result(True)

            def shutdown(event):
                """Shut down the Recorder."""
                if not hass_started.done():
                    hass_started.set_result(shutdown_task)
                self.queue.put(None)
                self.join()

            self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP,
                                            shutdown)

            if self.hass.state == CoreState.running:
                hass_started.set_result(None)
            else:
                @callback
                def notify_hass_started(event):
                    """Notify that hass has started."""
                    hass_started.set_result(None)

                self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START,
                                                notify_hass_started)

            if self.keep_days and self.purge_interval:
                @callback
                def async_purge(now):
                    """Trigger the purge and schedule the next run."""
                    self.queue.put(PurgeTask(self.keep_days))
                    self.hass.helpers.event.async_track_point_in_time(
                        async_purge, now + timedelta(days=self.purge_interval))

                earliest = dt_util.utcnow() + timedelta(minutes=30)
                run = latest = dt_util.utcnow() + \
                    timedelta(days=self.purge_interval)
                with session_scope(session=self.get_session()) as session:
                    event = session.query(Events).first()
                    if event is not None:
                        session.expunge(event)
                        run = dt_util.as_utc(event.time_fired) + \
                            timedelta(days=self.keep_days+self.purge_interval)
                run = min(latest, max(run, earliest))
                self.hass.helpers.event.async_track_point_in_time(
                    async_purge, run)

        self.hass.add_job(register)
        result = hass_started.result()

        # If shutdown happened before HASS finished starting
        if result is shutdown_task:
            return

        while True:
            event = self.queue.get()

            if event is None:
                self._close_run()
                self._close_connection()
                self.queue.task_done()
                return
            elif isinstance(event, PurgeTask):
                purge.purge_old_data(self, event.keep_days)
                self.queue.task_done()
                continue
            elif event.event_type == EVENT_TIME_CHANGED:
                self.queue.task_done()
                continue
            elif event.event_type in self.exclude_t:
                self.queue.task_done()
                continue

            entity_id = event.data.get(ATTR_ENTITY_ID)
            if entity_id is not None:
                if not self.entity_filter(entity_id):
                    self.queue.task_done()
                    continue

            tries = 1
            updated = False
            while not updated and tries <= 10:
                if tries != 1:
                    time.sleep(CONNECT_RETRY_WAIT)
                try:
                    with session_scope(session=self.get_session()) as session:
                        dbevent = Events.from_event(event)
                        session.add(dbevent)

                        if event.event_type == EVENT_STATE_CHANGED:
                            dbstate = States.from_event(event)
                            dbstate.event_id = dbevent.event_id
                            session.add(dbstate)
                    updated = True

                except exc.OperationalError as err:
                    _LOGGER.error("Error in database connectivity: %s. "
                                  "(retrying in %s seconds)", err,
                                  CONNECT_RETRY_WAIT)
                    tries += 1

            if not updated:
                _LOGGER.error("Error in database update. Could not save "
                              "after %d tries. Giving up", tries)

            self.queue.task_done()

    @callback
    def event_listener(self, event):
        """Listen for new events and put them in the process queue."""
        self.queue.put(event)

    def block_till_done(self):
        """Block till all events processed."""
        self.queue.join()

    def _setup_connection(self):
        """Ensure database is ready to fly."""
        from sqlalchemy import create_engine, event
        from sqlalchemy.engine import Engine
        from sqlalchemy.orm import scoped_session
        from sqlalchemy.orm import sessionmaker
        from sqlite3 import Connection

        from . import models

        kwargs = {}

        # pylint: disable=unused-variable
        @event.listens_for(Engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            """Set sqlite's WAL mode."""
            if isinstance(dbapi_connection, Connection):
                old_isolation = dbapi_connection.isolation_level
                dbapi_connection.isolation_level = None
                cursor = dbapi_connection.cursor()
                cursor.execute("PRAGMA journal_mode=WAL")
                cursor.close()
                dbapi_connection.isolation_level = old_isolation

        if self.db_url == 'sqlite://' or ':memory:' in self.db_url:
            from sqlalchemy.pool import StaticPool

            kwargs['connect_args'] = {'check_same_thread': False}
            kwargs['poolclass'] = StaticPool
            kwargs['pool_reset_on_return'] = None
        else:
            kwargs['echo'] = False

        if self.engine is not None:
            self.engine.dispose()

        self.engine = create_engine(self.db_url, **kwargs)
        models.Base.metadata.create_all(self.engine)
        self.get_session = scoped_session(sessionmaker(bind=self.engine))

    def _close_connection(self):
        """Close the connection."""
        self.engine.dispose()
        self.engine = None
        self.get_session = None

    def _setup_run(self):
        """Log the start of the current run."""
        from .models import RecorderRuns

        with session_scope(session=self.get_session()) as session:
            for run in session.query(RecorderRuns).filter_by(end=None):
                run.closed_incorrect = True
                run.end = self.recording_start
                _LOGGER.warning("Ended unfinished session (id=%s from %s)",
                                run.run_id, run.start)
                session.add(run)

            self.run_info = RecorderRuns(
                start=self.recording_start,
                created=dt_util.utcnow()
            )
            session.add(self.run_info)
            session.flush()
            session.expunge(self.run_info)

    def _close_run(self):
        """Save end time for current run."""
        with session_scope(session=self.get_session()) as session:
            self.run_info.end = dt_util.utcnow()
            session.add(self.run_info)
        self.run_info = None
<|endoftext|>


<commit_before>__version__ = '1.10.73'

class ValidationTimeout(Exception):
    """Validation has timed out.

    May be replaced by the exception type raised by an external timeout
    handler when run in a server environment."""

    def __init__(self, timeout):
        self.timeout = timeout

    def __str__(self):
        return 'Validation timeout after %d seconds' % self.timeout
<commit_msg>Release 1.10.74<commit_after>__version__ = '1.10.74'

class ValidationTimeout(Exception):
    """Validation has timed out.

    May be replaced by the exception type raised by an external timeout
    handler when run in a server environment."""

    def __init__(self, timeout):
        self.timeout = timeout

    def __str__(self):
        return 'Validation timeout after %d seconds' % self.timeout
<|endoftext|>


<commit_before>"""Frequency Finder

Analyzes frequency of letters in given message compared to the most common occurring
letters to determine if message is in the English language.

Attributes:
    ETAOIN (str): String containing uppercase latin letters in order from most to least common.
    LETTERS (str): String containing uppercase latin letters in alphabetical order.

Note:
    * Compares six most and six least common letters in the English language.
    * https://www.nostarch.com/crackingcodes/ (BSD Licensed)
"""

ETAOIN = 'ETAOINSHRDLCUMWFGYPBVKJXQZ'
LETTERS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

def getLetterCount(message: str) -> dict:
    """Get letter count

    Counts the frequency of all latin letters in a given message.

    Args:
        message: String containing message to analyze letter frequency.

    Returns:
        Dictionary with keys of single letters and values of the count of how many
        times they appear in the message parameter.
    """
    letterCount = {'A': 0, 'B': 0, 'C': 0, 'D': 0, 'E': 0, 'F': 0,
                   'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0, 'L': 0, 'M': 0, 'N': 0,
                   'O': 0, 'P': 0, 'Q': 0, 'R': 0, 'S': 0, 'T': 0, 'U': 0, 'V': 0,
                   'W': 0, 'X': 0, 'Y': 0, 'Z': 0}

    for letter in message.upper():
        if letter in LETTERS:
            letterCount[letter] += 1

    return letterCount

def getItemAtIndexZero(items: tuple):
    """Get element at index zero

    Helper function that returns the first element of a given tuple.

    Args:
        items: Tuple containing a latin letter and its frequency count.

    Returns:
        The first element of the given tuple: the latin letter.
    """
    return items[0]

def getFrequencyOrder(message: str) -> str:
    """Get frequency order

    Analyzes frequency of each letter in given message and returns string with each letter from most to least frequent.

    Args:
         message: String containing message to analyze frequency.

    Returns:
        String of the alphabet letters arranged in order of most frequently occurring in the message parameter.
    """

    # First, get a dictionary of each letter and its frequency count:
    letterToFreq = getLetterCount(message)

    # Second, make a dictionary of each frequency count to the letter(s)
    # with that frequency:
    freqToLetter = {}
    for letter in LETTERS:
        if letterToFreq[letter] not in freqToLetter:
            freqToLetter[letterToFreq[letter]] = [letter]
        else:
            freqToLetter[letterToFreq[letter]].append(letter)

    # Third, put each list of letters in reverse "ETAOIN" order, and then
    # convert it to a string:
    for freq in freqToLetter:
        freqToLetter[freq].sort(key=ETAOIN.find, reverse=True)
        freqToLetter[freq] = ''.join(freqToLetter[freq])

    # Fourth, convert the freqToLetter dictionary to a list of
    # tuple pairs (key, value), and then sort them:
    freqPairs = list(freqToLetter.items())
    freqPairs.sort(key=getItemAtIndexZero, reverse=True)

    # Fifth, now that the letters are ordered by frequency, extract all
    # the letters for the final string:
    freqOrder = []
    for freqPair in freqPairs:
        freqOrder.append(freqPair[1])

    return ''.join(freqOrder)

def englishFreqMatchScore(message: str) -> int:
    """English Frequency Match Score

    Calculates number of matches that the string in the message parameter has when its letter frequency is
    compared to English letter frequency.

    Args:
         message: String containing message to calculate English match score.

    Returns:
        Number representing message's matches to English letter frequency.

    Note:
        * A "match" is how many of its six most frequent and six least frequent letters are among the six
        most frequent and six least frequent letters for English.
        * A "perfect score" is 12
    """
    freqOrder = getFrequencyOrder(message)

    matchScore = 0
    # Find how many matches for the six most common letters there are:
    for commonLetter in ETAOIN[:6]:
        if commonLetter in freqOrder[:6]:
            matchScore += 1
    # Find how many matches for the six least common letters there are:
    for uncommonLetter in ETAOIN[-6:]:
        if uncommonLetter in freqOrder[-6:]:
            matchScore += 1

    return matchScore
<commit_msg>Update Chapter19/freqAnalysis.py fixed autodoc warnings<commit_after>"""Frequency Finder

Analyzes frequency of letters in given message compared to the most common occurring
letters to determine if message is in the English language.

Attributes:
    ETAOIN (str): String containing uppercase latin letters in order from most to least common.
    LETTERS (str): String containing uppercase latin letters in alphabetical order.

Note:
    * Compares six most and six least common letters in the English language.
    * https://www.nostarch.com/crackingcodes/ (BSD Licensed)
"""

ETAOIN = 'ETAOINSHRDLCUMWFGYPBVKJXQZ'
LETTERS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

def getLetterCount(message: str) -> dict:
    """Get letter count

    Counts the frequency of all latin letters in a given message.

    Args:
        message: String containing message to analyze letter frequency.

    Returns:
        Dictionary with keys of single letters and values of the count of how many
        times they appear in the message parameter.
    """
    letterCount = {'A': 0, 'B': 0, 'C': 0, 'D': 0, 'E': 0, 'F': 0,
                   'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0, 'L': 0, 'M': 0, 'N': 0,
                   'O': 0, 'P': 0, 'Q': 0, 'R': 0, 'S': 0, 'T': 0, 'U': 0, 'V': 0,
                   'W': 0, 'X': 0, 'Y': 0, 'Z': 0}

    for letter in message.upper():
        if letter in LETTERS:
            letterCount[letter] += 1

    return letterCount

def getItemAtIndexZero(items: tuple):
    """Get element at index zero

    Helper function that returns the first element of a given tuple.

    Args:
        items: Tuple containing a latin letter and its frequency count.

    Returns:
        The first element of the given tuple: the latin letter.
    """
    return items[0]

def getFrequencyOrder(message: str) -> str:
    """Get frequency order

    Analyzes frequency of each letter in given message and returns string with each letter from most to least frequent.

    Args:
         message: String containing message to analyze frequency.

    Returns:
        String of the alphabet letters arranged in order of most frequently occurring in the message parameter.
    """

    # First, get a dictionary of each letter and its frequency count:
    letterToFreq = getLetterCount(message)

    # Second, make a dictionary of each frequency count to the letter(s)
    # with that frequency:
    freqToLetter = {}
    for letter in LETTERS:
        if letterToFreq[letter] not in freqToLetter:
            freqToLetter[letterToFreq[letter]] = [letter]
        else:
            freqToLetter[letterToFreq[letter]].append(letter)

    # Third, put each list of letters in reverse "ETAOIN" order, and then
    # convert it to a string:
    for freq in freqToLetter:
        freqToLetter[freq].sort(key=ETAOIN.find, reverse=True)
        freqToLetter[freq] = ''.join(freqToLetter[freq])

    # Fourth, convert the freqToLetter dictionary to a list of
    # tuple pairs (key, value), and then sort them:
    freqPairs = list(freqToLetter.items())
    freqPairs.sort(key=getItemAtIndexZero, reverse=True)

    # Fifth, now that the letters are ordered by frequency, extract all
    # the letters for the final string:
    freqOrder = []
    for freqPair in freqPairs:
        freqOrder.append(freqPair[1])

    return ''.join(freqOrder)

def englishFreqMatchScore(message: str) -> int:
    """English Frequency Match Score

    Calculates number of matches that the string in the message parameter has when its letter frequency is
    compared to English letter frequency.

    Args:
         message: String containing message to calculate English match score.

    Returns:
        Number representing message's matches to English letter frequency.

    Note:
        * A "match" is how many of its six most frequent and six least frequent letters are among the six
          most frequent and six least frequent letters for English.
        * A "perfect score" is 12

    """
    freqOrder = getFrequencyOrder(message)

    matchScore = 0
    # Find how many matches for the six most common letters there are:
    for commonLetter in ETAOIN[:6]:
        if commonLetter in freqOrder[:6]:
            matchScore += 1
    # Find how many matches for the six least common letters there are:
    for uncommonLetter in ETAOIN[-6:]:
        if uncommonLetter in freqOrder[-6:]:
            matchScore += 1

    return matchScore
<|endoftext|>


<commit_before># Copyright 2019 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Rewrite script for TF->JAX."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
from absl import app
from absl import flags

flags.DEFINE_boolean('numpy_to_jax', False,
                     'Whether or not to rewrite numpy imports to jax.numpy')
flags.DEFINE_list('omit_deps', [], 'List of build deps being omitted.')

FLAGS = flags.FLAGS

TF_REPLACEMENTS = {
    'import tensorflow ':
        'from tensorflow_probability.python.internal.backend import numpy ',
    'import tensorflow.compat.v1':
        'from tensorflow_probability.python.internal.backend.numpy.compat '
        'import v1',
    'import tensorflow.compat.v2':
        'from tensorflow_probability.python.internal.backend.numpy.compat '
        'import v2',
    'import tensorflow_probability as tfp':
        'import tensorflow_probability as tfp; '
        'tfp = tfp.substrates.numpy',
    'from tensorflow.python.framework import tensor_shape':
        ('from tensorflow_probability.python.internal.backend.numpy.gen '
         'import tensor_shape'),
    'from tensorflow.python.framework import ops':
        ('from tensorflow_probability.python.internal.backend.numpy '
         'import ops'),
    'from tensorflow.python.framework import tensor_util':
        ('from tensorflow_probability.python.internal.backend.numpy '
         'import ops'),
    'from tensorflow.python.util import':
        'from tensorflow_probability.python.internal.backend.numpy import',
    'from tensorflow.python.util.all_util':
        'from tensorflow_probability.python.internal.backend.numpy.private',
    'from tensorflow.python.ops.linalg':
        'from tensorflow_probability.python.internal.backend.numpy.gen',
    'from tensorflow.python.ops import parallel_for':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import functional_ops as parallel_for',
    'from tensorflow.python.ops import control_flow_ops':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import control_flow as control_flow_ops',
    'from tensorflow.python.eager import context':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import private',
    ('from tensorflow.python.client '
     'import pywrap_tf_session as c_api'):
        'pass',
    ('from tensorflow.python '
     'import pywrap_tensorflow as c_api'):
        'pass'
}

DISABLED_BY_PKG = {
    'experimental':
        ('auto_batching', 'composite_tensor', 'linalg',
         'marginalize', 'nn', 'sequential', 'substrates', 'vi'),
}
LIBS = ('bijectors', 'distributions', 'experimental', 'math', 'mcmc',
        'optimizer', 'random', 'stats', 'util')
INTERNALS = ('assert_util', 'batched_rejection_sampler', 'broadcast_util',
             'cache_util', 'callable_util',
             'custom_gradient', 'distribution_util', 'dtype_util',
             'hypothesis_testlib', 'implementation_selection', 'monte_carlo',
             'name_util', 'nest_util', 'numerics_testing',
             'parameter_properties', 'prefer_static', 'samplers',
             'special_math', 'structural_tuple', 'tensor_util',
             'tensorshape_util', 'test_combinations', 'test_util', 'unnest',
             'variadic_reduce', 'vectorization_util')
OPTIMIZERS = ('linesearch',)
LINESEARCH = ('internal',)
SAMPLERS = ('categorical', 'normal', 'poisson', 'uniform', 'shuffle')

PRIVATE_TF_PKGS = ('array_ops', 'control_flow_util', 'gradient_checker_v2',
                   'numpy_text', 'random_ops')

def main(argv):

  disabled_by_pkg = dict(DISABLED_BY_PKG)
  for dep in FLAGS.omit_deps:
    pkg = dep.split('/python/')[1].split(':')[0].replace('/', '.')
    lib = dep.split(':')[1]
    if pkg.endswith('.{}'.format(lib)):
      pkg = pkg.replace('.{}'.format(lib), '')
      disabled_by_pkg.setdefault(pkg, ())
      disabled_by_pkg[pkg] += (lib,)
    else:
      disabled_by_pkg.setdefault(pkg, ())
      disabled_by_pkg[pkg] += (lib,)

  replacements = collections.OrderedDict(TF_REPLACEMENTS)
  for pkg, disabled in disabled_by_pkg.items():
    replacements.update({
        'from tensorflow_probability.python.{}.{} '.format(pkg, item):
        '# from tensorflow_probability.python.{}.{} '.format(pkg, item)
        for item in disabled
    })
    replacements.update({
        'from tensorflow_probability.python.{} import {}'.format(pkg, item):
        '# from tensorflow_probability.python.{} import {}'.format(pkg, item)
        for item in disabled
    })
  replacements.update({
      'tensorflow_probability.python.{}'.format(lib):
      'tensorflow_probability.substrates.numpy.{}'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      'tensorflow_probability.python import {} as'.format(lib):
      'tensorflow_probability.substrates.numpy import {} as'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      'tensorflow_probability.python import {}'.format(lib):
      'tensorflow_probability.substrates.numpy import {}'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      # Permits distributions.internal, psd_kernels.internal.
      # 'as psd_kernels as': 'as',
  })
  replacements.update({
      'tensorflow_probability.python.internal.{}'.format(internal):
      'tensorflow_probability.substrates.numpy.internal.{}'.format(internal)
      for internal in INTERNALS
  })
  # pylint: disable=g-complex-comprehension
  replacements.update({
      'tensorflow_probability.python.internal import {}'.format(internal):
          'tensorflow_probability.substrates.numpy.internal import {}'.format(
              internal)
      for internal in INTERNALS
  })
  replacements.update({
      'tensorflow.python.ops import {}'.format(private):
      'tensorflow_probability.python.internal.backend.numpy import private'
      ' as {}'.format(private)
      for private in PRIVATE_TF_PKGS
  })
  replacements.update({
      'tensorflow.python.framework.ops import {}'.format(
          private):
      'tensorflow_probability.python.internal.backend.numpy import private'
      ' as {}'.format(private)
      for private in PRIVATE_TF_PKGS
  })
  # pylint: enable=g-complex-comprehension

  # TODO(bjp): Delete this block after TFP uses stateless samplers.
  replacements.update({
      'tf.random.{}'.format(sampler): 'tf.random.stateless_{}'.format(sampler)
      for sampler in SAMPLERS
  })
  replacements.update({
      'self._maybe_assert_dtype': '# self._maybe_assert_dtype',
      'SKIP_DTYPE_CHECKS = False': 'SKIP_DTYPE_CHECKS = True',
      '@test_util.test_all_tf_execution_regimes':
          '# @test_util.test_all_tf_execution_regimes',
      '@test_util.test_graph_and_eager_modes':
          '# @test_util.test_graph_and_eager_modes',
      '@test_util.test_graph_mode_only':
          '# @test_util.test_graph_mode_only',
      'TestCombinationsTest(test_util.TestCase)':
          'TestCombinationsDoNotTest(object)',
      '@six.add_metaclass(TensorMetaClass)':
          '# @six.add_metaclass(TensorMetaClass)',
  })

  filename = argv[1]
  contents = open(filename, encoding='utf-8').read()
  if '__init__.py' in filename:
    # Comment out items from __all__.
    for pkg, disabled in disabled_by_pkg.items():
      for item in disabled:
        def disable_all(name):
          replacements.update({
              '"{}"'.format(name): '# "{}"'.format(name),
              '\'{}\''.format(name): '# \'{}\''.format(name),
          })
        if 'from tensorflow_probability.python.{} import {}'.format(
            pkg, item) in contents:
          disable_all(item)
        for segment in contents.split(
            'from tensorflow_probability.python.{}.{} import '.format(
                pkg, item)):
          disable_all(segment.split('\n')[0])

  for find, replace in replacements.items():
    contents = contents.replace(find, replace)

  disabler = 'JAX_DISABLE' if FLAGS.numpy_to_jax else 'NUMPY_DISABLE'
  lines = contents.split('\n')
  for i, l in enumerate(lines):
    if disabler in l:
      lines[i] = '# {}'.format(l)
  contents = '\n'.join(lines)

  if not FLAGS.numpy_to_jax:
    contents = contents.replace('NUMPY_MODE = False', 'NUMPY_MODE = True')
  if FLAGS.numpy_to_jax:
    contents = contents.replace('tfp.substrates.numpy', 'tfp.substrates.jax')
    contents = contents.replace('substrates.numpy', 'substrates.jax')
    contents = contents.replace('backend.numpy', 'backend.jax')

    contents = contents.replace('def _call_jax', 'def __call__')
    contents = contents.replace('JAX_MODE = False', 'JAX_MODE = True')
    contents = contents.replace('SKIP_DTYPE_CHECKS = True',
                                'SKIP_DTYPE_CHECKS = False')
    is_test = lambda x: x.endswith('_test.py') or x.endswith('_test_util.py')
    if is_test(argv[1]):  # Test-only rewrites.
      contents = contents.replace(
          'tf.test.main()',
          'from jax.config import config; '
          'config.update("jax_enable_x64", True); '
          'config.enable_omnistaging(); '
          'tf.test.main()')

  print('# ' + '@' * 78)
  print('# This file is auto-generated by substrates/meta/rewrite.py')
  print('# It will be surfaced by the build system as a symlink at:')
  substrate = 'jax' if FLAGS.numpy_to_jax else 'numpy'
  print('#   `tensorflow_probability/substrates/{substrate}/{path}`'.format(
      substrate=substrate, path=filename.split('/python/')[1]))
  print('# For more info, see substrate_runfiles_symlinks in build_defs.bzl')
  print('# ' + '@' * 78)
  print('\n# (This notice adds 10 to line numbering.)\n\n')

  print(contents, file=open(1, 'w', encoding='utf-8', closefd=False))

if __name__ == '__main__':
  app.run(main)
<commit_msg>Fix the header, which somehow became a footer after we added the explicit file= to the print(contents..) line for utf8 support.<commit_after># Copyright 2019 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Rewrite script for TF->JAX."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# Dependency imports
from absl import app
from absl import flags

flags.DEFINE_boolean('numpy_to_jax', False,
                     'Whether or not to rewrite numpy imports to jax.numpy')
flags.DEFINE_list('omit_deps', [], 'List of build deps being omitted.')

FLAGS = flags.FLAGS

TF_REPLACEMENTS = {
    'import tensorflow ':
        'from tensorflow_probability.python.internal.backend import numpy ',
    'import tensorflow.compat.v1':
        'from tensorflow_probability.python.internal.backend.numpy.compat '
        'import v1',
    'import tensorflow.compat.v2':
        'from tensorflow_probability.python.internal.backend.numpy.compat '
        'import v2',
    'import tensorflow_probability as tfp':
        'import tensorflow_probability as tfp; '
        'tfp = tfp.substrates.numpy',
    'from tensorflow.python.framework import tensor_shape':
        ('from tensorflow_probability.python.internal.backend.numpy.gen '
         'import tensor_shape'),
    'from tensorflow.python.framework import ops':
        ('from tensorflow_probability.python.internal.backend.numpy '
         'import ops'),
    'from tensorflow.python.framework import tensor_util':
        ('from tensorflow_probability.python.internal.backend.numpy '
         'import ops'),
    'from tensorflow.python.util import':
        'from tensorflow_probability.python.internal.backend.numpy import',
    'from tensorflow.python.util.all_util':
        'from tensorflow_probability.python.internal.backend.numpy.private',
    'from tensorflow.python.ops.linalg':
        'from tensorflow_probability.python.internal.backend.numpy.gen',
    'from tensorflow.python.ops import parallel_for':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import functional_ops as parallel_for',
    'from tensorflow.python.ops import control_flow_ops':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import control_flow as control_flow_ops',
    'from tensorflow.python.eager import context':
        'from tensorflow_probability.python.internal.backend.numpy '
        'import private',
    ('from tensorflow.python.client '
     'import pywrap_tf_session as c_api'):
        'pass',
    ('from tensorflow.python '
     'import pywrap_tensorflow as c_api'):
        'pass'
}

DISABLED_BY_PKG = {
    'experimental':
        ('auto_batching', 'composite_tensor', 'linalg',
         'marginalize', 'nn', 'sequential', 'substrates', 'vi'),
}
LIBS = ('bijectors', 'distributions', 'experimental', 'math', 'mcmc',
        'optimizer', 'random', 'stats', 'util')
INTERNALS = ('assert_util', 'batched_rejection_sampler', 'broadcast_util',
             'cache_util', 'callable_util',
             'custom_gradient', 'distribution_util', 'dtype_util',
             'hypothesis_testlib', 'implementation_selection', 'monte_carlo',
             'name_util', 'nest_util', 'numerics_testing',
             'parameter_properties', 'prefer_static', 'samplers',
             'special_math', 'structural_tuple', 'tensor_util',
             'tensorshape_util', 'test_combinations', 'test_util', 'unnest',
             'variadic_reduce', 'vectorization_util')
OPTIMIZERS = ('linesearch',)
LINESEARCH = ('internal',)
SAMPLERS = ('categorical', 'normal', 'poisson', 'uniform', 'shuffle')

PRIVATE_TF_PKGS = ('array_ops', 'control_flow_util', 'gradient_checker_v2',
                   'numpy_text', 'random_ops')

def main(argv):

  disabled_by_pkg = dict(DISABLED_BY_PKG)
  for dep in FLAGS.omit_deps:
    pkg = dep.split('/python/')[1].split(':')[0].replace('/', '.')
    lib = dep.split(':')[1]
    if pkg.endswith('.{}'.format(lib)):
      pkg = pkg.replace('.{}'.format(lib), '')
      disabled_by_pkg.setdefault(pkg, ())
      disabled_by_pkg[pkg] += (lib,)
    else:
      disabled_by_pkg.setdefault(pkg, ())
      disabled_by_pkg[pkg] += (lib,)

  replacements = collections.OrderedDict(TF_REPLACEMENTS)
  for pkg, disabled in disabled_by_pkg.items():
    replacements.update({
        'from tensorflow_probability.python.{}.{} '.format(pkg, item):
        '# from tensorflow_probability.python.{}.{} '.format(pkg, item)
        for item in disabled
    })
    replacements.update({
        'from tensorflow_probability.python.{} import {}'.format(pkg, item):
        '# from tensorflow_probability.python.{} import {}'.format(pkg, item)
        for item in disabled
    })
  replacements.update({
      'tensorflow_probability.python.{}'.format(lib):
      'tensorflow_probability.substrates.numpy.{}'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      'tensorflow_probability.python import {} as'.format(lib):
      'tensorflow_probability.substrates.numpy import {} as'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      'tensorflow_probability.python import {}'.format(lib):
      'tensorflow_probability.substrates.numpy import {}'.format(lib)
      for lib in LIBS
  })
  replacements.update({
      # Permits distributions.internal, psd_kernels.internal.
      # 'as psd_kernels as': 'as',
  })
  replacements.update({
      'tensorflow_probability.python.internal.{}'.format(internal):
      'tensorflow_probability.substrates.numpy.internal.{}'.format(internal)
      for internal in INTERNALS
  })
  # pylint: disable=g-complex-comprehension
  replacements.update({
      'tensorflow_probability.python.internal import {}'.format(internal):
          'tensorflow_probability.substrates.numpy.internal import {}'.format(
              internal)
      for internal in INTERNALS
  })
  replacements.update({
      'tensorflow.python.ops import {}'.format(private):
      'tensorflow_probability.python.internal.backend.numpy import private'
      ' as {}'.format(private)
      for private in PRIVATE_TF_PKGS
  })
  replacements.update({
      'tensorflow.python.framework.ops import {}'.format(
          private):
      'tensorflow_probability.python.internal.backend.numpy import private'
      ' as {}'.format(private)
      for private in PRIVATE_TF_PKGS
  })
  # pylint: enable=g-complex-comprehension

  # TODO(bjp): Delete this block after TFP uses stateless samplers.
  replacements.update({
      'tf.random.{}'.format(sampler): 'tf.random.stateless_{}'.format(sampler)
      for sampler in SAMPLERS
  })
  replacements.update({
      'self._maybe_assert_dtype': '# self._maybe_assert_dtype',
      'SKIP_DTYPE_CHECKS = False': 'SKIP_DTYPE_CHECKS = True',
      '@test_util.test_all_tf_execution_regimes':
          '# @test_util.test_all_tf_execution_regimes',
      '@test_util.test_graph_and_eager_modes':
          '# @test_util.test_graph_and_eager_modes',
      '@test_util.test_graph_mode_only':
          '# @test_util.test_graph_mode_only',
      'TestCombinationsTest(test_util.TestCase)':
          'TestCombinationsDoNotTest(object)',
      '@six.add_metaclass(TensorMetaClass)':
          '# @six.add_metaclass(TensorMetaClass)',
  })

  filename = argv[1]
  contents = open(filename, encoding='utf-8').read()
  if '__init__.py' in filename:
    # Comment out items from __all__.
    for pkg, disabled in disabled_by_pkg.items():
      for item in disabled:
        def disable_all(name):
          replacements.update({
              '"{}"'.format(name): '# "{}"'.format(name),
              '\'{}\''.format(name): '# \'{}\''.format(name),
          })
        if 'from tensorflow_probability.python.{} import {}'.format(
            pkg, item) in contents:
          disable_all(item)
        for segment in contents.split(
            'from tensorflow_probability.python.{}.{} import '.format(
                pkg, item)):
          disable_all(segment.split('\n')[0])

  for find, replace in replacements.items():
    contents = contents.replace(find, replace)

  disabler = 'JAX_DISABLE' if FLAGS.numpy_to_jax else 'NUMPY_DISABLE'
  lines = contents.split('\n')
  for i, l in enumerate(lines):
    if disabler in l:
      lines[i] = '# {}'.format(l)
  contents = '\n'.join(lines)

  if not FLAGS.numpy_to_jax:
    contents = contents.replace('NUMPY_MODE = False', 'NUMPY_MODE = True')
  if FLAGS.numpy_to_jax:
    contents = contents.replace('tfp.substrates.numpy', 'tfp.substrates.jax')
    contents = contents.replace('substrates.numpy', 'substrates.jax')
    contents = contents.replace('backend.numpy', 'backend.jax')

    contents = contents.replace('def _call_jax', 'def __call__')
    contents = contents.replace('JAX_MODE = False', 'JAX_MODE = True')
    contents = contents.replace('SKIP_DTYPE_CHECKS = True',
                                'SKIP_DTYPE_CHECKS = False')
    is_test = lambda x: x.endswith('_test.py') or x.endswith('_test_util.py')
    if is_test(argv[1]):  # Test-only rewrites.
      contents = contents.replace(
          'tf.test.main()',
          'from jax.config import config; '
          'config.update("jax_enable_x64", True); '
          'config.enable_omnistaging(); '
          'tf.test.main()')

  substrate = 'jax' if FLAGS.numpy_to_jax else 'numpy'
  path = filename.split('/python/')[1]
  header = '\n'.join([
      '# ' + '@' * 78,
      '# This file is auto-generated by substrates/meta/rewrite.py',
      '# It will be surfaced by the build system as a symlink at:',
      f'#   `tensorflow_probability/substrates/{substrate}/{path}`',
      '# For more info, see substrate_runfiles_symlinks in build_defs.bzl',
      '# ' + '@' * 78,
      '\n# (This notice adds 10 to line numbering.)\n\n\n'])

  print(header + contents, file=open(1, 'w', encoding='utf-8', closefd=False))

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


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

# Query remote devices via SNMP
# and update the shared storage dict.
#
# Note that we're not on python3.6 yet,
# so shared substructures don't work.
# This is why the indices are crappy-looking composite strings.

#   Copyright [2017] [James Fleming <james@electronic-quill.net]
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

# Third-party modules
import pysnmp.hlapi
import re

# Utility functions

def snmpGet(hostname, mib, oid, community, logger, port=161):
    '''
    Perform an SNMP GET for a single OID or scalar attribute.
    Return only the value.
    '''
    # Use pysnmp to retrieve the data
    errorIndication, errorStatus, errorIndex, varBinds = next(
            pysnmp.hlapi.getCmd(
                # Create the SNMP engine
                pysnmp.hlapi.SnmpEngine(),
                # Authentication: set the SNMP version (2c) and community-string
                pysnmp.hlapi.CommunityData(community, mpModel=1),
                # Set the transport and target: UDP, hostname:port
                pysnmp.hlapi.UdpTransportTarget((hostname, port)),
                # Context is a v3 thing, but appears to be required anyway
                pysnmp.hlapi.ContextData(),
                # Specify the MIB object to read.
                # The 0 means we're retrieving a scalar value.
                pysnmp.hlapi.ObjectType(pysnmp.hlapi.ObjectIdentity(mib, oid, 0))))
    # Handle the responses
    if errorIndication:
        logger.error(errorIndication)
        return False
    elif errorStatus:
        logger.error('%s at %s' % (errorStatus.prettyPrint(),
                        errorIndex and varBinds[int(errorIndex) - 1][0] or '?'))
        return False
    # If we actually got something, return it in human-readable form
    else:
        return varBinds[0][1].prettyPrint()

def snmpBulkGet(hostname, mib, attr, community, logger, port=161):
    '''
    Perform an SNMP BULKGET on mib::attr.
    Return a 2-level dict:
    - rowname
        - index = value
    This structure mirrors SNMP's representation of tables as rows with indexed values.
    '''
    logger.debug('Querying %s for %s::%s', hostname, mib, attr)
    # Number of nonrepeating MIB variables in the request
    nonRepeaters = 0
    # Maximum number of variables requested for each of the remaining MIB variables in the request
    maxRepetitions = 10000
    # Use pysnmp to retrieve the data
    data={} # Accumulator for the results
    for (errorIndication,
            errorStatus,
            errorIndex,
            varBinds) in pysnmp.hlapi.bulkCmd(
                    # Create the SNMP engine
                    pysnmp.hlapi.SnmpEngine(),
                    # Authentication: set the SNMP version (2c) and community-string
                    pysnmp.hlapi.CommunityData(community, mpModel=1),
                    # Set the transport and target: UDP, hostname:port
                    pysnmp.hlapi.UdpTransportTarget((hostname, port)),
                    # Context is a v3 thing, but appears to be required anyway
                    pysnmp.hlapi.ContextData(),
                    # Specify operational limits
                    nonRepeaters,
                    maxRepetitions,
                    # Specify the MIB object to read.
                    # The 0 means we're retrieving a scalar value.
                    pysnmp.hlapi.ObjectType(pysnmp.hlapi.ObjectIdentity(mib, attr)),
                    # Stop when we get results outside the scope we requested,
                    # instead of carrying on until the agent runs out of OIDs to send back.
                    lexicographicMode=False):
        # Handle the responses
        if errorIndication:
            logger.error(errorIndication)
            return False
        elif errorStatus:
            logger.error('%s at %s' % (errorStatus.prettyPrint(), errorIndex and varBinds[int(errorIndex) - 1][0] or '?'))
            return False
        # If we actually got something, return it in human-readable form
        else:
            for varBind in varBinds:
                # Extract the index values.
                # We're breaking down 'IF-MIB::ifType.530' into (row='ifType', index='530')
                keys = re.split('\.', re.split('::', varBind[0].prettyPrint())[1], maxsplit=1)
                row = keys[0]
                index = keys[1]
                # Now get the value
                value = varBind[1].prettyPrint()
                # Update the results table, ensuring the row is actually present
                logger.debug('%s.%s = %s', row, index, value)
                if row not in data:
                    data[row] = {}
                data[row][index] = value
    # Return what we found
    return data

def query_device(details, logger, state):
    '''
    When completed, will send an SNMP query to a device, and report on the results.
    '''
    hostname=details['hostname']
    logger.info('Querying host %s' % hostname)
    # Process the metrics in turn
    for metric in details['metrics']:
        # Shorter references
        oid=metric['oid']
        mib=metric['mib']
        metricname=metric['metricname']
        # Generate the index in the shared-state dict
        index='%s-%s' % (hostname, oid)
        # Fetch the result
        logger.debug('Fetching OID %s::%s from target %s' % (mib, oid, hostname))
        result=int(snmpGet(hostname, mib, oid, details['community'], logger))
        # Convert it to integer if necessary
        if 'type' in metric and metric['type'] in ['counter', 'gauge']:
            val=int(result)
        else:
            val=result
        logger.debug('%s - %s::%s (%s) = %s' % (hostname, mib, oid, metricname, result))
        #
        # Handling counter-type metrics
        if 'type' in metric and metric['type'] == 'counter':
            # If this is the first run, the index won't already be in the dict
            if index in state:
                # To simulate changes in SNMP counters, calculate the difference between
                # the current value and the stored one
                diff = val - state[index]
                logger.debug('%s delta = %s' % (index, diff))
            # This is the first run; log it so the operator knows what's going on.
            else:
                logger.debug('First run for %s; skipping the diff on this run' % index)
            # Set the result
            logger.debug('Adding %s:%s to the state dict' % (metric, val))
            state[index]=val
    return True
<commit_msg>Update snmp_query.query_device regarding metric types<commit_after>#!/usr/bin/env python

# Query remote devices via SNMP
# and update the shared storage dict.
#
# Note that we're not on python3.6 yet,
# so shared substructures don't work.
# This is why the indices are crappy-looking composite strings.

#   Copyright [2017] [James Fleming <james@electronic-quill.net]
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

# Third-party modules
import pysnmp.hlapi
import re

# Utility functions

def snmpGet(hostname, mib, oid, community, logger, port=161):
    '''
    Perform an SNMP GET for a single OID or scalar attribute.
    Return only the value.
    '''
    # Use pysnmp to retrieve the data
    errorIndication, errorStatus, errorIndex, varBinds = next(
            pysnmp.hlapi.getCmd(
                # Create the SNMP engine
                pysnmp.hlapi.SnmpEngine(),
                # Authentication: set the SNMP version (2c) and community-string
                pysnmp.hlapi.CommunityData(community, mpModel=1),
                # Set the transport and target: UDP, hostname:port
                pysnmp.hlapi.UdpTransportTarget((hostname, port)),
                # Context is a v3 thing, but appears to be required anyway
                pysnmp.hlapi.ContextData(),
                # Specify the MIB object to read.
                # The 0 means we're retrieving a scalar value.
                pysnmp.hlapi.ObjectType(pysnmp.hlapi.ObjectIdentity(mib, oid, 0))))
    # Handle the responses
    if errorIndication:
        logger.error(errorIndication)
        return False
    elif errorStatus:
        logger.error('%s at %s' % (errorStatus.prettyPrint(),
                        errorIndex and varBinds[int(errorIndex) - 1][0] or '?'))
        return False
    # If we actually got something, return it in human-readable form
    else:
        return varBinds[0][1].prettyPrint()

def snmpBulkGet(hostname, mib, attr, community, logger, port=161):
    '''
    Perform an SNMP BULKGET on mib::attr.
    Return a 2-level dict:
    - rowname
        - index = value
    This structure mirrors SNMP's representation of tables as rows with indexed values.
    '''
    logger.debug('Querying %s for %s::%s', hostname, mib, attr)
    # Number of nonrepeating MIB variables in the request
    nonRepeaters = 0
    # Maximum number of variables requested for each of the remaining MIB variables in the request
    maxRepetitions = 10000
    # Use pysnmp to retrieve the data
    data={} # Accumulator for the results
    for (errorIndication,
            errorStatus,
            errorIndex,
            varBinds) in pysnmp.hlapi.bulkCmd(
                    # Create the SNMP engine
                    pysnmp.hlapi.SnmpEngine(),
                    # Authentication: set the SNMP version (2c) and community-string
                    pysnmp.hlapi.CommunityData(community, mpModel=1),
                    # Set the transport and target: UDP, hostname:port
                    pysnmp.hlapi.UdpTransportTarget((hostname, port)),
                    # Context is a v3 thing, but appears to be required anyway
                    pysnmp.hlapi.ContextData(),
                    # Specify operational limits
                    nonRepeaters,
                    maxRepetitions,
                    # Specify the MIB object to read.
                    # The 0 means we're retrieving a scalar value.
                    pysnmp.hlapi.ObjectType(pysnmp.hlapi.ObjectIdentity(mib, attr)),
                    # Stop when we get results outside the scope we requested,
                    # instead of carrying on until the agent runs out of OIDs to send back.
                    lexicographicMode=False):
        # Handle the responses
        if errorIndication:
            logger.error(errorIndication)
            return False
        elif errorStatus:
            logger.error('%s at %s' % (errorStatus.prettyPrint(), errorIndex and varBinds[int(errorIndex) - 1][0] or '?'))
            return False
        # If we actually got something, return it in human-readable form
        else:
            for varBind in varBinds:
                # Extract the index values.
                # We're breaking down 'IF-MIB::ifType.530' into (row='ifType', index='530')
                keys = re.split('\.', re.split('::', varBind[0].prettyPrint())[1], maxsplit=1)
                row = keys[0]
                index = keys[1]
                # Now get the value
                value = varBind[1].prettyPrint()
                # Update the results table, ensuring the row is actually present
                logger.debug('%s.%s = %s', row, index, value)
                if row not in data:
                    data[row] = {}
                data[row][index] = value
    # Return what we found
    return data

def query_device(details, logger, state):
    '''
    Send an SNMP query to a device, and report on the results.
    '''
    hostname=details['hostname']
    logger.info('Querying host %s' % hostname)
    # Process the metrics in turn
    for metric in details['metrics']:
        # Shorter references
        oid=metric['oid']
        mib=metric['mib']
        metricname=metric['metricname']
        # Generate the index in the shared-state dict
        index='%s-%s' % (hostname, oid)
        # Fetch the result
        logger.debug('Fetching OID %s::%s from target %s' % (mib, oid, hostname))
        result=int(snmpGet(hostname, mib, oid, details['community'], logger))
        logger.debug('%s - %s::%s (%s) = %s' % (hostname, mib, oid, metricname, result))
        #
        # Handling counter-type metrics
        if 'counter' in metric:
            # If this is the first run, the index won't already be in the dict
            if index in state:
                # To simulate changes in SNMP counters, calculate the difference between
                # the current value and the stored one
                diff = val - state[index]
                logger.debug('%s delta = %s' % (index, diff))
            # This is the first run; log it so the operator knows what's going on.
            else:
                logger.debug('First run for %s; skipping the diff on this run' % index)
            # Set the result
            logger.debug('Adding %s:%s to the state dict' % (metric, val))
            state[index]=val
    return True
<|endoftext|>


<commit_before>#!/usr/bin/env python
# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2016, Numenta, Inc.  Unless you have purchased from
# Numenta, Inc. a separate commercial license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

helpStr = """
Methods and data for running NLP model API tests. The intent here is
to ensure and changes to the models does not decrease their classification
accuracies (see nlpModelAccuracies below). Three tests are supported:

  hello: Very simple, hello world classification test. There are two categories
    that can be discriminated using bag of words. The training set is 8 docs,
    and the test set is an additional 2 (i.e. 10 total) -- first is an
    incorrectly labeled version of a training sample, second is semantically
    similar to one of the training samples.
  query: Qualitative test where we query a trained model to see which data
    samples are the most and least similar.
  simple: Less simple classification test. The dataset used here must be
    specified in the command line args.

"""
import argparse
import numpy

from prettytable import PrettyTable
from textwrap import TextWrapper
from tqdm import tqdm

from htmresearch.frameworks.nlp.classification_model import ClassificationModel
from htmresearch.frameworks.nlp.model_factory import (
  createModel, getNetworkConfig)
from htmresearch.support.csv_helper import readDataAndReshuffle

# There should be one "htm" model for each htm config entry.
nlpModelTypes = [
  "docfp",
  "cioword",
  "htm",
  "htm",
  "htm",
  "keywords"]

# Network models use 4k retina.
htmConfigs = [
  ("HTM_sensor_knn", "../data/network_configs/sensor_knn.json"),
  ("HTM_sensor_simple_tp_knn", "../data/network_configs/sensor_simple_tp_knn.json"),
  ("HTM_sensor_tm_knn", "../data/network_configs/sensor_tm_knn.json"),
]

# Some values of k we know work well.
kValues = { "keywords": 21, "docfp": 3}

nlpModelAccuracies = {
  "hello": {
    "docfp": 90.0,
    "cioword": 90.0,
    "HTM_sensor_knn": 80.0,
    "HTM_sensor_simple_tp_knn": 90.0,
    "HTM_sensor_tm_knn": 90.0,
    "keywords": 80.0,
  },
  "query": {
    "docfp": "good but not great",
    "cioword": "passable",
    "HTM_sensor_knn": "passable",
    "HTM_sensor_simple_tp_knn": None,
    "HTM_sensor_tm_knn": None,
    "keywords": "passable",
  },
  "simple": {
    "docfp": 99.7, ## CioDocumentFingerprint gets 100.0
    "cioword": 100.0,
    "HTM_sensor_knn": 66.2,
    "HTM_sensor_simple_tp_knn": 99.7,
    "HTM_sensor_tm_knn": 0.0,
    "keywords": 16.9,
  },
}

wrapper = TextWrapper(width=80)

def executeModelLifecycle(args, trainingData, labelRefs):
  """ Execute model lifecycle: create a model, train it, save it, reload it.
  Returns both the original and new models.
  """
  model = instantiateModel(args)
  model = trainModel(model, trainingData, labelRefs, args.verbosity)
  model.save(args.modelDir)
  newModel = ClassificationModel.load(args.modelDir)
  return model, newModel

def instantiateModel(args):
  """
  Set some specific arguments and return an instance of the model we will use.
  """
  args.networkConfig = getNetworkConfig(args.networkConfigPath)
  args.k = kValues.get(args.modelName, 1)

  return createModel(**vars(args))

def trainModel(model, trainingData, labelRefs, verbosity=0):
  """
  Train the given model on trainingData. Return the trained model instance.
  """
  modelName = repr(model).split()[0].split(".")[-1]
  print
  print "===================Training {} on sample text================".format(
    modelName)
  if verbosity > 0:
    printTemplate = PrettyTable(["ID", "Document", "Label"])
    printTemplate.align = "l"
    printTemplate.header_style = "upper"
  for (document, labels, docId) in tqdm(trainingData):
    if verbosity > 0:
      docStr = unicode(document, errors='ignore')
      printTemplate.add_row([docId, wrapper.fill(docStr), labelRefs[labels[0]]])
    model.trainDocument(document, labels, docId)
  if verbosity > 0:
    print printTemplate

  return model

def testModel(model, testData, labelRefs, docCategoryMap=None, verbosity=0):
  """
  Test the given model on testData, print out and return accuracy percentage.

  Accuracy is calculated as follows. Each token in a document votes for a single
  category; it's possible for a token to contribute no votes. The document is
  classified with the category that received the most votes. Note that it is
  possible for a document to receive no votes, in which case it is counted as a
  misclassification.
  """
  modelName = repr(model).split()[0].split(".")[-1]
  print
  print "===================Testing {} on sample text==================".format(
    modelName)
  if verbosity > 0:
    print
    printTemplate = PrettyTable(
      ["ID", "Document", "Actual Label(s)", "Predicted Label"])
    printTemplate.align = "l"
    printTemplate.header_style = "upper"

  numCorrect = 0
  labelRefs.append("none")
  for (document, labels, docId) in tqdm(testData):

    categoryVotes, _, _ = model.inferDocument(document)

    if categoryVotes.sum() > 0:
      # We will count classification as correct if the best category is any
      # one of the categories associated with this docId
      predicted = categoryVotes.argmax()
      if predicted in docCategoryMap[docId]:
        numCorrect += 1
    else:
      # No classification possible for this doc
      predicted = -1

    if verbosity > 0:
      docStr = unicode(document, errors="ignore")
      printTemplate.add_row(
        [docId,
         wrapper.fill(docStr),
         [labelRefs[l] for l in labels],
         labelRefs[predicted]]
      )

  accuracyPct = numCorrect * 100.0 / len(testData)

  if verbosity > 0:
    print printTemplate
  print
  print "Total correct =", numCorrect, "out of", len(testData), "documents"
  print "Accuracy =", accuracyPct, "%"

  return accuracyPct

def printSummary(testName, accuracies):
  """ Compare the new acuracies against the current values.
  @param testName (str) One of the nlpModelAccuracies keys.
  @param accuracies (dict) Keys are model names, values are accuracy percents.
  """
  try:
    currentAccuracies = nlpModelAccuracies[testName]
  except KeyError as e:
    print "No accuracy values for test '{}'".format(testName)
    raise e

  printTemplate = PrettyTable(["NLP Model", "Current Accuracy", "New Accuracy"])
  printTemplate.align = "l"
  printTemplate.header_style = "upper"

  for modelName, accuracyPct in accuracies.iteritems():
    currentPct = currentAccuracies.get(modelName, "which model?")
    printTemplate.add_row([modelName, currentPct, accuracyPct])

  print
  print "Results summary:"
  print printTemplate
<commit_msg>docfp improves with k=1<commit_after>#!/usr/bin/env python
# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2016, Numenta, Inc.  Unless you have purchased from
# Numenta, Inc. a separate commercial license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

helpStr = """
Methods and data for running NLP model API tests. The intent here is
to ensure and changes to the models does not decrease their classification
accuracies (see nlpModelAccuracies below). Three tests are supported:

  hello: Very simple, hello world classification test. There are two categories
    that can be discriminated using bag of words. The training set is 8 docs,
    and the test set is an additional 2 (i.e. 10 total) -- first is an
    incorrectly labeled version of a training sample, second is semantically
    similar to one of the training samples.
  query: Qualitative test where we query a trained model to see which data
    samples are the most and least similar.
  simple: Less simple classification test. The dataset used here must be
    specified in the command line args.

"""
import argparse
import numpy

from prettytable import PrettyTable
from textwrap import TextWrapper
from tqdm import tqdm

from htmresearch.frameworks.nlp.classification_model import ClassificationModel
from htmresearch.frameworks.nlp.model_factory import (
  createModel, getNetworkConfig)
from htmresearch.support.csv_helper import readDataAndReshuffle

# There should be one "htm" model for each htm config entry.
nlpModelTypes = [
  "docfp",
  "cioword",
  "htm",
  "htm",
  "htm",
  "keywords"]

# Network models use 4k retina.
htmConfigs = [
  ("HTM_sensor_knn", "../data/network_configs/sensor_knn.json"),
  ("HTM_sensor_simple_tp_knn", "../data/network_configs/sensor_simple_tp_knn.json"),
  ("HTM_sensor_tm_knn", "../data/network_configs/sensor_tm_knn.json"),
]

# Some values of k we know work well.
kValues = { "keywords": 21, "docfp": 1}

nlpModelAccuracies = {
  "hello": {
    "docfp": 90.0,
    "cioword": 90.0,
    "HTM_sensor_knn": 80.0,
    "HTM_sensor_simple_tp_knn": 90.0,
    "HTM_sensor_tm_knn": 90.0,
    "keywords": 80.0,
  },
  "query": {
    "docfp": "good but not great",
    "cioword": "passable",
    "HTM_sensor_knn": "passable",
    "HTM_sensor_simple_tp_knn": None,
    "HTM_sensor_tm_knn": None,
    "keywords": "passable",
  },
  "simple": {
    "docfp": 100.0,
    "cioword": 100.0,
    "HTM_sensor_knn": 66.2,
    "HTM_sensor_simple_tp_knn": 99.7,
    "HTM_sensor_tm_knn": 0.0,
    "keywords": 16.9,
  },
}

wrapper = TextWrapper(width=80)

def executeModelLifecycle(args, trainingData, labelRefs):
  """ Execute model lifecycle: create a model, train it, save it, reload it.
  Returns both the original and new models.
  """
  model = instantiateModel(args)
  model = trainModel(model, trainingData, labelRefs, args.verbosity)
  model.save(args.modelDir)
  newModel = ClassificationModel.load(args.modelDir)
  return model, newModel

def instantiateModel(args):
  """
  Set some specific arguments and return an instance of the model we will use.
  """
  args.networkConfig = getNetworkConfig(args.networkConfigPath)
  args.k = kValues.get(args.modelName, 1)

  return createModel(**vars(args))

def trainModel(model, trainingData, labelRefs, verbosity=0):
  """
  Train the given model on trainingData. Return the trained model instance.
  """
  modelName = repr(model).split()[0].split(".")[-1]
  print
  print "===================Training {} on sample text================".format(
    modelName)
  if verbosity > 0:
    printTemplate = PrettyTable(["ID", "Document", "Label"])
    printTemplate.align = "l"
    printTemplate.header_style = "upper"
  for (document, labels, docId) in tqdm(trainingData):
    if verbosity > 0:
      docStr = unicode(document, errors='ignore')
      printTemplate.add_row([docId, wrapper.fill(docStr), labelRefs[labels[0]]])
    model.trainDocument(document, labels, docId)
  if verbosity > 0:
    print printTemplate

  return model

def testModel(model, testData, labelRefs, docCategoryMap=None, verbosity=0):
  """
  Test the given model on testData, print out and return accuracy percentage.

  Accuracy is calculated as follows. Each token in a document votes for a single
  category; it's possible for a token to contribute no votes. The document is
  classified with the category that received the most votes. Note that it is
  possible for a document to receive no votes, in which case it is counted as a
  misclassification.
  """
  modelName = repr(model).split()[0].split(".")[-1]
  print
  print "===================Testing {} on sample text==================".format(
    modelName)
  if verbosity > 0:
    print
    printTemplate = PrettyTable(
      ["ID", "Document", "Actual Label(s)", "Predicted Label"])
    printTemplate.align = "l"
    printTemplate.header_style = "upper"

  numCorrect = 0
  labelRefs.append("none")
  for (document, labels, docId) in tqdm(testData):

    categoryVotes, _, _ = model.inferDocument(document)

    if categoryVotes.sum() > 0:
      # We will count classification as correct if the best category is any
      # one of the categories associated with this docId
      predicted = categoryVotes.argmax()
      if predicted in docCategoryMap[docId]:
        numCorrect += 1
    else:
      # No classification possible for this doc
      predicted = -1

    if verbosity > 0:
      docStr = unicode(document, errors="ignore")
      printTemplate.add_row(
        [docId,
         wrapper.fill(docStr),
         [labelRefs[l] for l in labels],
         labelRefs[predicted]]
      )

  accuracyPct = numCorrect * 100.0 / len(testData)

  if verbosity > 0:
    print printTemplate
  print
  print "Total correct =", numCorrect, "out of", len(testData), "documents"
  print "Accuracy =", accuracyPct, "%"

  return accuracyPct

def printSummary(testName, accuracies):
  """ Compare the new acuracies against the current values.
  @param testName (str) One of the nlpModelAccuracies keys.
  @param accuracies (dict) Keys are model names, values are accuracy percents.
  """
  try:
    currentAccuracies = nlpModelAccuracies[testName]
  except KeyError as e:
    print "No accuracy values for test '{}'".format(testName)
    raise e

  printTemplate = PrettyTable(["NLP Model", "Current Accuracy", "New Accuracy"])
  printTemplate.align = "l"
  printTemplate.header_style = "upper"

  for modelName, accuracyPct in accuracies.iteritems():
    currentPct = currentAccuracies.get(modelName, "which model?")
    printTemplate.add_row([modelName, currentPct, accuracyPct])

  print
  print "Results summary:"
  print printTemplate
<|endoftext|>