|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# Copyright 2023 Google LLC |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import dataclasses |
| 18 | +import pprint |
| 19 | +import re |
| 20 | +import reprlib |
| 21 | +import textwrap |
| 22 | + |
| 23 | + |
| 24 | +def strip_oneof(docstring): |
| 25 | + lines = docstring.splitlines() |
| 26 | + lines = [line for line in lines if ".. _oneof:" not in line] |
| 27 | + lines = [line for line in lines if "This field is a member of `oneof`_" not in line] |
| 28 | + return "\n".join(lines) |
| 29 | + |
| 30 | + |
| 31 | +def prettyprint(cls): |
| 32 | + cls.__str__ = _prettyprint |
| 33 | + cls.__repr__ = _prettyprint |
| 34 | + return cls |
| 35 | + |
| 36 | + |
| 37 | +repr = reprlib.Repr() |
| 38 | + |
| 39 | + |
| 40 | +@reprlib.recursive_repr() |
| 41 | +def _prettyprint(self): |
| 42 | + """A dataclass prettyprint function you can use in __str__or __repr__. |
| 43 | +
|
| 44 | + Note: You can't set `__str__ = pprint.pformat` because it causes a recursion error. |
| 45 | +
|
| 46 | + Mostly identical to pprint but: |
| 47 | +
|
| 48 | + * This will contract long lists and dicts (> 10lines) to [...] and {...}. |
| 49 | + * This will contract long object reprs to ClassName(...). |
| 50 | + """ |
| 51 | + fields = [] |
| 52 | + for f in dataclasses.fields(self): |
| 53 | + s = pprint.pformat(getattr(self, f.name)) |
| 54 | + class_re = r"^(\w+)\(.*\)$" |
| 55 | + if s.count("\n") >= 10: |
| 56 | + if s.startswith("["): |
| 57 | + s = "[...]" |
| 58 | + elif s.startswith("{"): |
| 59 | + s = "{...}" |
| 60 | + elif re.match(class_re, s, flags=re.DOTALL): |
| 61 | + s = re.sub(class_re, r"\1(...)", s, flags=re.DOTALL) |
| 62 | + else: |
| 63 | + s = "..." |
| 64 | + else: |
| 65 | + width = len(f.name) + 1 |
| 66 | + s = textwrap.indent(s, " " * width).lstrip(" ") |
| 67 | + fields.append(f"{f.name}={s}") |
| 68 | + attrs = ",\n".join(fields) |
| 69 | + |
| 70 | + name = self.__class__.__name__ |
| 71 | + width = len(name) + 1 |
| 72 | + |
| 73 | + attrs = textwrap.indent(attrs, " " * width).lstrip(" ") |
| 74 | + return f"{name}({attrs})" |
0 commit comments