Pretty prints the given object which is of the Program type and any of its attribute’s types.
(obj: Any, indent: int = 0, out: Optional[TextIO] = None)
| 248 | |
| 249 | # pyre-ignore |
| 250 | def pretty_print(obj: Any, indent: int = 0, out: Optional[TextIO] = None) -> None: |
| 251 | """ |
| 252 | Pretty prints the given object which is of the Program type and any of its |
| 253 | attribute’s types. |
| 254 | """ |
| 255 | if isinstance(obj, torch.fx.GraphModule): |
| 256 | raise ExportError( |
| 257 | ExportErrorType.INVALID_INPUT_TYPE, |
| 258 | "pretty_print() does not accept GraphModule as input.", |
| 259 | ) |
| 260 | |
| 261 | # Instruction types are IntEnum object |
| 262 | if isinstance(obj, IntEnum): |
| 263 | print(int(obj), end="", file=out) |
| 264 | return |
| 265 | |
| 266 | primitives = (int, str, bool, float, type(None)) |
| 267 | if isinstance(obj, primitives): |
| 268 | print(obj, end="", file=out) |
| 269 | return |
| 270 | |
| 271 | if isinstance(obj, bytes): |
| 272 | r = reprlib.Repr() |
| 273 | r.maxother = 1024 |
| 274 | print(r.repr(obj), end="", file=out) |
| 275 | return |
| 276 | |
| 277 | if isinstance(obj, list): |
| 278 | if len(obj) < 10 and all(isinstance(elem, int) for elem in obj): |
| 279 | print(obj, end="", file=out) |
| 280 | return |
| 281 | print("[", file=out) |
| 282 | for index, elem in enumerate(obj): |
| 283 | print(" " * (indent + 1), end="", file=out) |
| 284 | pretty_print(elem, indent + 1, out=out) |
| 285 | print(f"(index={index}),", file=out) |
| 286 | print(" " * indent + "]", end="", file=out) |
| 287 | return |
| 288 | |
| 289 | inline = all( |
| 290 | isinstance(getattr(obj, field.name), primitives) for field in fields(obj) |
| 291 | ) |
| 292 | end = "" if inline else "\n" |
| 293 | print(f"{type(obj).__name__}(", end=end, file=out) |
| 294 | for i, _field in enumerate(fields(obj)): |
| 295 | if not inline: |
| 296 | print(" " * (indent + 1), end="", file=out) |
| 297 | print(_field.name + "=", end="", file=out) |
| 298 | pretty_print(getattr(obj, _field.name), indent + 1, out=out) |
| 299 | if i < len(fields(obj)) - 1: |
| 300 | print(", ", end="", file=out) |
| 301 | print("", end=end, file=out) |
| 302 | if not inline: |
| 303 | print(" " * indent, end="", file=out) |
| 304 | print(")", end="" if indent else "\n", file=out) |
| 305 | |
| 306 | |
| 307 | def pretty_print_stacktraces(obj: FrameList) -> str: |