From 591a5fcc010b5aca7fa15de9aa5f40ed23b86eb9 Mon Sep 17 00:00:00 2001 From: Kevin Zheng Date: Mon, 26 Jul 2021 05:04:30 -0400 Subject: [PATCH] Add builtin support for truncation --- api.go | 8 ++ interfaces/interfaces.go | 6 ++ internal/buffer/buffer.go | 135 ++++++++++++++++++++++++++++++++-- internal/markers/constants.go | 1 + internal/rfmt/print.go | 14 ++++ markers_test.go | 40 ++++++++++ 6 files changed, 198 insertions(+), 6 deletions(-) diff --git a/api.go b/api.go index 2c8eba6..55a2211 100644 --- a/api.go +++ b/api.go @@ -57,6 +57,9 @@ type SafeValue = i.SafeValue // TODO(knz): Remove this. type SafeMessager = i.SafeMessager +// TruncatedValue represents a value that is truncated when printed. +type TruncatedValue = i.TruncatedValue + // SafePrinter is a stateful helper that abstracts an output stream in // the context of printf-like formatting, but with the ability to // separate safe and unsafe bits of data. @@ -141,6 +144,11 @@ func Unsafe(a interface{}) interface{} { return w.Unsafe(a) } // The implementation is also slow. func Safe(a interface{}) SafeValue { return w.Safe(a) } +// Truncate limits the rune length of the outputted value when +// passed into the formatter. The excess runes are replaced +// with the marker `…`. +func Truncate(a interface{}, len int) TruncatedValue { return TruncatedValue{Value: a, Length: len} } + // RegisterRedactErrorFn registers an error redaction function for use // during automatic redaction by this package. // Provided e.g. by cockroachdb/errors. diff --git a/interfaces/interfaces.go b/interfaces/interfaces.go index 20fec2e..90ff624 100644 --- a/interfaces/interfaces.go +++ b/interfaces/interfaces.go @@ -147,3 +147,9 @@ type SafeValue interface { type SafeMessager = interface { SafeMessage() string } + +// TruncatedValue represents a value that is truncated when printed. +type TruncatedValue struct { + Value interface{} + Length int +} diff --git a/internal/buffer/buffer.go b/internal/buffer/buffer.go index 74fd9b8..d7e37a6 100644 --- a/internal/buffer/buffer.go +++ b/internal/buffer/buffer.go @@ -31,6 +31,12 @@ type Buffer struct { validUntil int // exclusive upper bound of data that's already validated mode OutputMode markerOpen bool + + // These fields relate to truncation and are used to enforce the limit of + // runes in the buffer. Markers do not contribute to the rune count. + runeCount int + runeLimit int + excess bool } // OutputMode determines how writes are processed in the Buffer. @@ -128,28 +134,96 @@ func (b *Buffer) Cap() int { // needed. The return value n is the length of p; err is always nil. If the // buffer becomes too large, Write will panic with ErrTooLarge. func (b *Buffer) Write(p []byte) (n int, err error) { + length := len(p) + if b.runeLimit != 0 { + if b.runeLimit == b.runeCount { + if len(p) != 0 { + b.excess = true + } + return + } + + runes, end, excess := limitRunes(p, b.runeLimit-b.runeCount) + b.runeCount += runes + if end < length { + length = end + } + b.excess = excess + } + b.startWrite() - m, ok := b.tryGrowByReslice(len(p)) + m, ok := b.tryGrowByReslice(length) if !ok { - m = b.grow(len(p)) + m = b.grow(length) } - return copy(b.buf[m:], p), nil + return copy(b.buf[m:], p[:length]), nil } // WriteString appends the contents of s to the buffer, growing the buffer as // needed. The return value n is the length of s; err is always nil. If the // buffer becomes too large, WriteString will panic with ErrTooLarge. func (b *Buffer) WriteString(s string) (n int, err error) { + length := len(s) + if b.runeLimit != 0 { + if b.runeLimit == b.runeCount { + if len(s) != 0 { + b.excess = true + } + return + } + + runes, end, excess := limitRunesInString(s, b.runeLimit-b.runeCount) + b.runeCount += runes + if end < length { + length = end + } + b.excess = excess + } + b.startWrite() - m, ok := b.tryGrowByReslice(len(s)) + m, ok := b.tryGrowByReslice(length) if !ok { - m = b.grow(len(s)) + m = b.grow(length) } - return copy(b.buf[m:], s), nil + return copy(b.buf[m:], s[:length]), nil +} + +// limitRunes counts the number of bytes in the provided +// byte slice up to a given rune limit. +func limitRunes(p []byte, limit int) (runes int, end int, excess bool) { + for runes < limit && len(p) > 0 { + _, s := utf8.DecodeRune(p) + runes++ + end += s + p = p[s:] + } + excess = len(p) != 0 + return +} + +// limitRunesInString counts the number of bytes in the provided +// string up to a given rune limit. +func limitRunesInString(str string, limit int) (runes int, end int, excess bool) { + for runes < limit && len(str) > 0 { + _, s := utf8.DecodeRuneInString(str) + runes++ + end += s + str = str[s:] + } + excess = len(str) != 0 + return } // WriteByte emits a single byte. func (b *Buffer) WriteByte(s byte) error { + if b.runeLimit != 0 { + if b.runeLimit == b.runeCount { + b.excess = true + return nil + } + b.runeCount++ + } + b.startWrite() if b.mode == UnsafeEscaped && (s >= utf8.RuneSelf || @@ -168,6 +242,14 @@ func (b *Buffer) WriteByte(s byte) error { // WriteRune emits a single rune. func (b *Buffer) WriteRune(s rune) error { + if b.runeLimit != 0 { + if b.runeLimit == b.runeCount { + b.excess = true + return nil + } + b.runeCount++ + } + b.startWrite() l := utf8.RuneLen(s) m, ok := b.tryGrowByReslice(l) @@ -178,6 +260,47 @@ func (b *Buffer) WriteRune(s rune) error { return nil } +// TruncateState is a struct used to store the previously set +// rune limit. This is used to keep track of rune limits during +// nested truncation operations. +type TruncateState struct { + buf *Buffer + limit int +} + +// Truncate sets the rune limit and returns a TruncateState +// that can be used to restore the previous rune limit. The +// limit cannot exceed the previous rune limit. +func (b *Buffer) Truncate(limit int) TruncateState { + prev := TruncateState{b, b.runeLimit} + if b.runeLimit == 0 || b.runeCount+limit < b.runeLimit { + b.runeLimit = b.runeCount + limit + } + return prev +} + +// Restore restores the previous rune limit. If the current rune limit +// is exceeded, a truncation marker is written. +func (s TruncateState) Restore() { + if s.buf.excess && (s.limit == 0 || s.limit != s.buf.runeLimit) { + s.buf.excess = false + s.buf.markTruncate() + } + s.buf.runeLimit = s.limit +} + +// markTruncate adds the truncation marker. +func (b *Buffer) markTruncate() { + if len(b.buf) == 0 { + return + } + p, ok := b.tryGrowByReslice(len(m.Truncate)) + if !ok { + p = b.grow(len(m.Truncate)) + } + copy(b.buf[p:], m.Truncate) +} + // finalize ensures that all the buffer is properly // marked. func (b *Buffer) finalize() { diff --git a/internal/markers/constants.go b/internal/markers/constants.go index 51f488f..f3ee5cb 100644 --- a/internal/markers/constants.go +++ b/internal/markers/constants.go @@ -25,6 +25,7 @@ const ( EscapeMark = '?' EscapeMarkS = string(EscapeMark) RedactedS = StartS + "×" + EndS + Truncate = `…` ) // Internal variables. diff --git a/internal/rfmt/print.go b/internal/rfmt/print.go index 16ba3ce..2820f21 100644 --- a/internal/rfmt/print.go +++ b/internal/rfmt/print.go @@ -647,6 +647,13 @@ func (p *pp) handleMethods(verb rune) (handled bool) { } } + if truncated, ok := p.arg.(i.TruncatedValue); ok { + handled = true + defer p.catchPanic(p.arg, verb, "TruncatedValue") + p.printTruncated(truncated, verb) + return + } + // Is it a Formatter? if formatter, ok := p.arg.(Formatter); ok { handled = true @@ -695,6 +702,13 @@ func (p *pp) handleMethods(verb rune) (handled bool) { return false } +func (p *pp) printTruncated(truncated i.TruncatedValue, verb rune) { + if truncated.Length > 0 { + defer p.buf.Truncate(truncated.Length).Restore() + } + p.printArg(truncated.Value, verb) +} + func (p *pp) printArg(arg interface{}, verb rune) { t := reflect.TypeOf(arg) if safeTypeRegistry[t] { diff --git a/markers_test.go b/markers_test.go index 18f216e..b510bc5 100644 --- a/markers_test.go +++ b/markers_test.go @@ -204,6 +204,46 @@ func TestPrinter(t *testing.T) { {func(w p) { w.Print(buf) }, "safe ‹unsafe›"}, {func(w p) { w.Printf("%v", &buf) }, "safe ‹unsafe›"}, {func(w p) { w.Print(&buf) }, "safe ‹unsafe›"}, + + // Truncation. + {func(w p) { w.Print(Truncate(Safe("säfe"), 3)) }, "säf…"}, + {func(w p) { w.Print(Truncate(Safe("safe"), 5)) }, "safe"}, + {func(w p) { w.Print(Truncate(Safe("safe"), 100)) }, "safe"}, + {func(w p) { w.Print(Truncate("unsäfe", 3)) }, "‹uns›…"}, + {func(w p) { w.Print(Truncate("unsafe", 100)) }, "‹unsafe›"}, + {func(w p) { + w.Printf("Unsafe: %v Safe: %v", Truncate("unsafe", 3), Truncate(Safe("safe"), 2)) + }, "Unsafe: ‹uns›… Safe: sa…"}, + // Nested Truncation. + {func(w p) { w.Print(Truncate(Truncate(Safe("safe"), 5), 3)) }, "saf…"}, + {func(w p) { w.Print(Truncate(Truncate(Safe("safe"), 3), 5)) }, "saf…"}, + {func(w p) { + w.Print(Truncate(compose{func(p2 p) { + p2.Printf( + "Inside %v %v", + Truncate(Safe("safe"), 3), + Truncate("unsafe", 3), + ) + }}, 10)) + }, "Inside saf…"}, + {func(w p) { + w.Print(Truncate(compose{func(p2 p) { + p2.Printf( + "Inside %v %v", + Truncate(Safe("safe"), 3), + Truncate("unsafe", 3), + ) + }}, 20)) + }, "Inside saf… ‹uns›…"}, + {func(w p) { + w.Print(Truncate(compose{func(p2 p) { + p2.Printf( + "Inside %v %v", + Truncate(Safe("safe"), 3), + Truncate("unsafe", 20), + ) + }}, 20)) + }, "Inside saf… ‹unsafe›"}, } var methods = []struct {