Coverage for shellcraft / mixins.py: 93%

1import warnings 

2 

3_INJECTED = False 

4 

5 

6def inject(): 

7 """Inject Rails-style class/instance methods onto Django's base Model. 

8 

9 Safe to call multiple times — idempotent after first call. 

10 Only call this from the shellcraft management command. 

11 

12 Raises RuntimeError when DEBUG is False so that an accidental import of 

13 this module in a production Celery worker or WSGI request cannot 

14 activate the shell shortcuts on every model in the process. 

15 """ 

16 global _INJECTED 

17 if _INJECTED: 

18 return 

19 

20 # Belt-and-suspenders production guard. The management command also 

21 # checks DEBUG before calling inject(), but inject() is a public function 

22 # and could be called from user code (e.g. a custom management command 

23 # that wraps shellcraft). Checking here ensures the guard can never be 

24 # bypassed simply by calling inject() directly. 

25 from django.conf import settings 

26 if not settings.DEBUG: 

27 raise RuntimeError( 

28 "shellcraft.inject() must not be called when DEBUG=False. " 

29 "The shellcraft shell is a development-only tool." 

30 ) 

31 

32 _INJECTED = True 

33 

34 from django.db import models 

35 

36 from shellcraft.printer import ap, format_fields 

37 

38 def _find(cls, pk): 

39 return cls.objects.get(pk=pk) 

40 

41 def _all(cls): 

42 # Return the queryset so callers can assign: `users = User.all()`. 

43 # Printing and returning are both useful; returning None was a silent 

44 # footgun for anyone who stored the result. 

45 qs = cls.objects.all() 

46 ap(qs) 

47 return qs 

48 

49 def _where(cls, **kwargs): 

50 qs = cls.objects.filter(**kwargs) 

51 ap(qs) 

52 return qs 

53 

54 def _first(cls): 

55 obj = cls.objects.first() 

56 if obj is not None: 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 ap(obj) 

58 return obj 

59 

60 def _last(cls): 

61 obj = cls.objects.last() 

62 if obj is not None: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 ap(obj) 

64 return obj 

65 

66 def _count(cls): 

67 return cls.objects.count() 

68 

69 def _fields(cls): 

70 print(format_fields(cls)) 

71 

72 def _update(self, **kwargs): 

73 # Block dunder and private-attribute keys outright. Allowing 

74 # user.update(__class__=X) or user.update(__dict__={...}) would let 

75 # shell users overwrite Python internals, bypassing Django's field 

76 # validation entirely and corrupting the object in memory. 

77 for key in kwargs: 

78 if key.startswith("_"): 

79 raise ValueError( 

80 f"shellcraft: cannot update private or dunder attribute '{key}'. " 

81 "Pass only model field names." 

82 ) 

83 

84 # Warn (don't raise) for keys that don't correspond to any known field 

85 # on this model. Django's save() will silently ignore them, which is 

86 # confusing during interactive development. 

87 valid_fields = {f.attname for f in self.__class__._meta.concrete_fields} 

88 valid_fields.update(f.name for f in self.__class__._meta.concrete_fields) 

89 valid_fields.update(f.name for f in self.__class__._meta.many_to_many) 

90 

91 for key in kwargs: 

92 if key not in valid_fields: 

93 warnings.warn( 

94 f"shellcraft: '{key}' is not a recognised field on " 

95 f"{self.__class__.__name__} — the attribute will be set " 

96 "but will not be persisted by Django's save().", 

97 stacklevel=2, 

98 ) 

99 

100 for key, value in kwargs.items(): 

101 setattr(self, key, value) 

102 self.save() 

103 return self 

104 

105 def _destroy(self): 

106 return self.delete() 

107 

108 class_methods = { 

109 'find': _find, 

110 'all': _all, 

111 'where': _where, 

112 'first': _first, 

113 'last': _last, 

114 'count': _count, 

115 'fields': _fields, 

116 } 

117 instance_methods = { 

118 'update': _update, 

119 'destroy': _destroy, 

120 } 

121 

122 for name, fn in class_methods.items(): 

123 _safe_inject(models.Model, name, classmethod(fn)) 

124 

125 for name, fn in instance_methods.items(): 

126 _safe_inject(models.Model, name, fn) 

127 

128 

129def _safe_inject(cls, name, method): 

130 if hasattr(cls, name): 

131 existing = getattr(cls, name) 

132 warnings.warn( 

133 f"shellcraft: '{name}' already exists on {cls.__name__} " 

134 f"(defined in {getattr(existing, '__module__', '?')}) — skipping injection", 

135 stacklevel=3, 

136 ) 

137 return 

138 setattr(cls, name, method)