# -*- coding: utf-8 -*- import os import six import sys import time import traceback from django.conf import settings from django.core.management.base import BaseCommand, CommandError from django.db.backends import utils from django.utils.six import PY3 from django.utils.datastructures import OrderedSet from django_extensions.management.shells import import_objects from django_extensions.management.utils import signalcommand def use_vi_mode(): editor = os.environ.get('EDITOR') if not editor: return False editor = os.path.basename(editor) return editor.startswith('vi') or editor.endswith('vim') class Command(BaseCommand): help = "Like the 'shell' command but autoloads the models of all installed Django apps." def add_arguments(self, parser): super(Command, self).add_arguments(parser) parser.add_argument( '--plain', action='store_true', dest='plain', help='Tells Django to use plain Python, not BPython nor IPython.') parser.add_argument( '--bpython', action='store_true', dest='bpython', help='Tells Django to use BPython, not IPython.') parser.add_argument( '--ptpython', action='store_true', dest='ptpython', help='Tells Django to use PTPython, not IPython.') parser.add_argument( '--ptipython', action='store_true', dest='ptipython', help='Tells Django to use PT-IPython, not IPython.') parser.add_argument( '--ipython', action='store_true', dest='ipython', help='Tells Django to use IPython, not BPython.') parser.add_argument( '--notebook', action='store_true', dest='notebook', help='Tells Django to use IPython Notebook.') parser.add_argument( '--kernel', action='store_true', dest='kernel', help='Tells Django to start an IPython Kernel.') parser.add_argument( '--connection-file', action='store', dest='connection_file', help='Specifies the connection file to use if using the --kernel option'), parser.add_argument( '--no-startup', action='store_true', dest='no_startup', help='When using plain Python, ignore the PYTHONSTARTUP environment variable and ~/.pythonrc.py script.') parser.add_argument( '--use-pythonrc', action='store_true', dest='use_pythonrc', help='When using plain Python, load the PYTHONSTARTUP environment variable and ~/.pythonrc.py script.') parser.add_argument( '--print-sql', action='store_true', default=False, help="Print SQL queries as they're executed") parser.add_argument( '--dont-load', action='append', dest='dont_load', default=[], help='Ignore autoloading of some apps/models. Can be used ' 'several times.') parser.add_argument( '--quiet-load', action='store_true', default=False, dest='quiet_load', help='Do not display loaded models messages') parser.add_argument( '--vi', action='store_true', default=use_vi_mode(), dest='vi_mode', help='Load Vi key bindings (for --ptpython and --ptipython)') parser.add_argument( '--no-browser', action='store_true', default=False, dest='no_browser', help='Don\'t open the notebook in a browser after startup.') def get_ipython_arguments(self, options): return getattr(settings, 'IPYTHON_ARGUMENTS', []) def get_notebook_arguments(self, options): notebook_args = 'NOTEBOOK_ARGUMENTS' arguments = getattr(settings, notebook_args, []) if not arguments: arguments = os.environ.get(notebook_args, '').split() return arguments def get_imported_objects(self, options): return import_objects(options, self.style) def get_kernel(self, options): try: from IPython import release if release.version_info[0] < 2: print(self.style.ERROR("--kernel requires at least IPython version 2.0")) return from IPython import start_kernel except ImportError: return traceback.format_exc() def run_kernel(): imported_objects = self.get_imported_objects(options) kwargs = dict( argv=[], user_ns=imported_objects, ) connection_file = options.get('connection_file') if connection_file: kwargs['connection_file'] = connection_file start_kernel(**kwargs) return run_kernel def get_notebook(self, options): from IPython import release try: from notebook.notebookapp import NotebookApp except ImportError: try: from IPython.html.notebookapp import NotebookApp except ImportError: if release.version_info[0] >= 3: raise try: from IPython.frontend.html.notebook import notebookapp NotebookApp = notebookapp.NotebookApp except ImportError: return traceback.format_exc() no_browser = options.get('no_browser', False) def install_kernel_spec(app, display_name, ipython_arguments): """install an IPython >= 3.0 kernelspec that loads django extensions""" ksm = app.kernel_spec_manager try_spec_names = getattr(settings, 'NOTEBOOK_KERNEL_SPEC_NAMES', [ 'python3' if PY3 else 'python2', 'python', ]) if isinstance(try_spec_names, six.string_types): try_spec_names = [try_spec_names] ks = None for spec_name in try_spec_names: try: ks = ksm.get_kernel_spec(spec_name) break except Exception: continue if not ks: raise CommandError("No notebook (Python) kernel specs found") ks.argv.extend(ipython_arguments) ks.display_name = display_name manage_py_dir, manage_py = os.path.split(os.path.realpath(sys.argv[0])) if manage_py == 'manage.py' and os.path.isdir(manage_py_dir) and manage_py_dir != os.getcwd(): pythonpath = ks.env.get('PYTHONPATH', os.environ.get('PYTHONPATH', '')) pythonpath = pythonpath.split(':') if manage_py_dir not in pythonpath: pythonpath.append(manage_py_dir) ks.env['PYTHONPATH'] = ':'.join(filter(None, pythonpath)) kernel_dir = os.path.join(ksm.user_kernel_dir, 'django_extensions') if not os.path.exists(kernel_dir): os.makedirs(kernel_dir) with open(os.path.join(kernel_dir, 'kernel.json'), 'w') as f: f.write(ks.to_json()) def run_notebook(): app = NotebookApp.instance() # Treat IPYTHON_ARGUMENTS from settings ipython_arguments = self.get_ipython_arguments(options) if 'django_extensions.management.notebook_extension' not in ipython_arguments: ipython_arguments.extend(['--ext', 'django_extensions.management.notebook_extension']) # Treat NOTEBOOK_ARGUMENTS from settings notebook_arguments = self.get_notebook_arguments(options) if no_browser and '--no-browser' not in notebook_arguments: notebook_arguments.append('--no-browser') if '--notebook-dir' not in notebook_arguments: notebook_arguments.extend(['--notebook-dir', '.']) # IPython < 3 passes through kernel args from notebook CLI if release.version_info[0] < 3: notebook_arguments.extend(ipython_arguments) app.initialize(notebook_arguments) # IPython >= 3 uses kernelspecs to specify kernel CLI args if release.version_info[0] >= 3: display_name = getattr(settings, 'IPYTHON_KERNEL_DISPLAY_NAME', "Django Shell-Plus") install_kernel_spec(app, display_name, ipython_arguments) app.start() return run_notebook def get_plain(self, options): # Using normal Python shell import code imported_objects = self.get_imported_objects(options) try: # Try activating rlcompleter, because it's handy. import readline except ImportError: pass else: # We don't have to wrap the following import in a 'try', because # we already know 'readline' was imported successfully. import rlcompleter readline.set_completer(rlcompleter.Completer(imported_objects).complete) readline.parse_and_bind("tab:complete") use_pythonrc = options.get('use_pythonrc', False) no_startup = options.get('no_startup', False) # We want to honor both $PYTHONSTARTUP and .pythonrc.py, so follow system # conventions and get $PYTHONSTARTUP first then .pythonrc.py. if use_pythonrc or not no_startup: for pythonrc in OrderedSet([os.environ.get("PYTHONSTARTUP"), os.path.expanduser('~/.pythonrc.py')]): if not pythonrc: continue if not os.path.isfile(pythonrc): continue # Match the behavior of the cpython shell where an error in # PYTHONSTARTUP prints an exception and continues. try: with open(pythonrc) as handle: exec(compile(handle.read(), pythonrc, 'exec'), imported_objects) in globals(), locals() except Exception: traceback.print_exc() def run_plain(): code.interact(local=imported_objects) return run_plain def get_bpython(self, options): try: from bpython import embed except ImportError: return traceback.format_exc() def run_bpython(): imported_objects = self.get_imported_objects(options) embed(imported_objects) return run_bpython def get_ipython(self, options): try: from IPython import start_ipython def run_ipython(): imported_objects = self.get_imported_objects(options) ipython_arguments = self.get_ipython_arguments(options) start_ipython(argv=ipython_arguments, user_ns=imported_objects) return run_ipython except ImportError: str_exc = traceback.format_exc() # IPython < 0.11 # Explicitly pass an empty list as arguments, because otherwise # IPython would use sys.argv from this script. # Notebook not supported for IPython < 0.11. try: from IPython.Shell import IPShell except ImportError: return str_exc + "\n" + traceback.format_exc() def run_ipython(): imported_objects = self.get_imported_objects(options) shell = IPShell(argv=[], user_ns=imported_objects) shell.mainloop() return run_ipython def get_ptpython(self, options): try: from ptpython.repl import embed, run_config except ImportError: tb = traceback.format_exc() try: # prompt_toolkit < v0.27 from prompt_toolkit.contrib.repl import embed, run_config except ImportError: return tb def run_ptpython(): imported_objects = self.get_imported_objects(options) history_filename = os.path.expanduser('~/.ptpython_history') embed(globals=imported_objects, history_filename=history_filename, vi_mode=options.get('vi_mode', False), configure=run_config) return run_ptpython def get_ptipython(self, options): try: from ptpython.repl import run_config from ptpython.ipython import embed except ImportError: tb = traceback.format_exc() try: # prompt_toolkit < v0.27 from prompt_toolkit.contrib.repl import run_config from prompt_toolkit.contrib.ipython import embed except ImportError: return tb def run_ptipython(): imported_objects = self.get_imported_objects(options) history_filename = os.path.expanduser('~/.ptpython_history') embed(user_ns=imported_objects, history_filename=history_filename, vi_mode=options.get('vi_mode', False), configure=run_config) return run_ptipython def set_application_name(self, options): """Set the application_name on PostgreSQL connection Use the fallback_application_name to let the user override it with PGAPPNAME env variable http://www.postgresql.org/docs/9.4/static/libpq-connect.html#LIBPQ-PARAMKEYWORDS # noqa """ supported_backends = ['django.db.backends.postgresql', 'django.db.backends.postgresql_psycopg2'] opt_name = 'fallback_application_name' default_app_name = 'django_shell' app_name = default_app_name dbs = getattr(settings, 'DATABASES', []) # lookup over all the databases entry for db in dbs.keys(): if dbs[db]['ENGINE'] in supported_backends: try: options = dbs[db]['OPTIONS'] except KeyError: options = {} # dot not override a defined value if opt_name in options.keys(): app_name = dbs[db]['OPTIONS'][opt_name] else: dbs[db].setdefault('OPTIONS', {}).update({opt_name: default_app_name}) app_name = default_app_name return app_name @signalcommand def handle(self, *args, **options): use_kernel = options.get('kernel', False) use_notebook = options.get('notebook', False) use_ipython = options.get('ipython', False) use_bpython = options.get('bpython', False) use_plain = options.get('plain', False) use_ptpython = options.get('ptpython', False) use_ptipython = options.get('ptipython', False) verbosity = int(options.get('verbosity', 1)) print_sql = getattr(settings, 'SHELL_PLUS_PRINT_SQL', False) truncate = getattr(settings, 'SHELL_PLUS_PRINT_SQL_TRUNCATE', 1000) if options.get("print_sql", False) or print_sql: # Code from http://gist.github.com/118990 try: import sqlparse except ImportError: sqlparse = None try: import pygments.lexers import pygments.formatters except ImportError: pygments = None class PrintQueryWrapper(utils.CursorDebugWrapper): def execute(self, sql, params=()): starttime = time.time() try: return self.cursor.execute(sql, params) finally: execution_time = time.time() - starttime raw_sql = self.db.ops.last_executed_query(self.cursor, sql, params) if sqlparse: raw_sql = raw_sql[:truncate] raw_sql = sqlparse.format(raw_sql, reindent_aligned=True, truncate_strings=500) if pygments: raw_sql = pygments.highlight( raw_sql, pygments.lexers.get_lexer_by_name("sql"), pygments.formatters.TerminalFormatter() ) print(raw_sql) print("") print('Execution time: %.6fs [Database: %s]' % (execution_time, self.db.alias)) print("") utils.CursorDebugWrapper = PrintQueryWrapper shells = ( ('ptipython', self.get_ptipython), ('ptpython', self.get_ptpython), ('bpython', self.get_bpython), ('ipython', self.get_ipython), ('plain', self.get_plain), ) SETTINGS_SHELL_PLUS = getattr(settings, 'SHELL_PLUS', None) shell = None shell_name = "any" self.set_application_name(options) if use_kernel: shell = self.get_kernel(options) shell_name = "IPython Kernel" elif use_notebook: shell = self.get_notebook(options) shell_name = "IPython Notebook" elif use_plain: shell = self.get_plain(options) shell_name = "plain" elif use_ipython: shell = self.get_ipython(options) shell_name = "IPython" elif use_bpython: shell = self.get_bpython(options) shell_name = "BPython" elif use_ptpython: shell = self.get_ptpython(options) shell_name = "ptpython" elif use_ptipython: shell = self.get_ptipython(options) shell_name = "ptipython" elif SETTINGS_SHELL_PLUS: shell_name = SETTINGS_SHELL_PLUS shell = dict(shells)[shell_name](options) else: for shell_name, func in shells: if verbosity > 2: print(self.style.NOTICE("Trying shell: %s" % shell_name)) shell = func(options) if callable(shell): if verbosity > 1: print(self.style.NOTICE("Using shell: %s" % shell_name)) break if not callable(shell): if shell: print(shell) print(self.style.ERROR("Could not load %s interactive Python environment." % shell_name)) return shell()