Coverage for django_query_capture/utils.py: 95%

1import io 

2import sys 

3from contextlib import ContextDecorator 

4 

5from django.utils import termcolors 

6 

7from django_query_capture.capture import CapturedQuery 

8from django_query_capture.settings import get_config 

9 

10 

11class CaptureStdOutToString(ContextDecorator): 

12 def __enter__(self) -> sys.stdout: 

13 self.old_stdout = sys.stdout 

14 new_stdout = io.StringIO() 

15 sys.stdout = new_stdout 

16 return sys.stdout 

17 

18 def __exit__(self, exc_type, exc_val, exc_tb): 

19 sys.stdout = self.old_stdout 

20 

21 

22def colorize(value: str, is_warning: bool) -> str: 

23 """ 

24 Utility to set a color for the output string when it exceeds the threshold. 

25 

26 Args: 

27 value: String to be output. 

28 is_warning: Whether it exceeds the threshold. 

29 

30 Returns: 

31 colorized string output 

32 """ 

33 if is_warning: 

34 return termcolors.make_style(fg=get_config()["PRINT_THRESHOLDS"]["COLOR"])( # type: ignore 

35 value 

36 ) 

37 return value 

38 

39 

40def get_stack_prefix(captured_query: CapturedQuery) -> str: 

41 """ 

42 Utilities that help you output call stacks consistently in [CapturedQuery][capture.CapturedQuery]. 

43 Args: 

44 captured_query: [CapturedQuery][capture.CapturedQuery] 

45 """ 

46 return f'[{captured_query["function_name"]}, {captured_query["file_name"]}:{captured_query["line_no"]}]' 

47 

48 

49def truncate_string(value: str, length: int) -> str: 

50 """ 

51 

52 Args: 

53 value: String to be output. 

54 length: Number of strings to output. 

55 

56 Returns: 

57 truncated string 

58 """ 

59 return (value[:length] + "..") if len(value) > length else value