diff --git a/decor/on_abort.go b/decor/on_abort.go index f71928f..50a1dfb 100644 --- a/decor/on_abort.go +++ b/decor/on_abort.go @@ -3,6 +3,8 @@ var ( _ Decorator = onAbortWrapper{} _ Wrapper = onAbortWrapper{} + _ Decorator = onAbortMetaWrapper{} + _ Wrapper = onAbortMetaWrapper{} ) // OnAbort wrap decorator. @@ -33,3 +35,34 @@ func (d onAbortWrapper) Unwrap() Decorator { return d.Decorator } + +// OnAbortMeta wrap decorator. +// Provided fn is supposed to wrap output of given decorator +// with meta information like ANSI escape codes for example. +// Primary usage intention is to set SGR display attributes. +// +// `decorator` Decorator to wrap +// `fn` func to apply meta information +func OnAbortMeta(decorator Decorator, fn func(string) string) Decorator { + if decorator == nil { + return nil + } + return onAbortMetaWrapper{decorator, fn} +} + +type onAbortMetaWrapper struct { + Decorator + fn func(string) string +} + +func (d onAbortMetaWrapper) Decor(s Statistics) (string, int) { + if s.Completed { + str, width := d.Decorator.Decor(s) + return d.fn(str), width + } + return d.Decorator.Decor(s) +} + +func (d onAbortMetaWrapper) Unwrap() Decorator { + return d.Decorator +}