1# -*- coding: utf-8 -*- 2import inspect 3import os 4import sys 5import traceback 6 7from django.conf import settings 8from django.core.management.base import BaseCommand, CommandError 9from django.utils.datastructures import OrderedSet 10 11from django_extensions.management.shells import import_objects 12from django_extensions.management.utils import signalcommand 13from django_extensions.management.debug_cursor import monkey_patch_cursordebugwrapper 14 15 16def use_vi_mode(): 17 editor = os.environ.get('EDITOR') 18 if not editor: 19 return False 20 editor = os.path.basename(editor) 21 return editor.startswith('vi') or editor.endswith('vim') 22 23 24def shell_runner(flags, name, help=None): 25 """ 26 Decorates methods with information about the application they are starting 27 28 :param flags: The flags used to start this runner via the ArgumentParser. 29 :param name: The name of this runner for the help text for the ArgumentParser. 30 :param help: The optional help for the ArgumentParser if the dynamically generated help is not sufficient. 31 """ 32 33 def decorator(fn): 34 fn.runner_flags = flags 35 fn.runner_name = name 36 fn.runner_help = help 37 38 return fn 39 40 return decorator 41 42 43class Command(BaseCommand): 44 help = "Like the 'shell' command but autoloads the models of all installed Django apps." 45 extra_args = None 46 tests_mode = False 47 48 def __init__(self): 49 super().__init__() 50 self.runners = [member for name, member in inspect.getmembers(self) 51 if hasattr(member, 'runner_flags')] 52 53 def add_arguments(self, parser): 54 super().add_arguments(parser) 55 56 group = parser.add_mutually_exclusive_group() 57 for runner in self.runners: 58 if runner.runner_help: 59 help = runner.runner_help 60 else: 61 help = 'Tells Django to use %s.' % runner.runner_name 62 63 group.add_argument( 64 *runner.runner_flags, action='store_const', dest='runner', const=runner, help=help) 65 66 parser.add_argument( 67 '--connection-file', action='store', dest='connection_file', 68 help='Specifies the connection file to use if using the --kernel option' 69 ) 70 parser.add_argument( 71 '--no-startup', action='store_true', dest='no_startup', 72 default=False, 73 help='When using plain Python, ignore the PYTHONSTARTUP environment variable and ~/.pythonrc.py script.' 74 ) 75 parser.add_argument( 76 '--use-pythonrc', action='store_true', dest='use_pythonrc', 77 default=False, 78 help='When using plain Python, load the PYTHONSTARTUP environment variable and ~/.pythonrc.py script.' 79 ) 80 parser.add_argument( 81 '--print-sql', action='store_true', 82 default=False, 83 help="Print SQL queries as they're executed" 84 ) 85 parser.add_argument( 86 '--print-sql-location', action='store_true', 87 default=False, 88 help="Show location in code where SQL query generated from" 89 ) 90 parser.add_argument( 91 '--dont-load', action='append', dest='dont_load', default=[], 92 help='Ignore autoloading of some apps/models. Can be used several times.' 93 ) 94 parser.add_argument( 95 '--quiet-load', action='store_true', 96 default=False, 97 dest='quiet_load', help='Do not display loaded models messages' 98 ) 99 parser.add_argument( 100 '--vi', action='store_true', default=use_vi_mode(), dest='vi_mode', 101 help='Load Vi key bindings (for --ptpython and --ptipython)' 102 ) 103 parser.add_argument( 104 '--no-browser', action='store_true', 105 default=False, 106 dest='no_browser', 107 help='Don\'t open the notebook in a browser after startup.' 108 ) 109 parser.add_argument( 110 '-c', '--command', 111 help='Instead of opening an interactive shell, run a command as Django and exit.', 112 ) 113 114 def run_from_argv(self, argv): 115 if '--' in argv[2:]: 116 idx = argv.index('--') 117 self.extra_args = argv[idx + 1:] 118 argv = argv[:idx] 119 return super().run_from_argv(argv) 120 121 def get_ipython_arguments(self, options): 122 ipython_args = 'IPYTHON_ARGUMENTS' 123 arguments = getattr(settings, ipython_args, []) 124 if not arguments: 125 arguments = os.environ.get(ipython_args, '').split() 126 return arguments 127 128 def get_notebook_arguments(self, options): 129 notebook_args = 'NOTEBOOK_ARGUMENTS' 130 arguments = getattr(settings, notebook_args, []) 131 if not arguments: 132 arguments = os.environ.get(notebook_args, '').split() 133 return arguments 134 135 def get_imported_objects(self, options): 136 imported_objects = import_objects(options, self.style) 137 if self.tests_mode: 138 # save imported objects so we can run tests against it later 139 self.tests_imported_objects = imported_objects 140 return imported_objects 141 142 @shell_runner(flags=['--kernel'], name='IPython Kernel') 143 def get_kernel(self, options): 144 try: 145 from IPython import release 146 if release.version_info[0] < 2: 147 print(self.style.ERROR("--kernel requires at least IPython version 2.0")) 148 return 149 from IPython import start_kernel 150 except ImportError: 151 return traceback.format_exc() 152 153 def run_kernel(): 154 imported_objects = self.get_imported_objects(options) 155 kwargs = dict( 156 argv=[], 157 user_ns=imported_objects, 158 ) 159 connection_file = options['connection_file'] 160 if connection_file: 161 kwargs['connection_file'] = connection_file 162 start_kernel(**kwargs) 163 return run_kernel 164 165 def load_base_kernel_spec(self, app): 166 """Finds and returns the base Python kernelspec to extend from.""" 167 ksm = app.kernel_spec_manager 168 try_spec_names = getattr(settings, 'NOTEBOOK_KERNEL_SPEC_NAMES', [ 169 'python3', 170 'python', 171 ]) 172 173 if isinstance(try_spec_names, str): 174 try_spec_names = [try_spec_names] 175 176 ks = None 177 for spec_name in try_spec_names: 178 try: 179 ks = ksm.get_kernel_spec(spec_name) 180 break 181 except Exception: 182 continue 183 if not ks: 184 raise CommandError("No notebook (Python) kernel specs found. Tried %r" % try_spec_names) 185 186 return ks 187 188 def generate_kernel_specs(self, app, ipython_arguments): 189 """Generate an IPython >= 3.0 kernelspec that loads django extensions""" 190 ks = self.load_base_kernel_spec(app) 191 ks.argv.extend(ipython_arguments) 192 ks.display_name = getattr(settings, 'IPYTHON_KERNEL_DISPLAY_NAME', "Django Shell-Plus") 193 194 manage_py_dir, manage_py = os.path.split(os.path.realpath(sys.argv[0])) 195 if manage_py == 'manage.py' and os.path.isdir(manage_py_dir): 196 pythonpath = ks.env.get('PYTHONPATH', os.environ.get('PYTHONPATH', '')) 197 pythonpath = pythonpath.split(os.pathsep) 198 if manage_py_dir not in pythonpath: 199 pythonpath.append(manage_py_dir) 200 201 ks.env['PYTHONPATH'] = os.pathsep.join(filter(None, pythonpath)) 202 203 return {'django_extensions': ks} 204 205 def run_notebookapp(self, app, options, use_kernel_specs=True): 206 no_browser = options['no_browser'] 207 208 if self.extra_args: 209 # if another '--' is found split the arguments notebook, ipython 210 if '--' in self.extra_args: 211 idx = self.extra_args.index('--') 212 notebook_arguments = self.extra_args[:idx] 213 ipython_arguments = self.extra_args[idx + 1:] 214 # otherwise pass the arguments to the notebook 215 else: 216 notebook_arguments = self.extra_args 217 ipython_arguments = [] 218 else: 219 notebook_arguments = self.get_notebook_arguments(options) 220 ipython_arguments = self.get_ipython_arguments(options) 221 222 # Treat IPYTHON_ARGUMENTS from settings 223 if 'django_extensions.management.notebook_extension' not in ipython_arguments: 224 ipython_arguments.extend(['--ext', 'django_extensions.management.notebook_extension']) 225 226 # Treat NOTEBOOK_ARGUMENTS from settings 227 if no_browser and '--no-browser' not in notebook_arguments: 228 notebook_arguments.append('--no-browser') 229 if '--notebook-dir' not in notebook_arguments and not any(e.startswith('--notebook-dir=') for e in notebook_arguments): 230 notebook_arguments.extend(['--notebook-dir', '.']) 231 232 # IPython < 3 passes through kernel args from notebook CLI 233 if not use_kernel_specs: 234 notebook_arguments.extend(ipython_arguments) 235 236 app.initialize(notebook_arguments) 237 238 # IPython >= 3 uses kernelspecs to specify kernel CLI args 239 if use_kernel_specs: 240 ksm = app.kernel_spec_manager 241 for kid, ks in self.generate_kernel_specs(app, ipython_arguments).items(): 242 roots = [os.path.dirname(ks.resource_dir), ksm.user_kernel_dir] 243 success = False 244 for root in roots: 245 kernel_dir = os.path.join(root, kid) 246 try: 247 if not os.path.exists(kernel_dir): 248 os.makedirs(kernel_dir) 249 250 with open(os.path.join(kernel_dir, 'kernel.json'), 'w') as f: 251 f.write(ks.to_json()) 252 253 success = True 254 break 255 except OSError: 256 continue 257 258 if not success: 259 raise CommandError("Could not write kernel %r in directories %r" % (kid, roots)) 260 261 app.start() 262 263 @shell_runner(flags=['--notebook'], name='IPython Notebook') 264 def get_notebook(self, options): 265 try: 266 from IPython import release 267 except ImportError: 268 return traceback.format_exc() 269 try: 270 from notebook.notebookapp import NotebookApp 271 except ImportError: 272 if release.version_info[0] >= 7: 273 return traceback.format_exc() 274 try: 275 from IPython.html.notebookapp import NotebookApp 276 except ImportError: 277 if release.version_info[0] >= 3: 278 return traceback.format_exc() 279 try: 280 from IPython.frontend.html.notebook import notebookapp 281 NotebookApp = notebookapp.NotebookApp 282 except ImportError: 283 return traceback.format_exc() 284 285 use_kernel_specs = release.version_info[0] >= 3 286 287 def run_notebook(): 288 app = NotebookApp.instance() 289 self.run_notebookapp(app, options, use_kernel_specs) 290 291 return run_notebook 292 293 @shell_runner(flags=['--lab'], name='JupyterLab Notebook') 294 def get_jupyterlab(self, options): 295 try: 296 from jupyterlab.labapp import LabApp 297 except ImportError: 298 return traceback.format_exc() 299 300 def run_jupyterlab(): 301 app = LabApp.instance() 302 self.run_notebookapp(app, options) 303 304 return run_jupyterlab 305 306 @shell_runner(flags=['--plain'], name='plain Python') 307 def get_plain(self, options): 308 # Using normal Python shell 309 import code 310 imported_objects = self.get_imported_objects(options) 311 try: 312 # Try activating rlcompleter, because it's handy. 313 import readline 314 except ImportError: 315 pass 316 else: 317 # We don't have to wrap the following import in a 'try', because 318 # we already know 'readline' was imported successfully. 319 import rlcompleter 320 readline.set_completer(rlcompleter.Completer(imported_objects).complete) 321 # Enable tab completion on systems using libedit (e.g. macOS). 322 # These lines are copied from Lib/site.py on Python 3.4. 323 readline_doc = getattr(readline, '__doc__', '') 324 if readline_doc is not None and 'libedit' in readline_doc: 325 readline.parse_and_bind("bind ^I rl_complete") 326 else: 327 readline.parse_and_bind("tab:complete") 328 329 use_pythonrc = options['use_pythonrc'] 330 no_startup = options['no_startup'] 331 332 # We want to honor both $PYTHONSTARTUP and .pythonrc.py, so follow system 333 # conventions and get $PYTHONSTARTUP first then .pythonrc.py. 334 if use_pythonrc or not no_startup: 335 for pythonrc in OrderedSet([os.environ.get("PYTHONSTARTUP"), os.path.expanduser('~/.pythonrc.py')]): 336 if not pythonrc: 337 continue 338 if not os.path.isfile(pythonrc): 339 continue 340 with open(pythonrc) as handle: 341 pythonrc_code = handle.read() 342 # Match the behavior of the cpython shell where an error in 343 # PYTHONSTARTUP prints an exception and continues. 344 try: 345 exec(compile(pythonrc_code, pythonrc, 'exec'), imported_objects) 346 except Exception: 347 traceback.print_exc() 348 if self.tests_mode: 349 raise 350 351 def run_plain(): 352 code.interact(local=imported_objects) 353 return run_plain 354 355 @shell_runner(flags=['--bpython'], name='BPython') 356 def get_bpython(self, options): 357 try: 358 from bpython import embed 359 except ImportError: 360 return traceback.format_exc() 361 362 def run_bpython(): 363 imported_objects = self.get_imported_objects(options) 364 kwargs = {} 365 if self.extra_args: 366 kwargs['args'] = self.extra_args 367 embed(imported_objects, **kwargs) 368 return run_bpython 369 370 @shell_runner(flags=['--ipython'], name='IPython') 371 def get_ipython(self, options): 372 try: 373 from IPython import start_ipython 374 375 def run_ipython(): 376 imported_objects = self.get_imported_objects(options) 377 ipython_arguments = self.extra_args or self.get_ipython_arguments(options) 378 start_ipython(argv=ipython_arguments, user_ns=imported_objects) 379 return run_ipython 380 except ImportError: 381 str_exc = traceback.format_exc() 382 # IPython < 0.11 383 # Explicitly pass an empty list as arguments, because otherwise 384 # IPython would use sys.argv from this script. 385 # Notebook not supported for IPython < 0.11. 386 try: 387 from IPython.Shell import IPShell 388 except ImportError: 389 return str_exc + "\n" + traceback.format_exc() 390 391 def run_ipython(): 392 imported_objects = self.get_imported_objects(options) 393 shell = IPShell(argv=[], user_ns=imported_objects) 394 shell.mainloop() 395 return run_ipython 396 397 @shell_runner(flags=['--ptpython'], name='PTPython') 398 def get_ptpython(self, options): 399 try: 400 from ptpython.repl import embed, run_config 401 except ImportError: 402 tb = traceback.format_exc() 403 try: # prompt_toolkit < v0.27 404 from prompt_toolkit.contrib.repl import embed, run_config 405 except ImportError: 406 return tb 407 408 def run_ptpython(): 409 imported_objects = self.get_imported_objects(options) 410 history_filename = os.path.expanduser('~/.ptpython_history') 411 embed(globals=imported_objects, history_filename=history_filename, 412 vi_mode=options['vi_mode'], configure=run_config) 413 return run_ptpython 414 415 @shell_runner(flags=['--ptipython'], name='PT-IPython') 416 def get_ptipython(self, options): 417 try: 418 from ptpython.repl import run_config 419 from ptpython.ipython import embed 420 except ImportError: 421 tb = traceback.format_exc() 422 try: # prompt_toolkit < v0.27 423 from prompt_toolkit.contrib.repl import run_config 424 from prompt_toolkit.contrib.ipython import embed 425 except ImportError: 426 return tb 427 428 def run_ptipython(): 429 imported_objects = self.get_imported_objects(options) 430 history_filename = os.path.expanduser('~/.ptpython_history') 431 embed(user_ns=imported_objects, history_filename=history_filename, 432 vi_mode=options['vi_mode'], configure=run_config) 433 return run_ptipython 434 435 @shell_runner(flags=['--idle'], name='Idle') 436 def get_idle(self, options): 437 from idlelib.pyshell import main 438 439 def run_idle(): 440 sys.argv = [ 441 sys.argv[0], 442 '-c', 443 """ 444from django_extensions.management import shells 445from django.core.management.color import no_style 446for k, m in shells.import_objects({}, no_style()).items(): 447 globals()[k] = m 448""", 449 ] 450 main() 451 452 return run_idle 453 454 def set_application_name(self, options): 455 """ 456 Set the application_name on PostgreSQL connection 457 458 Use the fallback_application_name to let the user override 459 it with PGAPPNAME env variable 460 461 http://www.postgresql.org/docs/9.4/static/libpq-connect.html#LIBPQ-PARAMKEYWORDS # noqa 462 """ 463 supported_backends = ['django.db.backends.postgresql', 464 'django.db.backends.postgresql_psycopg2'] 465 opt_name = 'fallback_application_name' 466 default_app_name = 'django_shell' 467 app_name = default_app_name 468 dbs = getattr(settings, 'DATABASES', []) 469 470 # lookup over all the databases entry 471 for db in dbs.keys(): 472 if dbs[db]['ENGINE'] in supported_backends: 473 try: 474 options = dbs[db]['OPTIONS'] 475 except KeyError: 476 options = {} 477 478 # dot not override a defined value 479 if opt_name in options.keys(): 480 app_name = dbs[db]['OPTIONS'][opt_name] 481 else: 482 dbs[db].setdefault('OPTIONS', {}).update({opt_name: default_app_name}) 483 app_name = default_app_name 484 485 return app_name 486 487 @signalcommand 488 def handle(self, *args, **options): 489 verbosity = options["verbosity"] 490 get_runner = options['runner'] 491 print_sql = getattr(settings, 'SHELL_PLUS_PRINT_SQL', False) 492 runner = None 493 runner_name = None 494 495 with monkey_patch_cursordebugwrapper(print_sql=options["print_sql"] or print_sql, print_sql_location=options["print_sql_location"], confprefix="SHELL_PLUS"): 496 SETTINGS_SHELL_PLUS = getattr(settings, 'SHELL_PLUS', None) 497 498 def get_runner_by_flag(flag): 499 for runner in self.runners: 500 if flag in runner.runner_flags: 501 return runner 502 return None 503 504 self.set_application_name(options) 505 506 if not get_runner and SETTINGS_SHELL_PLUS: 507 get_runner = get_runner_by_flag('--%s' % SETTINGS_SHELL_PLUS) 508 if not get_runner: 509 runner = None 510 runner_name = SETTINGS_SHELL_PLUS 511 512 if get_runner: 513 runner = get_runner(options) 514 runner_name = get_runner.runner_name 515 else: 516 def try_runner(get_runner): 517 runner_name = get_runner.runner_name 518 if verbosity > 2: 519 print(self.style.NOTICE("Trying: %s" % runner_name)) 520 521 runner = get_runner(options) 522 if callable(runner): 523 if verbosity > 1: 524 print(self.style.NOTICE("Using: %s" % runner_name)) 525 return runner 526 return None 527 528 tried_runners = set() 529 530 # try the runners that are least unexpected (normal shell runners) 531 preferred_runners = ['ptipython', 'ptpython', 'bpython', 'ipython', 'plain'] 532 for flag_suffix in preferred_runners: 533 get_runner = get_runner_by_flag('--%s' % flag_suffix) 534 tried_runners.add(get_runner) 535 runner = try_runner(get_runner) 536 if runner: 537 runner_name = get_runner.runner_name 538 break 539 540 # try any remaining runners if needed 541 if not runner: 542 for get_runner in self.runners: 543 if get_runner not in tried_runners: 544 runner = try_runner(get_runner) 545 if runner: 546 runner_name = get_runner.runner_name 547 break 548 549 if not callable(runner): 550 if runner: 551 print(runner) 552 if not runner_name: 553 raise CommandError("No shell runner could be found.") 554 raise CommandError("Could not load shell runner: '%s'." % runner_name) 555 556 if self.tests_mode: 557 return 130 558 559 if options['command']: 560 imported_objects = self.get_imported_objects(options) 561 exec(options['command'], {}, imported_objects) 562 return 563 564 runner() 565