# -*- coding: utf-8 -*- import inspect import os import sys import traceback import warnings from django.db import connections from django.conf import settings from django.core.management.base import BaseCommand, CommandError from django.utils.datastructures import OrderedSet from django_extensions.management.shells import import_objects from django_extensions.management.utils import signalcommand from django_extensions.management.debug_cursor import monkey_patch_cursordebugwrapper 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') def shell_runner(flags, name, help=None): """ Decorates methods with information about the application they are starting :param flags: The flags used to start this runner via the ArgumentParser. :param name: The name of this runner for the help text for the ArgumentParser. :param help: The optional help for the ArgumentParser if the dynamically generated help is not sufficient. """ def decorator(fn): fn.runner_flags = flags fn.runner_name = name fn.runner_help = help return fn return decorator class Command(BaseCommand): help = "Like the 'shell' command but autoloads the models of all installed Django apps." extra_args = None tests_mode = False def __init__(self): super().__init__() self.runners = [member for name, member in inspect.getmembers(self) if hasattr(member, 'runner_flags')] def add_arguments(self, parser): super().add_arguments(parser) group = parser.add_mutually_exclusive_group() for runner in self.runners: if runner.runner_help: help = runner.runner_help else: help = 'Tells Django to use %s.' % runner.runner_name group.add_argument( *runner.runner_flags, action='store_const', dest='runner', const=runner, help=help) 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', default=False, 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', default=False, 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( '--truncate-sql', action='store', type=int, help="Truncate SQL queries to a number of characters." ) parser.add_argument( '--print-sql-location', action='store_true', default=False, help="Show location in code where SQL query generated from" ) 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.' ) parser.add_argument( '-c', '--command', help='Instead of opening an interactive shell, run a command as Django and exit.', ) def run_from_argv(self, argv): if '--' in argv[2:]: idx = argv.index('--') self.extra_args = argv[idx + 1:] argv = argv[:idx] return super().run_from_argv(argv) def get_ipython_arguments(self, options): ipython_args = 'IPYTHON_ARGUMENTS' arguments = getattr(settings, ipython_args, []) if not arguments: arguments = os.environ.get(ipython_args, '').split() return 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): imported_objects = import_objects(options, self.style) if self.tests_mode: # save imported objects so we can run tests against it later self.tests_imported_objects = imported_objects return imported_objects @shell_runner(flags=['--kernel'], name='IPython Kernel') 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['connection_file'] if connection_file: kwargs['connection_file'] = connection_file start_kernel(**kwargs) return run_kernel def load_base_kernel_spec(self, app): """Finds and returns the base Python kernelspec to extend from.""" ksm = app.kernel_spec_manager try_spec_names = getattr(settings, 'NOTEBOOK_KERNEL_SPEC_NAMES', [ 'python3', 'python', ]) if isinstance(try_spec_names, str): 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. Tried %r" % try_spec_names) return ks def generate_kernel_specs(self, app, ipython_arguments): """Generate an IPython >= 3.0 kernelspec that loads django extensions""" ks = self.load_base_kernel_spec(app) ks.argv.extend(ipython_arguments) ks.display_name = getattr(settings, 'IPYTHON_KERNEL_DISPLAY_NAME', "Django Shell-Plus") 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): pythonpath = ks.env.get('PYTHONPATH', os.environ.get('PYTHONPATH', '')) pythonpath = pythonpath.split(os.pathsep) if manage_py_dir not in pythonpath: pythonpath.append(manage_py_dir) ks.env['PYTHONPATH'] = os.pathsep.join(filter(None, pythonpath)) return {'django_extensions': ks} def run_notebookapp(self, app_init, options, use_kernel_specs=True, history=True): no_browser = options['no_browser'] if self.extra_args: # if another '--' is found split the arguments notebook, ipython if '--' in self.extra_args: idx = self.extra_args.index('--') notebook_arguments = self.extra_args[:idx] ipython_arguments = self.extra_args[idx + 1:] # otherwise pass the arguments to the notebook else: notebook_arguments = self.extra_args ipython_arguments = [] else: notebook_arguments = self.get_notebook_arguments(options) ipython_arguments = self.get_ipython_arguments(options) # Treat IPYTHON_ARGUMENTS from settings if 'django_extensions.management.notebook_extension' not in ipython_arguments: ipython_arguments.extend(['--ext', 'django_extensions.management.notebook_extension']) # Treat NOTEBOOK_ARGUMENTS from settings if no_browser and '--no-browser' not in notebook_arguments: notebook_arguments.append('--no-browser') if '--notebook-dir' not in notebook_arguments and not any(e.startswith('--notebook-dir=') for e in notebook_arguments): notebook_arguments.extend(['--notebook-dir', '.']) # IPython < 3 passes through kernel args from notebook CLI if not use_kernel_specs: notebook_arguments.extend(ipython_arguments) # disable history if not already configured in some other way if not history and not any(arg.startswith('--HistoryManager') for arg in ipython_arguments): ipython_arguments.append('--HistoryManager.enabled=False') if not callable(app_init): app = app_init warnings.warn('Initialize should be a callable not an app instance', DeprecationWarning) app.initialize(notebook_arguments) else: app = app_init(notebook_arguments) # IPython >= 3 uses kernelspecs to specify kernel CLI args if use_kernel_specs: ksm = app.kernel_spec_manager for kid, ks in self.generate_kernel_specs(app, ipython_arguments).items(): roots = [os.path.dirname(ks.resource_dir), ksm.user_kernel_dir] for root in roots: kernel_dir = os.path.join(root, kid) try: 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()) break except OSError: continue else: raise CommandError('Could not write kernel %r in directories %r' % (kid, roots)) app.start() @shell_runner(flags=['--notebook'], name='IPython Notebook') def get_notebook(self, options): try: from IPython import release except ImportError: return traceback.format_exc() try: from notebook.notebookapp import NotebookApp except ImportError: if release.version_info[0] >= 7: return traceback.format_exc() try: from IPython.html.notebookapp import NotebookApp except ImportError: if release.version_info[0] >= 3: return traceback.format_exc() try: from IPython.frontend.html.notebook import notebookapp NotebookApp = notebookapp.NotebookApp except ImportError: return traceback.format_exc() use_kernel_specs = release.version_info[0] >= 3 def app_init(*args, **kwargs): app = NotebookApp.instance() app.initialize(*args, **kwargs) return app def run_notebook(): self.run_notebookapp(app_init, options, use_kernel_specs) return run_notebook @shell_runner(flags=['--lab'], name='JupyterLab Notebook') def get_jupyterlab(self, options): try: from jupyterlab.labapp import LabApp except ImportError: return traceback.format_exc() # check for JupyterLab 3.0 try: from notebook.notebookapp import NotebookApp except ImportError: NotebookApp = None if not NotebookApp or not issubclass(LabApp, NotebookApp): app_init = LabApp.initialize_server else: def app_init(*args, **kwargs): app = LabApp.instance() app.initialize(*args, **kwargs) return app def run_jupyterlab(): self.run_notebookapp(app_init, options, history=False) return run_jupyterlab @shell_runner(flags=['--plain'], name='plain Python') 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) # Enable tab completion on systems using libedit (e.g. macOS). # These lines are copied from Lib/site.py on Python 3.4. readline_doc = getattr(readline, '__doc__', '') if readline_doc is not None and 'libedit' in readline_doc: readline.parse_and_bind("bind ^I rl_complete") else: readline.parse_and_bind("tab:complete") use_pythonrc = options['use_pythonrc'] no_startup = options['no_startup'] # 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 with open(pythonrc) as handle: pythonrc_code = handle.read() # Match the behavior of the cpython shell where an error in # PYTHONSTARTUP prints an exception and continues. try: exec(compile(pythonrc_code, pythonrc, 'exec'), imported_objects) except Exception: traceback.print_exc() if self.tests_mode: raise def run_plain(): code.interact(local=imported_objects) return run_plain @shell_runner(flags=['--bpython'], name='BPython') 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) kwargs = {} if self.extra_args: kwargs['args'] = self.extra_args embed(imported_objects, **kwargs) return run_bpython @shell_runner(flags=['--ipython'], name='IPython') def get_ipython(self, options): try: from IPython import start_ipython def run_ipython(): imported_objects = self.get_imported_objects(options) ipython_arguments = self.extra_args or 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 @shell_runner(flags=['--ptpython'], name='PTPython') 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['vi_mode'], configure=run_config) return run_ptpython @shell_runner(flags=['--ptipython'], name='PT-IPython') 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['vi_mode'], configure=run_config) return run_ptipython @shell_runner(flags=['--idle'], name='Idle') def get_idle(self, options): from idlelib.pyshell import main def run_idle(): sys.argv = [ sys.argv[0], '-c', """ from django_extensions.management import shells from django.core.management.color import no_style for k, m in shells.import_objects({}, no_style()).items(): globals()[k] = m """, ] main() return run_idle 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' dbs = getattr(settings, 'DATABASES', []) for connection in connections.all(): alias = connection.alias mro = inspect.getmro(connection.__class__) if any(klass.__module__.startswith(supported_backends) for klass in mro): if 'OPTIONS' not in dbs[alias] or opt_name not in dbs[alias]['OPTIONS']: dbs[alias].setdefault('OPTIONS', {}).update({opt_name: default_app_name}) @signalcommand def handle(self, *args, **options): verbosity = options["verbosity"] get_runner = options['runner'] print_sql = getattr(settings, 'SHELL_PLUS_PRINT_SQL', False) runner = None runner_name = None truncate = None if options["truncate_sql"] == 0 else options["truncate_sql"] with monkey_patch_cursordebugwrapper(print_sql=options["print_sql"] or print_sql, truncate=truncate, print_sql_location=options["print_sql_location"], confprefix="SHELL_PLUS"): SETTINGS_SHELL_PLUS = getattr(settings, 'SHELL_PLUS', None) def get_runner_by_flag(flag): for runner in self.runners: if flag in runner.runner_flags: return runner return None self.set_application_name(options) if not get_runner and SETTINGS_SHELL_PLUS: get_runner = get_runner_by_flag('--%s' % SETTINGS_SHELL_PLUS) if not get_runner: runner = None runner_name = SETTINGS_SHELL_PLUS if get_runner: runner = get_runner(options) runner_name = get_runner.runner_name else: def try_runner(get_runner): runner_name = get_runner.runner_name if verbosity > 2: print(self.style.NOTICE("Trying: %s" % runner_name)) runner = get_runner(options) if callable(runner): if verbosity > 1: print(self.style.NOTICE("Using: %s" % runner_name)) return runner return None tried_runners = set() # try the runners that are least unexpected (normal shell runners) preferred_runners = ['ptipython', 'ptpython', 'bpython', 'ipython', 'plain'] for flag_suffix in preferred_runners: get_runner = get_runner_by_flag('--%s' % flag_suffix) tried_runners.add(get_runner) runner = try_runner(get_runner) if runner: runner_name = get_runner.runner_name break # try any remaining runners if needed if not runner: for get_runner in self.runners: if get_runner not in tried_runners: runner = try_runner(get_runner) if runner: runner_name = get_runner.runner_name break if not callable(runner): if runner: print(runner) if not runner_name: raise CommandError("No shell runner could be found.") raise CommandError("Could not load shell runner: '%s'." % runner_name) if self.tests_mode: return 130 if options['command']: imported_objects = self.get_imported_objects(options) exec(options['command'], imported_objects) return None runner()