Coverage for shellcraft / management / commands / shellcraft.py: 98%

1import warnings 

2 

3from django.conf import settings 

4from django.core.management.base import BaseCommand, CommandError 

5 

6BANNER = """\ 

7shellcraft {version} — Rails-style Django shell 

8 User.find(1) User.all() User.where(active=True) 

9 User.first() User.last() User.count() 

10 user.update(k=v) user.destroy() 

11 User.fields() tables() 

12""" 

13 

14 

15class Command(BaseCommand): 

16 help = "Rails-style Django shell with pretty printing and model shortcuts" 

17 

18 def handle(self, *_args, **_options): 

19 # Use CommandError instead of SystemExit so Django's management 

20 # framework can format the message via stderr and return a proper 

21 # non-zero exit code. SystemExit would skip that pipeline entirely. 

22 if not settings.DEBUG: 

23 raise CommandError( 

24 "shellcraft: refusing to run with DEBUG=False (production guard)" 

25 ) 

26 

27 from shellcraft.mixins import inject 

28 

29 inject() 

30 

31 namespace = self._build_namespace() 

32 self._start_shell(namespace) 

33 

34 def _build_namespace(self): 

35 from django.apps import apps 

36 

37 from shellcraft.printer import CYAN, RESET, ap 

38 

39 def tables(): 

40 all_models = apps.get_models() 

41 for model in all_models: 

42 label = model._meta.app_label 

43 self.stdout.write(f" {CYAN}{label}{RESET}.{model.__name__}") 

44 self.stdout.write(f"\n{len(all_models)} tables") 

45 

46 namespace = { 

47 "ap": ap, 

48 "tables": tables, 

49 } 

50 

51 # Build model-name → model mapping, warning when two apps define a 

52 # model with the same class name. Without the warning the second 

53 # model silently shadows the first, which is a confusing footgun when 

54 # working with multi-app projects (e.g. two apps both define `Profile`). 

55 seen: dict = {} 

56 for model in apps.get_models(): 

57 name = model.__name__ 

58 if name in seen: 

59 warnings.warn( 

60 f"shellcraft: model name conflict — '{name}' exists in both " 

61 f"'{seen[name]._meta.app_label}' and '{model._meta.app_label}'. " 

62 f"The '{model._meta.app_label}.{name}' version is used in the " 

63 "shell namespace. Access the other via " 

64 "apps.get_model('<app_label>', '<ModelName>').", 

65 stacklevel=2, 

66 ) 

67 seen[name] = model 

68 namespace[name] = model 

69 

70 return namespace 

71 

72 def _start_shell(self, namespace): 

73 from shellcraft import __version__ 

74 

75 banner = BANNER.format(version=__version__) 

76 self.stdout.write(banner) 

77 

78 # Detect which shell is available before entering it, so that 

79 # code.interact is never called inside an `except` block. If it were, 

80 # Python would set __context__ on every shell exception to the 

81 # ModuleNotFoundError, producing spurious "During handling of the above 

82 # exception" noise on every traceback the user sees. 

83 use_ipython = False 

84 try: 

85 import IPython # noqa: F401 

86 use_ipython = True 

87 except ImportError: 

88 self.stdout.write("(IPython not found — using standard Python shell)\n") 

89 

90 if use_ipython: 

91 import IPython 

92 IPython.start_ipython(argv=["--no-banner"], user_ns=namespace) 

93 else: 

94 import code 

95 code.interact(banner="", local=namespace)