<commit_before># -*- coding: iso8859-1 -*-
#
# Copyright (C) 2005 Edgewall Software
# Copyright (C) 2005 Christopher Lenz <cmlenz@gmx.de>
# Copyright (C) 2005 Matthew Good <trac@matt-good.net>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.com/license.html.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://projects.edgewall.com/trac/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>
#         Matthew Good <trac@matt-good.net>

import os

from trac.core import *
from trac.env import open_environment
from trac.perm import PermissionCache, PermissionError
from trac.util import escape, enum, format_datetime, http_date
from trac.web.api import absolute_url, Request, RequestDone, IAuthenticator, \
                         IRequestHandler
from trac.web.chrome import Chrome
from trac.web.clearsilver import HDFWrapper
from trac.web.href import Href
from trac.web.session import Session

# Environment cache for multithreaded front-ends:
try:
    import threading
except ImportError:
    has_threads = False
else:
    has_threads = True
    env_cache = {}
    env_cache_lock = threading.Lock()

def _open_environment(env_path, threaded=True):
    if not has_threads or not threaded:
        return open_environment(env_path)

    global env_cache, env_cache_lock
    env = None
    env_cache_lock.acquire()
    try:
        if not env_path in env_cache:
            env_cache[env_path] = open_environment(env_path)
        env = env_cache[env_path]
    finally:
        env_cache_lock.release()
    return env

class RequestDispatcher(Component):
    """Component responsible for dispatching requests to registered handlers."""

    authenticators = ExtensionPoint(IAuthenticator)
    handlers = ExtensionPoint(IRequestHandler)

    def authenticate(self, req):
        for authenticator in self.authenticators:
            authname = authenticator.authenticate(req)
            if authname:
                return authname
        else:
            return 'anonymous'

    def dispatch(self, req):
        """Find a registered handler that matches the request and let it process
        it.

        In addition, this method initializes the HDF data set and adds the web
        site chrome.
        """
        req.authname = self.authenticate(req)
        req.perm = PermissionCache(self.env, req.authname)

        chrome = Chrome(self.env)
        req.hdf = HDFWrapper(loadpaths=chrome.get_all_templates_dirs())
        populate_hdf(req.hdf, self.env, req)

        newsession = req.args.has_key('newsession')
        req.session = Session(self.env, req, newsession)

        # Select the component that should handle the request
        chosen_handler = None
        default_handler = None
        if not req.path_info or req.path_info == '/':
            default_handler = self.config.get('trac', 'default_handler')
        for handler in self.handlers:
            if handler.match_request(req) or \
               handler.__class__.__name__ == default_handler:
                chosen_handler = handler
                break

        chrome.populate_hdf(req, chosen_handler)

        if not chosen_handler:
            # FIXME: Should return '404 Not Found' to the client
            raise TracError, 'No handler matched request to %s' % req.path_info

        try:
            resp = chosen_handler.process_request(req)
            if resp:
                template, content_type = resp
                if not content_type:
                    content_type = 'text/html'

                req.display(template, content_type or 'text/html')
        finally:
            # Give the session a chance to persist changes
            req.session.save()

def dispatch_request(path_info, req, env):
    """Main entry point for the Trac web interface."""

    # Re-parse the configuration file if it changed since the last the time it
    # was parsed
    env.config.parse_if_needed()

    base_url = env.config.get('trac', 'base_url')
    if not base_url:
        base_url = absolute_url(req)
    req.base_url = base_url
    req.path_info = path_info

    env.href = Href(req.cgi_location)
    env.abs_href = Href(req.base_url)

    db = env.get_db_cnx()
    try:
        try:
            dispatcher = RequestDispatcher(env)
            dispatcher.dispatch(req)
        except RequestDone:
            pass
    finally:
        db.close()

def populate_hdf(hdf, env, req=None):
    """Populate the HDF data set with various information, such as common URLs,
    project information and request-related information.
    """
    from trac import __version__
    hdf['trac'] = {
        'version': __version__,
        'time': format_datetime(),
        'time.gmt': http_date()
    }
    hdf['trac.href'] = {
        'wiki': env.href.wiki(),
        'browser': env.href.browser('/'),
        'timeline': env.href.timeline(),
        'roadmap': env.href.roadmap(),
        'milestone': env.href.milestone(None),
        'report': env.href.report(),
        'query': env.href.query(),
        'newticket': env.href.newticket(),
        'search': env.href.search(),
        'about': env.href.about(),
        'about_config': env.href.about('config'),
        'login': env.href.login(),
        'logout': env.href.logout(),
        'settings': env.href.settings(),
        'homepage': 'http://trac.edgewall.com/'
    }

    hdf['project'] = {
        'name': env.config.get('project', 'name'),
        'name_encoded': escape(env.config.get('project', 'name')),
        'descr': env.config.get('project', 'descr'),
        'footer': env.config.get('project', 'footer',
                 'Visit the Trac open source project at
'
                 ''
                 'http://trac.edgewall.com/'),
        'url': env.config.get('project', 'url')
    }

    if req:
        hdf['base_url'] = req.base_url
        hdf['base_host'] = req.base_url[:req.base_url.rfind(req.cgi_location)]
        hdf['cgi_location'] = req.cgi_location
        hdf['trac.authname'] = escape(req.authname)

        for action in req.perm.permissions():
            req.hdf['trac.acl.' + action] = True

        for arg in [k for k in req.args.keys() if k]:
            if isinstance(req.args[arg], (list, tuple)):
                hdf['args.%s' % arg] = [v.value for v in req.args[arg]]
            else:
                hdf['args.%s' % arg] = req.args[arg].value

def send_pretty_error(e, env, req=None):
    """Send a "pretty" HTML error page to the client."""
    import traceback
    import StringIO
    tb = StringIO.StringIO()
    traceback.print_exc(file=tb)
    if not req:
        from trac.web.cgi_frontend import CGIRequest
        from trac.web.clearsilver import HDFWrapper
        req = CGIRequest()
        req.authname = ''
        req.hdf = HDFWrapper()
    try:
        if not env:
            from trac.env import open_environment
            env = open_environment()
            env.href = Href(req.cgi_location)
        if env and env.log:
            env.log.exception(e)
        populate_hdf(req.hdf, env, req)

        if isinstance(e, TracError):
            req.hdf['title'] = e.title or 'Error'
            req.hdf['error.title'] = e.title or 'Error'
            req.hdf['error.type'] = 'TracError'
            req.hdf['error.message'] = e.message
            if e.show_traceback:
                req.hdf['error.traceback'] = escape(tb.getvalue())
            req.display('error.cs', response=500)

        elif isinstance(e, PermissionError):
            req.hdf['title'] = 'Permission Denied'
            req.hdf['error.type'] = 'permission'
            req.hdf['error.action'] = e.action
            req.hdf['error.message'] = e
            req.display('error.cs', response=403)

        else:
            req.hdf['title'] = 'Oops'
            req.hdf['error.type'] = 'internal'
            req.hdf['error.message'] = escape(str(e))
            req.hdf['error.traceback'] = escape(tb.getvalue())
            req.display('error.cs', response=500)

    except RequestDone:
        pass
    except Exception, e2:
        if env and env.log:
            env.log.error('Failed to render pretty error page: %s', e2,
                          exc_info=True)
        try:
            req.send_response(500)
            req.send_header('Content-Type', 'text/plain')
            req.end_headers()
            req.write('Oops...\n\nTrac detected an internal error:\n\n')
            req.write(str(e))
            req.write('\n')
            req.write(tb.getvalue())
        except IOError:
            # Cannot send response, but the error has hopefully been logged
            pass

def send_project_index(req, options, env_paths=None):
    from trac.web.clearsilver import HDFWrapper

    if 'TRAC_ENV_INDEX_TEMPLATE' in options:
        tmpl_path, template = os.path.split(options['TRAC_ENV_INDEX_TEMPLATE'])

        from trac.config import default_dir
        req.hdf = HDFWrapper(loadpaths=[default_dir('templates'), tmpl_path])

        tmpl_vars = {}
        if 'TRAC_TEMPLATE_VARS' in options:
            for pair in options['TRAC_TEMPLATE_VARS'].split(','):
                key, val = pair.split('=')
                req.hdf[key] = val
    else:
        req.hdf = HDFWrapper()
        template = req.hdf.parse('''
<title>Available Projects</title>

Available Projects

<?cs
 each:project = projects ?>
<?cs
  if:project.href ?>
   <a href="<?cs var:project.href ?>" title="<?cs var:project.description ?>">
    <?cs var:project.name ?><?cs
  else ?>
   <?cs var:project.name ?>: Error

   (<?cs var:project.description ?>)<?cs
  /if ?>

<?cs
 /each ?>

''')

    if not env_paths and 'TRAC_ENV_PARENT_DIR' in options:
        dir = options['TRAC_ENV_PARENT_DIR']
        env_paths = [os.path.join(dir, f) for f in os.listdir(dir)]

    href = Href(req.idx_location)
    try:
        projects = []
        for env_path in env_paths:
            if not os.path.isdir(env_path):
                continue
            env_dir, project = os.path.split(env_path)
            try:
                env = _open_environment(env_path)
                proj = {
                    'name': env.config.get('project', 'name'),
                    'description': env.config.get('project', 'descr'),
                    'href': href(project)
                }
            except Exception, e:
                proj = {'name': project, 'description': str(e)}
            projects.append(proj)
        projects.sort(lambda x, y: cmp(x['name'].lower(), y['name'].lower()))
        req.hdf['projects'] = projects

        # TODO maybe this should be 404 if index wasn't specifically requested
        req.display(template, response=200)
    except RequestDone:
        pass

