Coverage for shellcraft / printer.py: 95%
107 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-12 09:09 +0530
1import datetime
2import re
3from decimal import Decimal
5# ANSI color codes
6BLUE = '\033[34m'
7YELLOW = '\033[33m'
8GREEN = '\033[32m'
9RED = '\033[31m'
10CYAN = '\033[36m'
11GRAY = '\033[90m'
12RESET = '\033[0m'
14_FIELD_TYPE_COLORS = {
15 'AutoField': CYAN, 'BigAutoField': CYAN, 'SmallAutoField': CYAN,
16 'IntegerField': BLUE, 'BigIntegerField': BLUE, 'SmallIntegerField': BLUE,
17 'FloatField': BLUE, 'DecimalField': BLUE,
18 'CharField': YELLOW, 'TextField': YELLOW, 'SlugField': YELLOW,
19 'EmailField': YELLOW, 'URLField': YELLOW,
20 'BooleanField': GREEN, 'NullBooleanField': GREEN,
21 'DateField': CYAN, 'DateTimeField': CYAN, 'TimeField': CYAN,
22 'ForeignKey': GRAY, 'OneToOneField': GRAY, 'ManyToManyField': GRAY,
23}
25# Matches any ANSI CSI escape sequence — covers colors, cursor movement,
26# clear-screen codes, etc. Used to sanitize user-supplied string values
27# before they reach the terminal to prevent terminal injection.
28_ANSI_ESCAPE_RE = re.compile(r'\x1b\[[0-9;]*[A-Za-z]')
31def _strip_ansi(s: str) -> str:
32 """Remove ANSI escape sequences embedded in a user-supplied string value.
34 A database value such as ``"\\033[2J\\033[0;0H"`` would otherwise clear the
35 terminal when printed. We strip rather than escape because the point is
36 display-only output; the raw data is unmodified in the database.
37 """
38 return _ANSI_ESCAPE_RE.sub('', s)
41def colorize(value):
42 """Return ANSI-colored string representation of a Python value."""
43 # bool must be checked before int because bool is a subclass of int.
44 if isinstance(value, bool):
45 return f"{GREEN}true{RESET}" if value else f"{RED}false{RESET}"
46 if isinstance(value, (int, float, Decimal)):
47 return f"{BLUE}{value}{RESET}"
48 if isinstance(value, str):
49 # Sanitize before wrapping in color codes so embedded escape sequences
50 # cannot bleed into the terminal (terminal injection prevention).
51 return f'{YELLOW}"{_strip_ansi(value)}"{RESET}'
52 if value is None:
53 return f"{RED}nil{RESET}"
54 # datetime.datetime is a subclass of datetime.date, so both are matched
55 # here. datetime.time is a separate type and also colored cyan to match
56 # the field-type color for TimeField.
57 if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
58 return f"{CYAN}{value}{RESET}"
59 return str(value)
62def _get_fields(obj):
63 """Return list of (attname, value) for concrete fields of a model instance."""
64 return [
65 (field.attname, getattr(obj, field.attname, None))
66 for field in obj.__class__._meta.concrete_fields
67 ]
70def _format_record(obj):
71 """Format a single model instance as an awesome_print-style string."""
72 cls = obj.__class__
73 header = f"#<{cls.__name__}:{hex(id(obj))}>"
74 fields = _get_fields(obj)
76 max_key_len = max((len(name) for name, _ in fields), default=0)
78 lines = [f"{header} {{"]
79 for name, value in fields:
80 key = f":{name}".ljust(max_key_len + 1)
81 lines.append(f" {key} => {colorize(value)},")
82 lines.append("}")
83 return '\n'.join(lines)
86def _is_queryset_like(obj) -> bool:
87 """Return True if *obj* is a Django QuerySet or a compatible iterable.
89 Design:
90 - We prefer a concrete ``isinstance`` check when Django is importable so
91 that objects such as ``ModelForm`` (which has both ``.__iter__`` and
92 ``.model``) are not mistakenly treated as querysets.
93 - A duck-typing fallback is kept so that lightweight test mocks (which
94 carry ``__iter__`` and ``.model`` but are not real QuerySet instances)
95 continue to work without requiring a live Django setup in every test.
96 - The ``not hasattr(obj, '_meta')`` guard prevents a single model
97 instance from being routed into the list path if it ever exposes
98 ``.model`` (uncommon, but possible with custom model mixins).
99 """
100 try:
101 from django.db.models.query import QuerySet
102 if isinstance(obj, QuerySet):
103 return True
104 except ImportError:
105 pass
107 # Duck-typing fallback for test mocks and custom queryset-like objects.
108 return (
109 hasattr(obj, '__iter__')
110 and hasattr(obj, 'model')
111 and not hasattr(obj, '_meta')
112 )
115def ap(obj):
116 """Pretty-print a model instance or queryset in awesome_print style."""
117 if obj is None:
118 print(colorize(None))
119 return
121 if _is_queryset_like(obj):
122 _print_queryset(list(obj))
123 return
125 # Single model instance
126 if hasattr(obj, '_meta'):
127 print(_format_record(obj))
128 return
130 print(repr(obj))
133def _print_queryset(records):
134 if not records:
135 print("[]")
136 return
138 if len(records) == 1:
139 print(_format_record(records[0]))
140 return
142 # Precompute the width of the widest index label (e.g. "[99]" = 4 chars)
143 # so that all headers and field lines stay aligned for any result count.
144 max_idx_width = len(f"[{len(records) - 1}]")
145 # Prefix for the record header line: " [N…] "
146 content_start = 4 + max_idx_width + 1
147 # Field lines are indented 4 more chars than the header start (inside {}).
148 field_indent = " " * (content_start + 4)
149 close_indent = " " * content_start
151 print("[")
152 for i, record in enumerate(records):
153 is_last = i == len(records) - 1
154 fields = _get_fields(record)
155 cls = record.__class__
156 header = f"#<{cls.__name__}:{hex(id(record))}>"
157 max_key_len = max((len(name) for name, _ in fields), default=0)
159 # Pad the plain label before applying color so ljust counts visible
160 # characters only (ANSI codes are zero-width).
161 idx_label = f"[{i}]".ljust(max_idx_width)
162 idx_colored = f"{GRAY}{idx_label}{RESET}"
164 print(f" {idx_colored} {header} {{")
165 for name, value in fields:
166 key = f":{name}".ljust(max_key_len + 1)
167 print(f"{field_indent}{key} => {colorize(value)},")
168 comma = "" if is_last else ","
169 print(f"{close_indent}}}{comma}")
170 print("]")
171 count = len(records)
172 print(f"{count} {'record' if count == 1 else 'records'}")
175def format_fields(model_class):
176 """Return a colored schema summary for a model class."""
177 concrete = list(model_class._meta.concrete_fields)
178 m2m = list(model_class._meta.many_to_many)
180 all_names = [f.attname for f in concrete] + [f.name for f in m2m]
181 max_name_len = max((len(n) for n in all_names), default=0)
183 lines = [f"{CYAN}{model_class.__name__}{RESET} {{"]
185 for field in concrete:
186 type_name = field.__class__.__name__
187 color = _FIELD_TYPE_COLORS.get(type_name, RESET)
188 name = f":{field.attname}".ljust(max_name_len + 1)
189 nullable = getattr(field, 'null', False)
190 null_tag = f" {GRAY}null: true{RESET}" if nullable else ""
191 lines.append(f" {name} {color}{type_name}{RESET}{null_tag}")
193 for field in m2m: 193 ↛ 194line 193 didn't jump to line 194 because the loop on line 193 never started
194 type_name = field.__class__.__name__
195 color = _FIELD_TYPE_COLORS.get(type_name, GRAY)
196 name = f":{field.name}".ljust(max_name_len + 1)
197 lines.append(f" {name} {color}{type_name}{RESET}")
199 lines.append("}")
200 return '\n'.join(lines)