writer.go 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package hclog
  2. import (
  3. "bytes"
  4. "io"
  5. )
  6. type writer struct {
  7. b bytes.Buffer
  8. w io.Writer
  9. }
  10. func newWriter(w io.Writer) *writer {
  11. return &writer{w: w}
  12. }
  13. func (w *writer) Flush(level Level) (err error) {
  14. if lw, ok := w.w.(LevelWriter); ok {
  15. _, err = lw.LevelWrite(level, w.b.Bytes())
  16. } else {
  17. _, err = w.w.Write(w.b.Bytes())
  18. }
  19. w.b.Reset()
  20. return err
  21. }
  22. func (w *writer) Write(p []byte) (int, error) {
  23. return w.b.Write(p)
  24. }
  25. func (w *writer) WriteByte(c byte) error {
  26. return w.b.WriteByte(c)
  27. }
  28. func (w *writer) WriteString(s string) (int, error) {
  29. return w.b.WriteString(s)
  30. }
  31. // LevelWriter is the interface that wraps the LevelWrite method.
  32. type LevelWriter interface {
  33. LevelWrite(level Level, p []byte) (n int, err error)
  34. }
  35. // LeveledWriter writes all log messages to the standard writer,
  36. // except for log levels that are defined in the overrides map.
  37. type LeveledWriter struct {
  38. standard io.Writer
  39. overrides map[Level]io.Writer
  40. }
  41. // NewLeveledWriter returns an initialized LeveledWriter.
  42. //
  43. // standard will be used as the default writer for all log levels,
  44. // except for log levels that are defined in the overrides map.
  45. func NewLeveledWriter(standard io.Writer, overrides map[Level]io.Writer) *LeveledWriter {
  46. return &LeveledWriter{
  47. standard: standard,
  48. overrides: overrides,
  49. }
  50. }
  51. // Write implements io.Writer.
  52. func (lw *LeveledWriter) Write(p []byte) (int, error) {
  53. return lw.standard.Write(p)
  54. }
  55. // LevelWrite implements LevelWriter.
  56. func (lw *LeveledWriter) LevelWrite(level Level, p []byte) (int, error) {
  57. w, ok := lw.overrides[level]
  58. if !ok {
  59. w = lw.standard
  60. }
  61. return w.Write(p)
  62. }