def get_environment(req, options, threaded=True):
    if 'TRAC_ENV' in options:
        env_path = options['TRAC_ENV']
    elif 'TRAC_ENV_PARENT_DIR' in options:
        env_parent_dir = options['TRAC_ENV_PARENT_DIR']
        env_name = req.cgi_location.split('/')[-1]
        env_path = os.path.join(env_parent_dir, env_name)
        if not len(env_name) or not os.path.exists(env_path):
            return None
    else:
        raise TracError, \
              'The environment options "TRAC_ENV" or "TRAC_ENV_PARENT_DIR" ' \
              'or the mod_python options "TracEnv" or "TracEnvParentDir" ' \
              'are missing.  Trac requires one of these options to locate ' \
              'the Trac environment(s).'

    return _open_environment(env_path, threaded)
<commit_msg>Always transcode the request path info to UTF-8. Seems to fix #1239.<commit_after># -*- coding: iso8859-1 -*-
#
# Copyright (C) 2005 Edgewall Software
# Copyright (C) 2005 Christopher Lenz <cmlenz@gmx.de>
# Copyright (C) 2005 Matthew Good <trac@matt-good.net>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.com/license.html.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://projects.edgewall.com/trac/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>
#         Matthew Good <trac@matt-good.net>

import os

from trac.core import *
from trac.env import open_environment
from trac.perm import PermissionCache, PermissionError
from trac.util import escape, enum, format_datetime, http_date, to_utf8
from trac.web.api import absolute_url, Request, RequestDone, IAuthenticator, \
                         IRequestHandler
from trac.web.chrome import Chrome
from trac.web.clearsilver import HDFWrapper
from trac.web.href import Href
from trac.web.session import Session

# Environment cache for multithreaded front-ends:
try:
    import threading
except ImportError:
    has_threads = False
else:
    has_threads = True
    env_cache = {}
    env_cache_lock = threading.Lock()

def _open_environment(env_path, threaded=True):
    if not has_threads or not threaded:
        return open_environment(env_path)

    global env_cache, env_cache_lock
    env = None
    env_cache_lock.acquire()
    try:
        if not env_path in env_cache:
            env_cache[env_path] = open_environment(env_path)
        env = env_cache[env_path]
    finally:
        env_cache_lock.release()
    return env

class RequestDispatcher(Component):
    """Component responsible for dispatching requests to registered handlers."""

    authenticators = ExtensionPoint(IAuthenticator)
    handlers = ExtensionPoint(IRequestHandler)

    def authenticate(self, req):
        for authenticator in self.authenticators:
            authname = authenticator.authenticate(req)
            if authname:
                return authname
        else:
            return 'anonymous'

    def dispatch(self, req):
        """Find a registered handler that matches the request and let it process
        it.

        In addition, this method initializes the HDF data set and adds the web
        site chrome.
        """
        req.authname = self.authenticate(req)
        req.perm = PermissionCache(self.env, req.authname)

        chrome = Chrome(self.env)
        req.hdf = HDFWrapper(loadpaths=chrome.get_all_templates_dirs())
        populate_hdf(req.hdf, self.env, req)

        newsession = req.args.has_key('newsession')
        req.session = Session(self.env, req, newsession)

        # Select the component that should handle the request
        chosen_handler = None
        default_handler = None
        if not req.path_info or req.path_info == '/':
            default_handler = self.config.get('trac', 'default_handler')
        for handler in self.handlers:
            if handler.match_request(req) or \
               handler.__class__.__name__ == default_handler:
                chosen_handler = handler
                break

        chrome.populate_hdf(req, chosen_handler)

        if not chosen_handler:
            # FIXME: Should return '404 Not Found' to the client
            raise TracError, 'No handler matched request to %s' % req.path_info

        try:
            resp = chosen_handler.process_request(req)
            if resp:
                template, content_type = resp
                if not content_type:
                    content_type = 'text/html'

                req.display(template, content_type or 'text/html')
        finally:
            # Give the session a chance to persist changes
            req.session.save()

def dispatch_request(path_info, req, env):
    """Main entry point for the Trac web interface."""

    # Re-parse the configuration file if it changed since the last the time it
    # was parsed
    env.config.parse_if_needed()

    base_url = env.config.get('trac', 'base_url')
    if not base_url:
        base_url = absolute_url(req)
    req.base_url = base_url
    req.path_info = to_utf8(path_info)

    env.href = Href(req.cgi_location)
    env.abs_href = Href(req.base_url)

    db = env.get_db_cnx()
    try:
        try:
            dispatcher = RequestDispatcher(env)
            dispatcher.dispatch(req)
        except RequestDone:
            pass
    finally:
        db.close()

def populate_hdf(hdf, env, req=None):
    """Populate the HDF data set with various information, such as common URLs,
    project information and request-related information.
    """
    from trac import __version__
    hdf['trac'] = {
        'version': __version__,
        'time': format_datetime(),
        'time.gmt': http_date()
    }
    hdf['trac.href'] = {
        'wiki': env.href.wiki(),
        'browser': env.href.browser('/'),
        'timeline': env.href.timeline(),
        'roadmap': env.href.roadmap(),
        'milestone': env.href.milestone(None),
        'report': env.href.report(),
        'query': env.href.query(),
        'newticket': env.href.newticket(),
        'search': env.href.search(),
        'about': env.href.about(),
        'about_config': env.href.about('config'),
        'login': env.href.login(),
        'logout': env.href.logout(),
        'settings': env.href.settings(),
        'homepage': 'http://trac.edgewall.com/'
    }

    hdf['project'] = {
        'name': env.config.get('project', 'name'),
        'name_encoded': escape(env.config.get('project', 'name')),
        'descr': env.config.get('project', 'descr'),
        'footer': env.config.get('project', 'footer',
                 'Visit the Trac open source project at
'
                 ''
                 'http://trac.edgewall.com/'),
        'url': env.config.get('project', 'url')
    }

    if req:
        hdf['base_url'] = req.base_url
        hdf['base_host'] = req.base_url[:req.base_url.rfind(req.cgi_location)]
        hdf['cgi_location'] = req.cgi_location
        hdf['trac.authname'] = escape(req.authname)

        for action in req.perm.permissions():
            req.hdf['trac.acl.' + action] = True

        for arg in [k for k in req.args.keys() if k]:
            if isinstance(req.args[arg], (list, tuple)):
                hdf['args.%s' % arg] = [v.value for v in req.args[arg]]
            else:
                hdf['args.%s' % arg] = req.args[arg].value

def send_pretty_error(e, env, req=None):
    """Send a "pretty" HTML error page to the client."""
    import traceback
    import StringIO
    tb = StringIO.StringIO()
    traceback.print_exc(file=tb)
    if not req:
        from trac.web.cgi_frontend import CGIRequest
        from trac.web.clearsilver import HDFWrapper
        req = CGIRequest()
        req.authname = ''
        req.hdf = HDFWrapper()
    try:
        if not env:
            from trac.env import open_environment
            env = open_environment()
            env.href = Href(req.cgi_location)
        if env and env.log:
            env.log.exception(e)
        populate_hdf(req.hdf, env, req)

        if isinstance(e, TracError):
            req.hdf['title'] = e.title or 'Error'
            req.hdf['error.title'] = e.title or 'Error'
            req.hdf['error.type'] = 'TracError'
            req.hdf['error.message'] = e.message
            if e.show_traceback:
                req.hdf['error.traceback'] = escape(tb.getvalue())
            req.display('error.cs', response=500)

        elif isinstance(e, PermissionError):
            req.hdf['title'] = 'Permission Denied'
            req.hdf['error.type'] = 'permission'
            req.hdf['error.action'] = e.action
            req.hdf['error.message'] = e
            req.display('error.cs', response=403)

        else:
            req.hdf['title'] = 'Oops'
            req.hdf['error.type'] = 'internal'
            req.hdf['error.message'] = escape(str(e))
            req.hdf['error.traceback'] = escape(tb.getvalue())
            req.display('error.cs', response=500)

    except RequestDone:
        pass
    except Exception, e2:
        if env and env.log:
            env.log.error('Failed to render pretty error page: %s', e2,
                          exc_info=True)
        try:
            req.send_response(500)
            req.send_header('Content-Type', 'text/plain')
            req.end_headers()
            req.write('Oops...\n\nTrac detected an internal error:\n\n')
            req.write(str(e))
            req.write('\n')
            req.write(tb.getvalue())
        except IOError:
            # Cannot send response, but the error has hopefully been logged
            pass

def send_project_index(req, options, env_paths=None):
    from trac.web.clearsilver import HDFWrapper

    if 'TRAC_ENV_INDEX_TEMPLATE' in options:
        tmpl_path, template = os.path.split(options['TRAC_ENV_INDEX_TEMPLATE'])

        from trac.config import default_dir
        req.hdf = HDFWrapper(loadpaths=[default_dir('templates'), tmpl_path])

        tmpl_vars = {}
        if 'TRAC_TEMPLATE_VARS' in options:
            for pair in options['TRAC_TEMPLATE_VARS'].split(','):
                key, val = pair.split('=')
                req.hdf[key] = val
    else:
        req.hdf = HDFWrapper()
        template = req.hdf.parse('''
<title>Available Projects</title>

Available Projects

<?cs
 each:project = projects ?>
<?cs
  if:project.href ?>
   <a href="<?cs var:project.href ?>" title="<?cs var:project.description ?>">
    <?cs var:project.name ?><?cs
  else ?>
   <?cs var:project.name ?>: Error

   (<?cs var:project.description ?>)<?cs
  /if ?>

<?cs
 /each ?>

''')

    if not env_paths and 'TRAC_ENV_PARENT_DIR' in options:
        dir = options['TRAC_ENV_PARENT_DIR']
        env_paths = [os.path.join(dir, f) for f in os.listdir(dir)]

    href = Href(req.idx_location)
    try:
        projects = []
        for env_path in env_paths:
            if not os.path.isdir(env_path):
                continue
            env_dir, project = os.path.split(env_path)
            try:
                env = _open_environment(env_path)
                proj = {
                    'name': env.config.get('project', 'name'),
                    'description': env.config.get('project', 'descr'),
                    'href': href(project)
                }
            except Exception, e:
                proj = {'name': project, 'description': str(e)}
            projects.append(proj)
        projects.sort(lambda x, y: cmp(x['name'].lower(), y['name'].lower()))
        req.hdf['projects'] = projects

        # TODO maybe this should be 404 if index wasn't specifically requested
        req.display(template, response=200)
    except RequestDone:
        pass

