Coverage for core/management/commands/nbshell.py: 0%

1import code 

2import platform 

3from collections import defaultdict 

4from types import SimpleNamespace 

5 

6from colorama import Fore, Style 

7from django import get_version 

8from django.apps import apps 

9from django.conf import settings 

10from django.core.management.base import BaseCommand 

11from django.utils.module_loading import import_string 

12 

13from netbox.constants import CORE_APPS 

14from netbox.plugins.utils import get_installed_plugins 

15 

16 

17def color(color: str, text: str): 

18 return getattr(Fore, color.upper()) + text + Style.RESET_ALL 

19 

20 

21def bright(text: str): 

22 return Style.BRIGHT + text + Style.RESET_ALL 

23 

24 

25def get_models(app_config): 

26 """ 

27 Return a list of all non-private models within an app. 

28 """ 

29 return [ 

30 model for model in app_config.get_models() 

31 if not getattr(model, '_netbox_private', False) 

32 ] 

33 

34 

35def get_constants(app_config): 

36 """ 

37 Return a dictionary mapping of all constants defined within an app. 

38 """ 

39 try: 

40 constants = import_string(f'{app_config.name}.constants') 

41 except ImportError: 

42 return {} 

43 return { 

44 name: value for name, value in vars(constants).items() 

45 } 

46 

47 

48class Command(BaseCommand): 

49 help = "Start the Django shell with all NetBox models already imported" 

50 django_models = {} 

51 

52 def add_arguments(self, parser): 

53 parser.add_argument( 

54 '-c', '--command', 

55 help='Python code to execute (instead of starting an interactive shell)', 

56 ) 

57 

58 def _lsapps(self): 

59 for app_label in self.django_models.keys(): 

60 app_name = apps.get_app_config(app_label).verbose_name 

61 print(f'{app_label} - {app_name}') 

62 

63 def _lsmodels(self, app_label=None): 

64 """ 

65 Return a list of all models within each app. 

66 

67 Args: 

68 app_label: The name of a specific app 

69 """ 

70 if app_label: 

71 if app_label not in self.django_models: 

72 print(f"No models listed for {app_label}") 

73 return 

74 app_labels = [app_label] 

75 else: 

76 app_labels = self.django_models.keys() # All apps 

77 

78 for app_label in app_labels: 

79 app_name = apps.get_app_config(app_label).verbose_name 

80 print(f'{app_name}:') 

81 for model in self.django_models[app_label]: 

82 print(f' {app_label}.{model}') 

83 

84 def get_namespace(self): 

85 namespace = defaultdict(SimpleNamespace) 

86 

87 # Iterate through all core apps & plugins to compile namespace of models and constants 

88 for app_name in [*CORE_APPS, *get_installed_plugins().keys()]: 

89 app_config = apps.get_app_config(app_name) 

90 

91 # Populate models 

92 if models := get_models(app_config): 

93 for model in models: 

94 setattr(namespace[app_name], model.__name__, model) 

95 self.django_models[app_name] = sorted([ 

96 model.__name__ for model in models 

97 ]) 

98 

99 # Populate constants 

100 for const_name, const_value in get_constants(app_config).items(): 

101 setattr(namespace[app_name], const_name, const_value) 

102 

103 return { 

104 **namespace, 

105 'lsapps': self._lsapps, 

106 'lsmodels': self._lsmodels, 

107 } 

108 

109 @staticmethod 

110 def get_banner_text(): 

111 lines = [ 

112 '{title} ({hostname})'.format( 

113 title=bright('NetBox interactive shell'), 

114 hostname=platform.node(), 

115 ), 

116 '{python} | {django} | {netbox}'.format( 

117 python=color('green', f'Python v{platform.python_version()}'), 

118 django=color('green', f'Django v{get_version()}'), 

119 netbox=color('green', settings.RELEASE.name), 

120 ), 

121 ] 

122 

123 if installed_plugins := get_installed_plugins(): 

124 plugin_list = ', '.join([ 

125 color('cyan', f'{name} v{version}') for name, version in installed_plugins.items() 

126 ]) 

127 lines.append( 

128 'Plugins: {plugin_list}'.format( 

129 plugin_list=plugin_list 

130 ) 

131 ) 

132 

133 lines.append( 

134 'lsapps() & lsmodels() will show available models. Use help(<model>) for more info.' 

135 ) 

136 

137 return '\n'.join([ 

138 f'### {line}' for line in lines 

139 ]) 

140 

141 def handle(self, **options): 

142 namespace = self.get_namespace() 

143 

144 # If Python code has been passed, execute it and exit. 

145 if options['command']: 

146 exec(options['command'], namespace) 

147 return None 

148 

149 # Try to enable tab-complete 

150 try: 

151 import readline 

152 import rlcompleter 

153 except ModuleNotFoundError: # pragma: no cover 

154 pass 

155 else: 

156 readline.set_completer(rlcompleter.Completer(namespace).complete) 

157 readline.parse_and_bind('tab: complete') 

158 

159 # Run interactive shell 

160 return code.interact(banner=self.get_banner_text(), local=namespace)