indico/cli/shell.py

   1 # This file is part of Indico.
   2 # Copyright (C) 2002 - 2015 European Organization for Nuclear Research (CERN).
   3 #
   4 # Indico is free software; you can redistribute it and/or
   5 # modify it under the terms of the GNU General Public License as
   6 # published by the Free Software Foundation; either version 3 of the
   7 # License, or (at your option) any later version.
   8 #
   9 # Indico is distributed in the hope that it will be useful, but
  10 # WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12 # General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License
  15 # along with Indico; if not, see <http://www.gnu.org/licenses/>.
  16
  17 import datetime
  18 import itertools
  19 from functools import partial
  20 from operator import itemgetter, attrgetter
  21
  22 import transaction
  23 from flask import current_app
  24 from flask_script import Shell, Option
  25 from werkzeug.local import LocalProxy
  26
  27 import MaKaC
  28 from indico.core import signals
  29 from indico.core.celery import celery
  30 from indico.core.config import Config
  31 from indico.util.console import colored, strip_ansi, cformat
  32 from indico.core.db import DBMgr, db
  33 from indico.core.index import Catalog
  34 from indico.core.plugins import plugin_engine
  35 from MaKaC.common import HelperMaKaCInfo
  36 from MaKaC.common.indexes import IndexesHolder
  37 from MaKaC.conference import Conference, ConferenceHolder, CategoryManager
  38 from indico.web.flask.util import IndicoConfigWrapper
  39
  40
  41 def _add_to_context(namespace, info, element, name=None, doc=None, color='green'):
  42     if not name:
  43         name = element.__name__
  44     namespace[name] = element
  45     if doc:
  46         info.append(cformat('+ %%{%s}{}%%{white!} ({})' % color).format(name, doc))
  47     else:
  48         info.append(cformat('+ %%{%s}{}' % color).format(name))
  49
  50
  51 def _add_to_context_multi(namespace, info, elements, names=None, doc=None, color='green'):
  52     if not names:
  53         names = [x.__name__ for x in elements]
  54     for name, element in zip(names, elements):
  55         namespace[name] = element
  56     if doc:
  57         info.append(cformat('+ %%{white!}{}:%%{reset} %%{%s}{}' % color).format(doc, ', '.join(names)))
  58     else:
  59         info.append(cformat('+ %%{%s}{}' % color).format(', '.join(names)))
  60
  61
  62 def _add_to_context_smart(namespace, info, objects, get_name=attrgetter('__name__'), color='cyan'):
  63     def _get_module(obj):
  64         segments = tuple(obj.__module__.split('.'))
  65         if segments[0].startswith('indico_'):  # plugin
  66             return 'plugin:{}'.format(segments[0])
  67         elif segments[:2] == ('indico', 'modules'):
  68             return 'module:{}'.format(segments[2])
  69         elif segments[:2] == ('indico', 'core'):
  70             return 'core:{}'.format(segments[2])
  71         else:
  72             return '.'.join(segments[:-1] if len(segments) > 1 else segments)
  73
  74     items = [(_get_module(obj), get_name(obj), obj) for obj in objects]
  75     for module, items in itertools.groupby(sorted(items, key=itemgetter(0, 1)), key=itemgetter(0)):
  76         names, elements = zip(*((x[1], x[2]) for x in items))
  77         _add_to_context_multi(namespace, info, elements, names, doc=module, color=color)
  78
  79
  80 class IndicoShell(Shell):
  81     def __init__(self):
  82         banner = colored('\nIndico v{} is ready for your commands!\n'.format(MaKaC.__version__),
  83                          'yellow', attrs=['bold'])
  84         super(IndicoShell, self).__init__(banner=banner, use_bpython=False)
  85         self._context = None
  86         self._info = None
  87         self._quiet = False
  88
  89     def run(self, no_ipython, use_bpython, quiet):
  90         context = self.get_context()
  91         if not quiet:
  92             self.banner = '\n'.join(self._info + [self.banner])
  93         if use_bpython:
  94             # bpython does not support escape sequences :(
  95             # https://github.com/bpython/bpython/issues/396
  96             self.banner = strip_ansi(self.banner)
  97         with context['dbi'].global_connection():
  98             super(IndicoShell, self).run(no_ipython or use_bpython, not use_bpython)
  99
 100     def get_options(self):
 101         return (
 102             Option('--no-ipython', action='store_true', dest='no_ipython', default=False,
 103                    help="Do not use the IPython shell"),
 104             Option('--use-bpython', action='store_true', dest='use_bpython', default=False,
 105                    help="Use the BPython shell"),
 106             Option('--quiet', '-q', action='store_true', dest='quiet', default=False,
 107                    help="Do not print the shell context")
 108         )
 109
 110     def get_context(self):
 111         if self._context is None:
 112             self._context = context = {}
 113             self._info = []
 114
 115             add_to_context = partial(_add_to_context, context, self._info)
 116             add_to_context_multi = partial(_add_to_context_multi, context, self._info)
 117             add_to_context_smart = partial(_add_to_context_smart, context, self._info)
 118             # Common stdlib modules
 119             self._info.append(cformat('*** %{magenta!}stdlib%{reset} ***'))
 120             add_to_context_multi([getattr(datetime, attr) for attr in ('date', 'time', 'datetime', 'timedelta')] +
 121                                  [itertools],
 122                                  color='yellow')
 123             # Legacy Indico
 124             self._info.append(cformat('*** %{magenta!}Legacy%{reset} ***'))
 125             add_to_context_multi([Conference, ConferenceHolder, CategoryManager, Catalog, IndexesHolder], color='green')
 126             add_to_context(LocalProxy(HelperMaKaCInfo.getMaKaCInfoInstance), 'minfo', color='green')
 127             # Models
 128             self._info.append(cformat('*** %{magenta!}Models%{reset} ***'))
 129             models = [cls for name, cls in sorted(db.Model._decl_class_registry.items(), key=itemgetter(0))
 130                       if hasattr(cls, '__table__')]
 131             add_to_context_smart(models)
 132             # Tasks
 133             self._info.append(cformat('*** %{magenta!}Tasks%{reset} ***'))
 134             celery.loader.import_default_modules()  # load all tasks
 135             tasks = [task for task in sorted(celery.tasks.values()) if not task.name.startswith('celery.')]
 136             add_to_context_smart(tasks, get_name=lambda x: x.name.replace('.', '_'), color='blue!')
 137             # Plugins
 138             self._info.append(cformat('*** %{magenta!}Plugins%{reset} ***'))
 139             plugins = [type(plugin) for plugin in sorted(plugin_engine.get_active_plugins().values(),
 140                                                          key=attrgetter('name'))]
 141             add_to_context_multi(plugins, color='yellow!')
 142             # Utils
 143             self._info.append(cformat('*** %{magenta!}Misc%{reset} ***'))
 144             add_to_context(celery, 'celery', doc='celery app', color='blue!')
 145             add_to_context(DBMgr.getInstance(), 'dbi', doc='zodb db interface', color='cyan!')
 146             add_to_context(db, 'db', doc='sqlalchemy db interface', color='cyan!')
 147             add_to_context(transaction, doc='transaction module', color='cyan!')
 148             add_to_context(IndicoConfigWrapper(Config.getInstance()), 'config', doc='indico config')
 149             add_to_context(current_app, 'app', doc='flask app')
 150             # Stuff from plugins
 151             signals.plugin.shell_context.send(add_to_context=add_to_context, add_to_context_multi=add_to_context_multi)
 152
 153         return self._context