def get_environment(req, options, threaded=True):
    if 'TRAC_ENV' in options:
        env_path = options['TRAC_ENV']
    elif 'TRAC_ENV_PARENT_DIR' in options:
        env_parent_dir = options['TRAC_ENV_PARENT_DIR']
        env_name = req.cgi_location.split('/')[-1]
        env_path = os.path.join(env_parent_dir, env_name)
        if not len(env_name) or not os.path.exists(env_path):
            return None
    else:
        raise TracError, \
              'The environment options "TRAC_ENV" or "TRAC_ENV_PARENT_DIR" ' \
              'or the mod_python options "TracEnv" or "TracEnvParentDir" ' \
              'are missing.  Trac requires one of these options to locate ' \
              'the Trac environment(s).'

    return _open_environment(env_path, threaded)
<|endoftext|>


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

import unittest
import json

# from TinCanPython.agent import Agent, AgentAccount
from TinCanPython.tcapi_version import TCAPIVersion

class TinCanBaseTestCase(unittest.TestCase):
    # def get_agent(self, name, id_type, id_fields):
    #     agent = Agent()
    #     agent.name = name
    #
    #     if 'mbox' == id_type:
    #         agent.mbox = id_fields
    #     elif 'openid' == id_type:
    #         agent.openid = id_fields
    #     elif 'mbox_sha1sum' == id_type:
    #         agent.mbox_sha1sum = id_fields
    #     elif 'account' == id_type:
    #         parts = id_fields.split('|')
    #         acct = AgentAccount()
    #         acct.home_page = parts[0]
    #         acct.name = parts[1]
    #         agent.account = acct
    #
    #     return agent

    # PEP8 says this should be lowercase, but unittest breaks this rule
    def assertSerializeDeserialize(self, obj, version=None):
        """
        Verifies that objects are packed to JSON and unpacked
        correctly for the version specified. If no version is
        specified, tests for all versions in TCAPIVersion.
        :param obj: A TinCan object to be tested.
        :param version: The version according to whose schema we will test.
        :type version: str
        """
        tested_versions = [version] if version is not None else TCAPIVersion.values()
        for version in tested_versions:
            json_obj = obj.to_json(version)
            constructor = obj.__class__ if not isinstance(obj, dict) else dict
            unpacked_json = json.loads(json_obj)
            clone = constructor(**unpacked_json)

            self.assertEqual(isinstance(obj, dict), isinstance(clone, dict))

            orig_dict = obj if isinstance(obj, dict) else obj.__dict__
            clone_dict = clone if isinstance(clone, dict) else clone.__dict__

            self.assertEqual(orig_dict, clone_dict)
<commit_msg>added TODOs for code to be re-enabled once agent.py is available<commit_after>#!/usr/bin/env python

import unittest
import json

## TODO: re-enable this
# from TinCanPython.agent import Agent, AgentAccount
from TinCanPython.tcapi_version import TCAPIVersion

class TinCanBaseTestCase(unittest.TestCase):
    ## TODO: re-enable this
    # def get_agent(self, name, id_type, id_fields):
    #     agent = Agent()
    #     agent.name = name
    #
    #     if 'mbox' == id_type:
    #         agent.mbox = id_fields
    #     elif 'openid' == id_type:
    #         agent.openid = id_fields
    #     elif 'mbox_sha1sum' == id_type:
    #         agent.mbox_sha1sum = id_fields
    #     elif 'account' == id_type:
    #         parts = id_fields.split('|')
    #         acct = AgentAccount()
    #         acct.home_page = parts[0]
    #         acct.name = parts[1]
    #         agent.account = acct
    #
    #     return agent

    # PEP8 says this should be lowercase, but unittest breaks this rule
    def assertSerializeDeserialize(self, obj, version=None):
        """
        Verifies that objects are packed to JSON and unpacked
        correctly for the version specified. If no version is
        specified, tests for all versions in TCAPIVersion.
        :param obj: A TinCan object to be tested.
        :param version: The version according to whose schema we will test.
        :type version: str
        """
        tested_versions = [version] if version is not None else TCAPIVersion.values()
        for version in tested_versions:
            json_obj = obj.to_json(version)
            constructor = obj.__class__ if not isinstance(obj, dict) else dict
            unpacked_json = json.loads(json_obj)
            clone = constructor(**unpacked_json)

            self.assertEqual(isinstance(obj, dict), isinstance(clone, dict))

            orig_dict = obj if isinstance(obj, dict) else obj.__dict__
            clone_dict = clone if isinstance(clone, dict) else clone.__dict__

            self.assertEqual(orig_dict, clone_dict)
<|endoftext|>


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

###
# Copyright (c) 2002, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

from test import *

import utils

class UtilsTest(unittest.TestCase):
    def testTimeElapsed(self):
        self.assertRaises(ValueError, utils.timeElapsed, 0, 0, seconds=False)
        then = 0
        now = 0
        for now, expected in [(0, '0 seconds'),
                              (1, '1 second'),
                              (60, '1 minute and 0 seconds'),
                              (61, '1 minute and 1 second'),
                              (62, '1 minute and 2 seconds'),
                              (122, '2 minutes and 2 seconds'),
                              (3722, '1 hour, 2 minutes, and 2 seconds'),
                              (7322, '2 hours, 2 minutes, and 2 seconds')]:
            self.assertEqual(utils.timeElapsed(now - then), expected)

    def testEachSubstring(self):
        s = 'foobar'
        L = ['f', 'fo', 'foo', 'foob', 'fooba', 'foobar']
        self.assertEqual(list(utils.eachSubstring(s)), L)

    def testDistance(self):
        self.assertEqual(utils.distance('', ''), 0)
        self.assertEqual(utils.distance('a', 'b'), 1)
        self.assertEqual(utils.distance('a', 'a'), 0)
        self.assertEqual(utils.distance('foobar', 'jemfinch'), 8)
        self.assertEqual(utils.distance('a', 'ab'), 1)
        self.assertEqual(utils.distance('foo', ''), 3)
        self.assertEqual(utils.distance('', 'foo'), 3)
        self.assertEqual(utils.distance('appel', 'nappe'), 2)
        self.assertEqual(utils.distance('nappe', 'appel'), 2)

    def testAbbrev(self):
        L = ['abc', 'bcd', 'bbe', 'foo', 'fool']
        d = utils.abbrev(L)
        def getItem(s):
            return d[s]
        self.assertRaises(KeyError, getItem, 'f')
        self.assertRaises(KeyError, getItem, 'fo')
        self.assertRaises(KeyError, getItem, 'b')
        self.assertEqual(d['bb'], 'bbe')
        self.assertEqual(d['bc'], 'bcd')
        self.assertEqual(d['a'], 'abc')
        self.assertEqual(d['ab'], 'abc')
        self.assertEqual(d['fool'], 'fool')
        self.assertEqual(d['foo'], 'foo')

    def testSoundex(self):
        L = [('Euler', 'E460'),
             ('Ellery', 'E460'),
             ('Gauss', 'G200'),
             ('Ghosh', 'G200'),
             ('Hilbert', 'H416'),
             ('Heilbronn', 'H416'),
             ('Knuth', 'K530'),
             ('Kant', 'K530'),
             ('Lloyd', 'L300'),
             ('Ladd', 'L300'),
             ('Lukasiewicz', 'L222'),
             ('Lissajous', 'L222')]
        for (name, key) in L:
            soundex = utils.soundex(name)
            self.assertEqual(soundex, key,
                             '%s was %s, not %s' % (name, soundex, key))

    def testDQRepr(self):
        L = ['foo', 'foo\'bar', 'foo"bar', '"', '\\', '', '\x00']
        for s in L:
            r = utils.dqrepr(s)
            self.assertEqual(s, eval(r), s)
            self.failUnless(r[0] == '"' and r[-1] == '"', s)

    def testPerlReToPythonRe(self):
        r = utils.perlReToPythonRe('m/foo/')
        self.failUnless(r.search('foo'))
        r = utils.perlReToPythonRe('/foo/')
        self.failUnless(r.search('foo'))
        r = utils.perlReToPythonRe('m/\\//')
        self.failUnless(r.search('/'))

    def testPerlReToReplacer(self):
        f = utils.perlReToReplacer('s/foo/bar/')
        self.assertEqual(f('foobarbaz'), 'barbarbaz')
        f = utils.perlReToReplacer('s/fool/bar/')
        self.assertEqual(f('foobarbaz'), 'foobarbaz')
        f = utils.perlReToReplacer('s/foo//')
        self.assertEqual(f('foobarbaz'), 'barbaz')
        f = utils.perlReToReplacer('s/ba//')
        self.assertEqual(f('foobarbaz'), 'foorbaz')
        f = utils.perlReToReplacer('s/ba//g')
        self.assertEqual(f('foobarbaz'), 'foorz')
        f = utils.perlReToReplacer('s/ba\\///g')
        self.assertEqual(f('fooba/rba/z'), 'foorz')

    def testFindBinaryInPath(self):
        if os.name == 'posix':
            self.assertEqual(None, utils.findBinaryInPath('asdfhjklasdfhjkl'))
            self.assertEqual('/bin/sh', utils.findBinaryInPath('sh'))

    def testCommaAndify(self):
        L = ['foo']
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo')
        self.assertEqual(L, original)
        L.append('bar')
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo and bar')
        self.assertEqual(L, original)
        L.append('baz')
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo, bar, and baz')
        self.assertEqual(L, original)

# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:

<commit_msg>Updated the tests to stress more possibilities<commit_after>#!/usr/bin/env python

###
# Copyright (c) 2002, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

from test import *
import sets

import utils

class UtilsTest(unittest.TestCase):
    def testTimeElapsed(self):
        self.assertRaises(ValueError, utils.timeElapsed, 0, 0, seconds=False)
        then = 0
        now = 0
        for now, expected in [(0, '0 seconds'),
                              (1, '1 second'),
                              (60, '1 minute and 0 seconds'),
                              (61, '1 minute and 1 second'),
                              (62, '1 minute and 2 seconds'),
                              (122, '2 minutes and 2 seconds'),
                              (3722, '1 hour, 2 minutes, and 2 seconds'),
                              (7322, '2 hours, 2 minutes, and 2 seconds')]:
            self.assertEqual(utils.timeElapsed(now - then), expected)

    def testEachSubstring(self):
        s = 'foobar'
        L = ['f', 'fo', 'foo', 'foob', 'fooba', 'foobar']
        self.assertEqual(list(utils.eachSubstring(s)), L)

    def testDistance(self):
        self.assertEqual(utils.distance('', ''), 0)
        self.assertEqual(utils.distance('a', 'b'), 1)
        self.assertEqual(utils.distance('a', 'a'), 0)
        self.assertEqual(utils.distance('foobar', 'jemfinch'), 8)
        self.assertEqual(utils.distance('a', 'ab'), 1)
        self.assertEqual(utils.distance('foo', ''), 3)
        self.assertEqual(utils.distance('', 'foo'), 3)
        self.assertEqual(utils.distance('appel', 'nappe'), 2)
        self.assertEqual(utils.distance('nappe', 'appel'), 2)

    def testAbbrev(self):
        L = ['abc', 'bcd', 'bbe', 'foo', 'fool']
        d = utils.abbrev(L)
        def getItem(s):
            return d[s]
        self.assertRaises(KeyError, getItem, 'f')
        self.assertRaises(KeyError, getItem, 'fo')
        self.assertRaises(KeyError, getItem, 'b')
        self.assertEqual(d['bb'], 'bbe')
        self.assertEqual(d['bc'], 'bcd')
        self.assertEqual(d['a'], 'abc')
        self.assertEqual(d['ab'], 'abc')
        self.assertEqual(d['fool'], 'fool')
        self.assertEqual(d['foo'], 'foo')

    def testSoundex(self):
        L = [('Euler', 'E460'),
             ('Ellery', 'E460'),
             ('Gauss', 'G200'),
             ('Ghosh', 'G200'),
             ('Hilbert', 'H416'),
             ('Heilbronn', 'H416'),
             ('Knuth', 'K530'),
             ('Kant', 'K530'),
             ('Lloyd', 'L300'),
             ('Ladd', 'L300'),
             ('Lukasiewicz', 'L222'),
             ('Lissajous', 'L222')]
        for (name, key) in L:
            soundex = utils.soundex(name)
            self.assertEqual(soundex, key,
                             '%s was %s, not %s' % (name, soundex, key))

    def testDQRepr(self):
        L = ['foo', 'foo\'bar', 'foo"bar', '"', '\\', '', '\x00']
        for s in L:
            r = utils.dqrepr(s)
            self.assertEqual(s, eval(r), s)
            self.failUnless(r[0] == '"' and r[-1] == '"', s)

    def testPerlReToPythonRe(self):
        r = utils.perlReToPythonRe('m/foo/')
        self.failUnless(r.search('foo'))
        r = utils.perlReToPythonRe('/foo/')
        self.failUnless(r.search('foo'))
        r = utils.perlReToPythonRe('m/\\//')
        self.failUnless(r.search('/'))

    def testPerlReToReplacer(self):
        f = utils.perlReToReplacer('s/foo/bar/')
        self.assertEqual(f('foobarbaz'), 'barbarbaz')
        f = utils.perlReToReplacer('s/fool/bar/')
        self.assertEqual(f('foobarbaz'), 'foobarbaz')
        f = utils.perlReToReplacer('s/foo//')
        self.assertEqual(f('foobarbaz'), 'barbaz')
        f = utils.perlReToReplacer('s/ba//')
        self.assertEqual(f('foobarbaz'), 'foorbaz')
        f = utils.perlReToReplacer('s/ba//g')
        self.assertEqual(f('foobarbaz'), 'foorz')
        f = utils.perlReToReplacer('s/ba\\///g')
        self.assertEqual(f('fooba/rba/z'), 'foorz')

    def testFindBinaryInPath(self):
        if os.name == 'posix':
            self.assertEqual(None, utils.findBinaryInPath('asdfhjklasdfhjkl'))
            self.assertEqual('/bin/sh', utils.findBinaryInPath('sh'))

    def testCommaAndify(self):
        L = ['foo']
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo')
        self.assertEqual(L, original)
        L.append('bar')
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo and bar')
        self.assertEqual(L, original)
        L.append('baz')
        original = L[:]
        self.assertEqual(utils.commaAndify(L), 'foo, bar, and baz')
        self.assertEqual(L, original)
        self.failUnless(utils.commaAndify(sets.Set(L)))

# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:

<|endoftext|>


<commit_before>"""tracking loop for telescope control.

Track provides the classes required to point a telescope with software using a
feedback control loop.

"""

from __future__ import print_function
import time
import abc
from track.telem import TelemSource
from track.mathutils import clamp

class ErrorSource(object):
    """Abstract parent class for error sources.

    This class provides some abstract methods to provide a common interface
    for error sources to be used in the tracking loop. All error sources must
    inheret from this class and implement the methods defined.
    """
    __metaclass__ = abc.ABCMeta

    class NoSignalException(Exception):
        """Raised when no signal is available for error calculation."""
        pass

    @abc.abstractmethod
    def get_axis_names(self):
        """Get axis names

        Returns:
            A list of strings giving abbreviated names of each axis. These must
            be the keys of the dict returned by the compute_error method.
        """
        pass

    @abc.abstractmethod
    def compute_error(self):
        """Computes the error signal.

        Returns:
            The pointing error as a dict with entries for each axis. The units
            should be degrees.

        Raises:
            NoSignalException: If the error cannot be computed.
        """
        pass

class TelescopeMount(object):
    """Abstract parent class for telescope mounts.

    This class provides some abstract methods to provide a common interface
    for telescope mounts to be used in the tracking loop. All mounts must
    inheret from this class and implement the methods defined. Only Az-Alt
    mounts are currently supported.
    """
    __metaclass__ = abc.ABCMeta

    class AxisLimitException(Exception):
        """Raised when an axis reaches or exceeds limits."""
        def __init__(self, axes):
            self.axes = axes

    @abc.abstractmethod
    def get_axis_names(self):
        """Get axis names

        Returns:
            A list of strings giving abbreviated names of each axis. These must
            be the keys of the dicts used by the get_position,
            get_aligned_slew_dir, remove_backlash, and get_max_slew_rates
            methods and the names accepted for the axis argument of the slew
            method.
        """

    @abc.abstractmethod
    def get_position(self, max_cache_age=0.0):
        """Gets the current position of the mount.

        Args:
            max_cache_age: If the position has been read from the mount less
                than this many seconds ago, the function may return a cached
                position value in lieu of reading the position from the mount.
                In cases where reading from the mount is relatively slow this
                may allow the function to return much more quickly. The default
                value is set to 0 seconds, in which case the function will
                never return a cached value.

        Returns:
            A dict with keys for each axis where the values are the positions
            in degrees.
        """
        pass

    @abc.abstractmethod
    def slew(self, axis, rate):
        """Command the mount to slew on one axis.

        Commands the mount to slew at a paritcular rate in one axis. A mount may enforce limits
        on the slew rate such as a maximum rate limit, acceleration limit, or max change in
        rate since the last command. Enforcement of such limits may result in the mount moving at
        a rate that is different than requested. The return values indicate whether a limit was
        exceeded and what slew rate was actually achieved.

        Args:
            axis: A string indicating the axis.
            rate: A float giving the requested slew rate in degrees per second. The sign of the
                value indicates the direction of the slew.

        Returns:
            A two-element tuple where the first element is a float giving the actual slew rate
                achieved by the mount. The actual rate could differ from the requested rate due
                to quantization or due to enforcement of slew rate or acceleration limits. The
                second element is a boolean that is set to True if any of the limits enforced by
                the mount were exceeded by the requested slew rate.

        Raises:
            AltitudeLimitException: Implementation dependent.
        """
        pass

class LoopFilter(object):
    """Proportional plus integral (PI) loop filter.

    This class implements a standard proportional plus integral loop filter.
    The proportional and integral coefficients are computed on the fly in order
    to support a dynamic loop period.

    Attributes:
        bandwidth: Loop bandwidth in Hz.
        damping_factor: Loop damping factor.
        max_update_period: Maximum tolerated loop update period in seconds.
        rate_limit: Maximum allowed slew rate in degrees per second.
        accel_limit: Maximum allowed slew acceleration in degrees per second
            squared.
        step_limit: Maximum allowed change in slew rate per update in degrees per second.
        int: Integrator value.
        last_iteration_time: Unix time of last call to update().
        last_rate: Output value returned on last call to update().
    """
    def __init__(
            self,
            bandwidth,
            damping_factor,
            max_update_period=0.1
        ):
        """Inits a Loop Filter object.

        Args:
            bandwidth: Loop bandwidth in Hz.
            damping_factor: Loop damping factor.
            rate_limit: Slew rate limit in degrees per second. Set to None to remove limit.
            accel_limit: Slew acceleration limit in degrees per second squared. Set to None to
                remove limit.
            step_limit: Maximum change in slew rate allowed per update in degrees per second. Set
                to None to remove limit.
            max_update_period: Maximum tolerated loop update period in seconds.
        """
        self.bandwidth = bandwidth
        self.damping_factor = damping_factor
        self.max_update_period = max_update_period
        self.int = 0.0
        self.last_iteration_time = None
        self.last_rate = 0.0

    def clamp_integrator(self, rate):
        """Clamps the integrator magnitude to not exceed a particular rate.

        This function should be called when the slew rate given by the return value of the update()
        method trips a slew rate or acceleration limit in the mount. Otherwise the integrator can
        grow in an unbounded fashion and the system will not respond appropriately.

        Args:
            rate: The actual rate achieved by the mount after the most recent call to update(), in
                degrees per second.
        """
        self.int = clamp(self.int, abs(rate))

    def update(self, error):
        """Update the loop filter using new error signal input.

        Updates the loop filter using new error signal information. The loop
        filter proportional and integral coefficients are calculated on each
        call based on the time elapsed since the previous call. This allows
        the loop response to remain consistent even if the loop period is
        changing dynamically or can't be predicted in advance.

        If this method was last called more than max_update_period seconds ago
        a warning will be printed and the stored integrator value will be
        returned. The error signal will be ignored. This is meant to protect
        against edge cases where long periods between calls to update() could
        cause huge disturbances to the loop behavior.

        Args:
            error: The error in phase units (typically degrees).

        Returns:
            The slew rate to be applied in the same units as the error signal
                (typically degrees).
        """

        # can't measure loop period on first update
        if self.last_iteration_time is None:
            self.last_iteration_time = time.time()
            return 0.0

        update_period = time.time() - self.last_iteration_time
        self.last_iteration_time = time.time()
        if update_period > self.max_update_period:
            print('Warning: loop filter update period was {:.4f} s, limit is {:.4f} s.'.format(
                update_period,
                self.max_update_period,
            ))
            return self.int

        # compute loop filter gains based on loop period
        bt = self.bandwidth * update_period
        k0 = update_period
        denom = self.damping_factor + 1.0 / (4.0 * self.damping_factor)
        prop_gain = 4.0 * self.damping_factor / denom * bt / k0
        int_gain = 4.0 / denom**2.0 * bt**2.0 / k0

        # proportional term
        prop = prop_gain * -error

        # integral term
        self.int = self.int + int_gain * -error

        # New candidate slew rate is the sum of P and I terms. Note that the mount may enforce
        # rate or acceleration limits such that the actual rate achieved is different from this
        # rate. When this is the case, the control loop should ensure that the loop filter
        # integrator value is also saturated to prevent it from growing in an unbounded manner.
        return prop + self.int

class Tracker(TelemSource):
    """Main tracking loop class.

    This class is the core of the track package. A tracking loop or control
    loop is a system that uses feedback for control. In this case, the thing
    under control is a telescope mount. Slew commands are sent to the mount
    at regular intervals to keep it pointed in the direction of some object.
    In order to know which direction the mount should slew, the tracking loop
    needs feedback which comes in the form of an error signal. The error signal
    is a measure of the difference between where the telescope is pointed now
    compared to where the object is. The control loop tries to drive the error
    signal magnitude to zero. The final component in the loop is a loop filter
    ("loop controller" might be a better name, but "loop filter" is the name
    everyone uses). The loop filter controls the response characteristics of
    the system including the bandwidth (responsiveness to changes in input)
    and the damping factor.

    Attributes:
        loop_filter: A dict with keys for each axis where the values are
            LoopFilter objects. Each axis has its own independent loop filter.
        mount: An object of type TelescopeMount. This represents the interface
            to the mount.
        error_source: An object of type ErrorSource. The error can be computed
            in many ways so an abstract class with a generic interface is used.
        error: Cached error value returned by the error_source object's
            compute_error() method. This is cached so that callback methods
            can make use of it if needed. A dict with keys for each axis.
        slew_rate: Cached slew rates from the most recent loop filter output.
            Cached to make it available to callbacks. A dict with keys for each
            axis.
        num_iterations: A running count of the number of iterations.
        callback: A callback function. The callback will be called once at the
            end of every control loop iteration. Set to None if no callback is
            registered.
        stop: Boolean value. The control loop checks this on every iteration
            and stops if the value is True.
    """

    def __init__(self, mount, error_source, loop_bandwidth, damping_factor):
        """Inits a Tracker object.

        Initializes a Tracker object by constructing loop filters and
        initializing state information.

        Args:
            mount: Object of type TelescopeMount. Must use the same set of axes
                as error_source.
            error_source: Object of type ErrorSource. Must use the same set of
                axes as mount.
            loop_bandwidth: The loop bandwidth in Hz.
            damping_factor: The damping factor. Keep in mind that the motors in
                the mount will not respond instantaneously to slew commands,
                therefore the damping factor may need to be higher than an
                ideal system would suggest. Common values like sqrt(2)/2 may be
                too small to prevent oscillations.
        """
        if set(mount.get_axis_names()) != set(error_source.get_axis_names()):
            raise ValueError('error_source and mount must use same set of axes')
        self.loop_filter = {}
        self.error = {}
        self.slew_rate = {}
        for axis in mount.get_axis_names():
            self.loop_filter[axis] = LoopFilter(
                bandwidth=loop_bandwidth,
                damping_factor=damping_factor,
            )
            self.error[axis] = None
            self.slew_rate[axis] = 0.0
        self.mount = mount
        self.error_source = error_source
        self.num_iterations = 0
        self.callback = None
        self.stop = False
        self.stop_on_timer = False
        self.max_run_time = 0.0
        self.stop_when_converged = False
        self.converge_max_error_mag = 50.0 / 3600.0
        self.converge_min_iterations = 50
        self.converge_error_state = None

    def register_callback(self, callback):
        """Register a callback function.

        Registers a callback function to be called near the end of each loop
        iteration. The callback function is called with no arguments.

        Args:
            callback: The function to call. None to un-register.
        """
        self.callback = callback

    def run(self, axes=None):
        """Run the control loop.

        Call this method to start the control loop. This function is blocking
        and will not return until an error occurs or until the stop attribute
        has been set to True. Immediately after invocation this function will
        set the stop attribute to False.

        Args:
            axes: A list of strings indicating which axes should be under
                active tracking control. The slew rate on any axis not included
                in the list will not be commanded by the control loop. If an
                empty list is passed the function returns immediately. If None
                is passed all axes will be active.
        """
        self.stop = False
        low_error_iterations = 0
        start_time = time.time()

        if axes is None:
            axes = self.mount.get_axis_names()

        if len(axes) == 0:
            return 'no axes selected'

        while True:

            if self.stop:
                return 'stop flag set'

            if self.stop_when_converged and low_error_iterations >= self.converge_min_iterations:
                return 'converged'

            if self.stop_on_timer and time.time() - start_time > self.max_run_time:
                return 'timer expired'

            # compute error
            try:
                self.error = self.error_source.compute_error()
            except ErrorSource.NoSignalException:
                for axis in self.mount.get_axis_names():
                    self.error[axis] = None
                self.finish_control_cycle()
                continue

            if self.error['mag'] > self.converge_max_error_mag:
                low_error_iterations = 0
            else:
                if self.converge_error_state is None:
                    low_error_iterations += 1
                elif self.error_source.state == self.converge_error_state:
                    low_error_iterations +=1
                else:
                    low_error_iterations = 0

            # update loop filters
            for axis in axes:
                self.slew_rate[axis] = self.loop_filter[axis].update(self.error[axis])

            # set mount slew rates
            for axis in axes:
                try:
                    (actual_rate, limit_exceeded) = self.mount.slew(axis, self.slew_rate[axis])
                    if limit_exceeded:
                        self.loop_filter[axis].clamp_integrator(actual_rate)
                except TelescopeMount.AxisLimitException:
                    self.loop_filter[axis].int = 0.0

            self.finish_control_cycle()

    def finish_control_cycle(self):
        """Final tasks to perform at the end of each control cycle."""
        if self.callback is not None:
            self.callback(self)
        self.num_iterations += 1

    def get_telem_channels(self):
        chans = {}
        chans['num_iterations'] = self.num_iterations
        for axis in self.mount.get_axis_names():
            chans['rate_' + axis] = self.slew_rate[axis]
            chans['error_' + axis] = self.error[axis]
            chans['loop_filt_int_' + axis] = self.loop_filter[axis].int
        return chans
<commit_msg>#102: Fix slew_rate telemetry channel values<commit_after>"""tracking loop for telescope control.

Track provides the classes required to point a telescope with software using a
feedback control loop.

"""

from __future__ import print_function
import time
import abc
from track.telem import TelemSource
from track.mathutils import clamp

class ErrorSource(object):
    """Abstract parent class for error sources.

    This class provides some abstract methods to provide a common interface
    for error sources to be used in the tracking loop. All error sources must
    inheret from this class and implement the methods defined.
    """
    __metaclass__ = abc.ABCMeta

    class NoSignalException(Exception):
        """Raised when no signal is available for error calculation."""
        pass

    @abc.abstractmethod
    def get_axis_names(self):
        """Get axis names

        Returns:
            A list of strings giving abbreviated names of each axis. These must
            be the keys of the dict returned by the compute_error method.
        """
        pass

    @abc.abstractmethod
    def compute_error(self):
        """Computes the error signal.

        Returns:
            The pointing error as a dict with entries for each axis. The units
            should be degrees.

        Raises:
            NoSignalException: If the error cannot be computed.
        """
        pass

class TelescopeMount(object):
    """Abstract parent class for telescope mounts.

    This class provides some abstract methods to provide a common interface
    for telescope mounts to be used in the tracking loop. All mounts must
    inheret from this class and implement the methods defined. Only Az-Alt
    mounts are currently supported.
    """
    __metaclass__ = abc.ABCMeta

    class AxisLimitException(Exception):
        """Raised when an axis reaches or exceeds limits."""
        def __init__(self, axes):
            self.axes = axes

    @abc.abstractmethod
    def get_axis_names(self):
        """Get axis names

        Returns:
            A list of strings giving abbreviated names of each axis. These must
            be the keys of the dicts used by the get_position,
            get_aligned_slew_dir, remove_backlash, and get_max_slew_rates
            methods and the names accepted for the axis argument of the slew
            method.
        """

    @abc.abstractmethod
    def get_position(self, max_cache_age=0.0):
        """Gets the current position of the mount.

        Args:
            max_cache_age: If the position has been read from the mount less
                than this many seconds ago, the function may return a cached
                position value in lieu of reading the position from the mount.
                In cases where reading from the mount is relatively slow this
                may allow the function to return much more quickly. The default
                value is set to 0 seconds, in which case the function will
                never return a cached value.

        Returns:
            A dict with keys for each axis where the values are the positions
            in degrees.
        """
        pass

    @abc.abstractmethod
    def slew(self, axis, rate):
        """Command the mount to slew on one axis.

        Commands the mount to slew at a paritcular rate in one axis. A mount may enforce limits
        on the slew rate such as a maximum rate limit, acceleration limit, or max change in
        rate since the last command. Enforcement of such limits may result in the mount moving at
        a rate that is different than requested. The return values indicate whether a limit was
        exceeded and what slew rate was actually achieved.

        Args:
            axis: A string indicating the axis.
            rate: A float giving the requested slew rate in degrees per second. The sign of the
                value indicates the direction of the slew.

        Returns:
            A two-element tuple where the first element is a float giving the actual slew rate
                achieved by the mount. The actual rate could differ from the requested rate due
                to quantization or due to enforcement of slew rate or acceleration limits. The
                second element is a boolean that is set to True if any of the limits enforced by
                the mount were exceeded by the requested slew rate.

        Raises:
            AltitudeLimitException: Implementation dependent.
        """
        pass

class LoopFilter(object):
    """Proportional plus integral (PI) loop filter.

    This class implements a standard proportional plus integral loop filter.
    The proportional and integral coefficients are computed on the fly in order
    to support a dynamic loop period.

    Attributes:
        bandwidth: Loop bandwidth in Hz.
        damping_factor: Loop damping factor.
        max_update_period: Maximum tolerated loop update period in seconds.
        rate_limit: Maximum allowed slew rate in degrees per second.
        accel_limit: Maximum allowed slew acceleration in degrees per second
            squared.
        step_limit: Maximum allowed change in slew rate per update in degrees per second.
        int: Integrator value.
        last_iteration_time: Unix time of last call to update().
        last_rate: Output value returned on last call to update().
    """
    def __init__(
            self,
            bandwidth,
            damping_factor,
            max_update_period=0.1
        ):
        """Inits a Loop Filter object.

        Args:
            bandwidth: Loop bandwidth in Hz.
            damping_factor: Loop damping factor.
            rate_limit: Slew rate limit in degrees per second. Set to None to remove limit.
            accel_limit: Slew acceleration limit in degrees per second squared. Set to None to
                remove limit.
            step_limit: Maximum change in slew rate allowed per update in degrees per second. Set
                to None to remove limit.
            max_update_period: Maximum tolerated loop update period in seconds.
        """
        self.bandwidth = bandwidth
        self.damping_factor = damping_factor
        self.max_update_period = max_update_period
        self.int = 0.0
        self.last_iteration_time = None
        self.last_rate = 0.0

    def clamp_integrator(self, rate):
        """Clamps the integrator magnitude to not exceed a particular rate.

        This function should be called when the slew rate given by the return value of the update()
        method trips a slew rate or acceleration limit in the mount. Otherwise the integrator can
        grow in an unbounded fashion and the system will not respond appropriately.

        Args:
            rate: The actual rate achieved by the mount after the most recent call to update(), in
                degrees per second.
        """
        self.int = clamp(self.int, abs(rate))

    def update(self, error):
        """Update the loop filter using new error signal input.

        Updates the loop filter using new error signal information. The loop
        filter proportional and integral coefficients are calculated on each
        call based on the time elapsed since the previous call. This allows
        the loop response to remain consistent even if the loop period is
        changing dynamically or can't be predicted in advance.

        If this method was last called more than max_update_period seconds ago
        a warning will be printed and the stored integrator value will be
        returned. The error signal will be ignored. This is meant to protect
        against edge cases where long periods between calls to update() could
        cause huge disturbances to the loop behavior.

        Args:
            error: The error in phase units (typically degrees).

        Returns:
            The slew rate to be applied in the same units as the error signal
                (typically degrees).
        """

        # can't measure loop period on first update
        if self.last_iteration_time is None:
            self.last_iteration_time = time.time()
            return 0.0

        update_period = time.time() - self.last_iteration_time
        self.last_iteration_time = time.time()
        if update_period > self.max_update_period:
            print('Warning: loop filter update period was {:.4f} s, limit is {:.4f} s.'.format(
                update_period,
                self.max_update_period,
            ))
            return self.int

        # compute loop filter gains based on loop period
        bt = self.bandwidth * update_period
        k0 = update_period
        denom = self.damping_factor + 1.0 / (4.0 * self.damping_factor)
        prop_gain = 4.0 * self.damping_factor / denom * bt / k0
        int_gain = 4.0 / denom**2.0 * bt**2.0 / k0

        # proportional term
        prop = prop_gain * -error

        # integral term
        self.int = self.int + int_gain * -error

        # New candidate slew rate is the sum of P and I terms. Note that the mount may enforce
        # rate or acceleration limits such that the actual rate achieved is different from this
        # rate. When this is the case, the control loop should ensure that the loop filter
        # integrator value is also saturated to prevent it from growing in an unbounded manner.
        return prop + self.int

class Tracker(TelemSource):
    """Main tracking loop class.

    This class is the core of the track package. A tracking loop or control
    loop is a system that uses feedback for control. In this case, the thing
    under control is a telescope mount. Slew commands are sent to the mount
    at regular intervals to keep it pointed in the direction of some object.
    In order to know which direction the mount should slew, the tracking loop
    needs feedback which comes in the form of an error signal. The error signal
    is a measure of the difference between where the telescope is pointed now
    compared to where the object is. The control loop tries to drive the error
    signal magnitude to zero. The final component in the loop is a loop filter
    ("loop controller" might be a better name, but "loop filter" is the name
    everyone uses). The loop filter controls the response characteristics of
    the system including the bandwidth (responsiveness to changes in input)
    and the damping factor.

    Attributes:
        loop_filter: A dict with keys for each axis where the values are
            LoopFilter objects. Each axis has its own independent loop filter.
        mount: An object of type TelescopeMount. This represents the interface
            to the mount.
        error_source: An object of type ErrorSource. The error can be computed
            in many ways so an abstract class with a generic interface is used.
        error: Cached error value returned by the error_source object's
            compute_error() method. This is cached so that callback methods
            can make use of it if needed. A dict with keys for each axis.
        slew_rate: Cached slew rates from the most recent loop filter output.
            Cached to make it available to callbacks. A dict with keys for each
            axis.
        num_iterations: A running count of the number of iterations.
        callback: A callback function. The callback will be called once at the
            end of every control loop iteration. Set to None if no callback is
            registered.
        stop: Boolean value. The control loop checks this on every iteration
            and stops if the value is True.
    """

    def __init__(self, mount, error_source, loop_bandwidth, damping_factor):
        """Inits a Tracker object.

        Initializes a Tracker object by constructing loop filters and
        initializing state information.

        Args:
            mount: Object of type TelescopeMount. Must use the same set of axes
                as error_source.
            error_source: Object of type ErrorSource. Must use the same set of
                axes as mount.
            loop_bandwidth: The loop bandwidth in Hz.
            damping_factor: The damping factor. Keep in mind that the motors in
                the mount will not respond instantaneously to slew commands,
                therefore the damping factor may need to be higher than an
                ideal system would suggest. Common values like sqrt(2)/2 may be
                too small to prevent oscillations.
        """
        if set(mount.get_axis_names()) != set(error_source.get_axis_names()):
            raise ValueError('error_source and mount must use same set of axes')
        self.loop_filter = {}
        self.loop_filter_output = {}
        self.error = {}
        self.slew_rate = {}
        for axis in mount.get_axis_names():
            self.loop_filter[axis] = LoopFilter(
                bandwidth=loop_bandwidth,
                damping_factor=damping_factor,
            )
            self.loop_filter_output[axis] = 0.0
            self.error[axis] = None
            self.slew_rate[axis] = 0.0
        self.mount = mount
        self.error_source = error_source
        self.num_iterations = 0
        self.callback = None
        self.stop = False
        self.stop_on_timer = False
        self.max_run_time = 0.0
        self.stop_when_converged = False
        self.converge_max_error_mag = 50.0 / 3600.0
        self.converge_min_iterations = 50
        self.converge_error_state = None

    def register_callback(self, callback):
        """Register a callback function.

        Registers a callback function to be called near the end of each loop
        iteration. The callback function is called with no arguments.

        Args:
            callback: The function to call. None to un-register.
        """
        self.callback = callback

    def run(self, axes=None):
        """Run the control loop.

        Call this method to start the control loop. This function is blocking
        and will not return until an error occurs or until the stop attribute
        has been set to True. Immediately after invocation this function will
        set the stop attribute to False.

        Args:
            axes: A list of strings indicating which axes should be under
                active tracking control. The slew rate on any axis not included
                in the list will not be commanded by the control loop. If an
                empty list is passed the function returns immediately. If None
                is passed all axes will be active.
        """
        self.stop = False
        low_error_iterations = 0
        start_time = time.time()

        if axes is None:
            axes = self.mount.get_axis_names()

        if len(axes) == 0:
            return 'no axes selected'

        while True:

            if self.stop:
                return 'stop flag set'

            if self.stop_when_converged and low_error_iterations >= self.converge_min_iterations:
                return 'converged'

            if self.stop_on_timer and time.time() - start_time > self.max_run_time:
                return 'timer expired'

            # compute error
            try:
                self.error = self.error_source.compute_error()
            except ErrorSource.NoSignalException:
                for axis in self.mount.get_axis_names():
                    self.error[axis] = None
                self.finish_control_cycle()
                continue

            if self.error['mag'] > self.converge_max_error_mag:
                low_error_iterations = 0
            else:
                if self.converge_error_state is None:
                    low_error_iterations += 1
                elif self.error_source.state == self.converge_error_state:
                    low_error_iterations +=1
                else:
                    low_error_iterations = 0

            # update loop filters
            for axis in axes:
                self.loop_filter_output[axis] = self.loop_filter[axis].update(self.error[axis])

            # set mount slew rates
            for axis in axes:
                try:
                    (self.slew_rate[axis], limit_exceeded) = self.mount.slew(axis, self.loop_filter_output[axis])
                    if limit_exceeded:
                        self.loop_filter[axis].clamp_integrator(self.slew_rate[axis])
                except TelescopeMount.AxisLimitException:
                    self.loop_filter[axis].int = 0.0

            self.finish_control_cycle()

    def finish_control_cycle(self):
        """Final tasks to perform at the end of each control cycle."""
        if self.callback is not None:
            self.callback(self)
        self.num_iterations += 1

    def get_telem_channels(self):
        chans = {}
        chans['num_iterations'] = self.num_iterations
        for axis in self.mount.get_axis_names():
            chans['rate_' + axis] = self.slew_rate[axis]
            chans['error_' + axis] = self.error[axis]
            chans['loop_filt_int_' + axis] = self.loop_filter[axis].int
            chans['loop_filt_out_' + axis] = self.loop_filter_output[axis]
        return chans
<|endoftext|>


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

import yaml
import time

from astroquery.simbad import Simbad
from astropy.coordinates import SkyCoord
from astropy import units as u

import panoptes
pan = panoptes.Panoptes(connect_on_startup=True)

tracking_file = 'tracking_log.txt'

def track(targets, default_exp_time=60, interval=2):

    for target in targets:
        # Sidereal
        c = SkyCoord(target.get('coords'), frame='icrs')

        exp_time = default_exp_time

        # Run commands
        for cmd in target.get('commands'):
            cmd = cmd.split(' ')
            pan.observatory.mount.serial_query(*cmd)  # split for params

        # Set the mount to target
        pan.observatory.mount.set_target_coordinates(c)
        pan.observatory.mount.slew_to_target()

        # Wait while slewing
        while pan.observatory.mount.is_slewing():
            pan.logger.debug("Mount is slewing. Sleeping for two seconds...")
            time.sleep(2)

        with open(tracking_file, 'w+') as f:
            print("{}".format('*' * 20), file=f, flush=True)
            print("Target Coords: {}\nMount Commands: {}\nExp Time: {}".format(
            	target.get('coords'), target.get('commands'), default_exp_time
            	),
                file=f, flush=True)

            line = '\t'.join(['local_date', 'local_time', 'sidereal_time', 'ra', 'dec', 'alt', 'az', 'pier_position'])
            print(line, file=f, flush=True)

            pan.logger.info("Exposing for {} seconds".format(exp_time))

            # Track for exp_time, reporting along the way
            while exp_time > 0:
                alt = pan.observatory.mount.serial_query('get_alt')
                az = pan.observatory.mount.serial_query('get_az')

                ra = pan.observatory.mount.serial_query('get_ra')
                dec = pan.observatory.mount.serial_query('get_dec')

                local_time = pan.observatory.mount.serial_query('get_local_time')
                local_date = pan.observatory.mount.serial_query('get_local_date')
                sidereal_time = pan.observatory.mount.serial_query('get_sidereal_time')

                pier_position = pan.observatory.mount.pier_position()

                line = '\t'.join([local_date, local_time, sidereal_time, ra, dec, alt, az, pier_position])
                print(line, file=f, flush=True)

                time.sleep(interval)
                exp_time -= interval
                pan.logger.info("Exposing for {} more seconds".format(exp_time))

    pan.observatory.mount.slew_to_home()

if __name__ == '__main__':
    with open('targets.yaml', 'r') as f:
        targets = yaml.load(f.read())

    exp_time = 60 * 5  # seconds * minutes

    track(targets, exp_time)

    """
-Acquire target
	- track 1 hr sidereal
	- query ra for similar
- Acquire target
	- track 1 hr custom ra rate
	- query ra/dec
- Acquire target
	- track 1 hr custom dec rate
	- query ra/dec
	"""
<commit_msg>Test script updates for new scheduler code<commit_after>#!/usr/bin/env python

import yaml
import time

from astroquery.simbad import Simbad
from astropy.coordinates import SkyCoord
from astropy import units as u

import panoptes

def track(pan):

    interval = 2

    target = pan.scheduler.get_target(pan.observatory)

    c = target.position
    print("Moving to target: {} {}".format(target.name, c))

    for visit in target.visit:
        # Sidereal

        exp_time = visit.master_exptime

        # Set the mount to visit
        pan.observatory.mount.set_target_coordinates(c)
        pan.observatory.mount.slew_to_target()

        # Wait while slewing
        while pan.observatory.mount.is_slewing():
            pan.logger.debug("Mount is slewing. Sleeping for two seconds...")
            time.sleep(interval)

        with open(tracking_file, 'w+') as f:
            print("{}".format('*' * 20), file=f, flush=True)
            print("Target Coords: {}\nMount Commands: {}\nExp Time: {}".format(
            	target.position, '', exp_time
            	),
                file=f, flush=True)

            line = '\t'.join(['local_date', 'local_time', 'sidereal_time', 'ra', 'dec', 'alt', 'az', 'pier_position'])
            print(line, file=f, flush=True)

            pan.logger.info("Exposing for {} seconds".format(exp_time))

            # Track for exp_time, reporting along the way
            while exp_time > 0:
                alt = pan.observatory.mount.serial_query('get_alt')
                az = pan.observatory.mount.serial_query('get_az')

                ra = pan.observatory.mount.serial_query('get_ra')
                dec = pan.observatory.mount.serial_query('get_dec')

                local_time = pan.observatory.mount.serial_query('get_local_time')
                local_date = pan.observatory.mount.serial_query('get_local_date')
                sidereal_time = pan.observatory.mount.serial_query('get_sidereal_time')

                pier_position = pan.observatory.mount.pier_position()

                line = '\t'.join([local_date, local_time, sidereal_time, ra, dec, alt, az, pier_position])
                print(line, file=f, flush=True)

                time.sleep(interval)
                exp_time -= interval * u.s
                pan.logger.info("Exposing for {} more seconds".format(exp_time))

    pan.observatory.mount.slew_to_home()

if __name__ == '__main__':
    pan = panoptes.Panoptes(connect_on_startup=True)

    tracking_file = 'tracking_log.txt'

    track(pan)

    """
-Acquire target
	- track 1 hr sidereal
	- query ra for similar
- Acquire target
	- track 1 hr custom ra rate
	- query ra/dec
- Acquire target
	- track 1 hr custom dec rate
	- query ra/dec
	"""
<|endoftext|>


<commit_before>"""Helper utilities and decorators."""
from flask import flash

def flash_errors(form, category='warning'):
    """Flash all errors for a form."""
    for field, errors in form.errors.items():
        for error in errors:
            flash('{0} - {1}'.format(getattr(form, field).label.text, error), category)
<commit_msg>Remove more bootstrapping cruft<commit_after><|endoftext|>


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

import sys
import os
import platform
import subprocess
import re
import shutil
import datetime
import logging
from optparse import OptionParser

####################################
# Option parser and constants
TRACTOR_PURGE_VERSION = 'v1.0.0'
parser = OptionParser(version='%prog ' + TRACTOR_PURGE_VERSION)
parser.add_option('-t', '--tq', dest='tq',
                  default='/opt/pixar/Tractor-2.2/bin/tq',
                  help='Absolute path to tq [default: %default]')
parser.add_option('-c', '--cmdlogsdir', dest='cmdlogsdir',
                  default='/var/spool/tractor/cmd-logs',
                  help='Absolute path to cmd-logs dir [default: %default]')
parser.add_option('-l', '--log', dest='logfile',
                  default='/var/tmp/tractor-purge.log',
                  help='Absolute path to tractor-purge log file [default: %default]')
parser.add_option('-d', '--days', dest='days', default='30',
                  help='Number of days worth of jobs/logs to keep [default: %default]')
parser.add_option('--deletejobs', action='store_true', dest='deletejobs',
                  default=False,
                  help='Delete jobs from psql database after log deletion')
parser.add_option('--dryrun', action='store_true', dest='dryrun',
                  default=False,
                  help='Do not perform actual deletion, instead just preview \
                      deletions')
(options, args) = parser.parse_args()

TQ = options.tq
CMD_LOGS_DIR = options.cmdlogsdir
PURGE_LOG = options.logfile
DAYS = options.days
DELETE_JOBS = options.deletejobs
DRY_RUN = options.dryrun

####################################
# General setup

# Logging
logger = logging.getLogger('Tractor 2.2 purger')
hdlr = logging.FileHandler(PURGE_LOG)
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)
logger.setLevel(logging.INFO)

# Logging to stdout
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.INFO)
ch.setFormatter(formatter)
logger.addHandler(ch)

####################################
# Functions

def jobs_to_delete(days):
    """Create list of all jids (equivalient of all jobs to be deleted)
    """
    jids = []
    command = [TQ, 'jobs',
               'not active and not ready and spooltime  < -' + days + 'd',
               '--noheader', '--archives', '-c', 'jid']
    p = subprocess.Popen(command, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)

    try:
        for line in iter(p.stdout.readline, b''):
            sys.stdout.flush()
            jid = line.rstrip()
            jids.append(int(jid))
            logger.info('Added job for deletion: ' + jid)
    except:
        logger.warning('Failed to read stdout.')

    return jids

def get_all_job_folders(cmd_logs_dir):
    """Create list of all job folders
    """
    job_folders = []
    for root, directories, files in os.walk(cmd_logs_dir):
        if len(directories) > 0:
            for directory in directories:
                match = re.search('J\d*', directory)
                if match:
                    job_folder = root + '/' + directory
                    job_folders.append(job_folder)
    return job_folders

def get_job_deletion_list(job_folders, jids):
    """Compare job folders list against jids list, create deletion list
    """
    delete_list = []
    for job_folder in job_folders:
        jid_match = False
        for jid in jids:
            if job_folder.endswith('J' + str(jid)):
                jid_match = True

        if jid_match:
            delete_list.append(job_folder)
            logger.info('Added log folder for deletion: ' + job_folder)

    return delete_list

def delete_logs(delete_list):
    """Delete the actual log folders
    """
    for job_folder in delete_list:
        if not DRY_RUN:
            logger.info('Deleting ' + job_folder)
            shutil.rmtree(job_folder)
        else:
            logger.info('Dry run: (not) deleting ' + job_folder)

def delete_tractor_jobs(days):
    """Delete jobs from Tractor. You can also delete jobs manually using:
    tractor-dbctl --purge-archive-to-year-month YY-MM
    """
    if not DRY_RUN:
        logger.info('Executing tq command to delete jobs...')
        command = [TQ, '--force', '--yes', 'delete',
                   'not active and not ready and spooltime  < -' + days + 'd']
    else:
        logger.info('Executing tq command to (not) delete jobs...')
        command = [TQ, 'jobs', '--archives',
                   'not active and not ready and spooltime  < -' + days + 'd']

    p = subprocess.Popen(command, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)
    try:
        for line in iter(p.stdout.readline, b''):
            sys.stdout.flush()
            logger.info(line.rstrip())
    except:
        logger.warning('Failed reading stdout.')

####################################
# Main

if __name__ == '__main__':

    if not DRY_RUN:
        logger.info('Tractor purge initiated.')
    else:
        logger.info('Tractor purge initiated in "dry run" mode.')

    # Queries
    jids = jobs_to_delete(days=DAYS)
    job_folders = get_all_job_folders(cmd_logs_dir=CMD_LOGS_DIR)
    delete_list = get_job_deletion_list(job_folders=job_folders, jids=jids)

    # Summary
    logger.info('Jobs to be deleted: ' + str(len(jids)))
    logger.info('Job log folders found: ' + str(len(job_folders)))
    logger.info('Job log folders to be deleted: ' + str(len(delete_list)))

    # Delete logs
    if len(jids) > 0:
        delete_logs(delete_list=delete_list)
    else:
        logger.info('No logs to delete.')

    # Delete jobs
    if DELETE_JOBS:
        if len(delete_list) > 0:
            delete_tractor_jobs(days=DAYS)
        else:
            logger.info('No jobs to delete.')

    logger.info('Tractor purge done.\n')
<commit_msg>Fix additional formatting<commit_after>#!/usr/bin/env python

import sys
import os
import platform
import subprocess
import re
import shutil
import datetime
import logging
from optparse import OptionParser

####################################
# Option parser and constants
TRACTOR_PURGE_VERSION = 'v1.0.0'
parser = OptionParser(version='%prog ' + TRACTOR_PURGE_VERSION)
parser.add_option('-t', '--tq', dest='tq',
                  default='/opt/pixar/Tractor-2.2/bin/tq',
                  help='Absolute path to tq [default: %default]')
parser.add_option('-c', '--cmdlogsdir', dest='cmdlogsdir',
                  default='/var/spool/tractor/cmd-logs',
                  help='Absolute path to cmd-logs dir [default: %default]')
parser.add_option('-l', '--log', dest='logfile',
                  default='/var/tmp/tractor-purge.log',
                  help='Absolute path to tractor-purge log file '
                       '[default: %default]')
parser.add_option('-d', '--days', dest='days', default='30',
                  help='Number of days worth of jobs/logs to keep '
                       '[default: %default]')
parser.add_option('--deletejobs', action='store_true', dest='deletejobs',
                  default=False,
                  help='Delete jobs from psql database after log deletion. '
                       'If DBArchiving is True in Tractor config, archive '
                       'jobs instead.')
parser.add_option('--dryrun', action='store_true', dest='dryrun',
                  default=False,
                  help='Do not perform actual deletion, instead just preview \
                      deletions')
(options, args) = parser.parse_args()

TQ = options.tq
CMD_LOGS_DIR = options.cmdlogsdir
PURGE_LOG = options.logfile
DAYS = options.days
DELETE_JOBS = options.deletejobs
DRY_RUN = options.dryrun

####################################
# General setup

# Logging
logger = logging.getLogger('Tractor 2.2 purger')
hdlr = logging.FileHandler(PURGE_LOG)
formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)
logger.setLevel(logging.INFO)

# Logging to stdout
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.INFO)
ch.setFormatter(formatter)
logger.addHandler(ch)

####################################
# Functions

def jobs_to_delete(days):
    """Create list of all jids (equivalient of all jobs to be deleted)
    """
    jids = []
    command = [TQ, 'jobs',
               'not active and not ready and spooltime  < -' + days + 'd',
               '--noheader', '--archives', '-c', 'jid']
    p = subprocess.Popen(command, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)

    try:
        for line in iter(p.stdout.readline, b''):
            sys.stdout.flush()
            jid = line.rstrip()
            jids.append(int(jid))
            logger.info('Added job for deletion: ' + jid)
    except:
        logger.warning('Failed to read stdout.')

    return jids

def get_all_job_folders(cmd_logs_dir):
    """Create list of all job folders
    """
    job_folders = []
    for root, directories, files in os.walk(cmd_logs_dir):
        if len(directories) > 0:
            for directory in directories:
                match = re.search('J\d*', directory)
                if match:
                    job_folder = root + '/' + directory
                    job_folders.append(job_folder)
    return job_folders

def get_job_deletion_list(job_folders, jids):
    """Compare job folders list against jids list, create deletion list
    """
    delete_list = []
    for job_folder in job_folders:
        jid_match = False
        for jid in jids:
            if job_folder.endswith('J' + str(jid)):
                jid_match = True

        if jid_match:
            delete_list.append(job_folder)
            logger.info('Added log folder for deletion: ' + job_folder)

    return delete_list

def delete_logs(delete_list):
    """Delete the actual log folders
    """
    for job_folder in delete_list:
        if not DRY_RUN:
            logger.info('Deleting ' + job_folder)
            shutil.rmtree(job_folder)
        else:
            logger.info('Dry run: (not) deleting ' + job_folder)

def delete_tractor_jobs(days):
    """Delete jobs from Tractor. You can also delete jobs manually using:
    tractor-dbctl --purge-archive-to-year-month YY-MM
    """
    if not DRY_RUN:
        logger.info('Executing tq command to delete jobs...')
        command = [TQ, '--force', '--yes', 'delete',
                   'not active and not ready and spooltime  < -' + days + 'd']
    else:
        logger.info('Executing tq command to (not) delete jobs...')
        command = [TQ, 'jobs', '--archives',
                   'not active and not ready and spooltime  < -' + days + 'd']

    p = subprocess.Popen(command, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)
    try:
        for line in iter(p.stdout.readline, b''):
            sys.stdout.flush()
            logger.info(line.rstrip())
    except:
        logger.warning('Failed reading stdout.')

####################################
# Main

if __name__ == '__main__':

    if not DRY_RUN:
        logger.info('Tractor purge initiated.')
    else:
        logger.info('Tractor purge initiated in "dry run" mode.')

    # Queries
    jids = jobs_to_delete(days=DAYS)
    job_folders = get_all_job_folders(cmd_logs_dir=CMD_LOGS_DIR)
    delete_list = get_job_deletion_list(job_folders=job_folders, jids=jids)

    # Summary
    logger.info('Jobs to be deleted: ' + str(len(jids)))
    logger.info('Job log folders found: ' + str(len(job_folders)))
    logger.info('Job log folders to be deleted: ' + str(len(delete_list)))

    # Delete logs
    if len(jids) > 0:
        delete_logs(delete_list=delete_list)
    else:
        logger.info('No logs to delete.')

    # Delete jobs
    if DELETE_JOBS:
        if len(delete_list) > 0:
            delete_tractor_jobs(days=DAYS)
        else:
            logger.info('No jobs to delete.')

    logger.info('Tractor purge done.\n')
<|endoftext|>