asteris-llc/converge

// Copyright © 2016 Asteris, LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package logging

import (
    "bytes"
    "fmt"
    "os"
    "runtime"
    "sort"
    "strings"
    "sync"
    "time"

    "github.com/Sirupsen/logrus"
)

// Formatter is a compact output for logrus logs
type Formatter struct {
    DisableColors bool
    fancy         bool
    replacer      *strings.Replacer
    colors        map[string]func(string) string

    initter sync.Once
}

func (f *Formatter) init() {
    isColorTerminal := logrus.IsTerminal(os.Stderr) && (runtime.GOOS != "windows")
    f.fancy = (!f.DisableColors && isColorTerminal)

    f.replacer = strings.NewReplacer(
        `"`, `\"`,
        `\n`, `\\n`,
    )

    // set up colors
    reset := "\x1b[0m"
    f.colors = map[string]func(string) string{
        "blue":   func(in string) string { return "\x1b[34m" + in + reset },
        "red":    func(in string) string { return "\x1b[31m" + in + reset },
        "white":  func(in string) string { return "\x1b[37;1m" + in + reset },
        "yellow": func(in string) string { return "\x1b[33m" + in + reset },
        "dim":    func(in string) string { return "\x1b[2m" + in + reset },
    }
}

// Format an entry for printing
func (f *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
    f.initter.Do(f.init)

    var b bytes.Buffer

    // first, a timestamp
    b.WriteString(f.timestamp(entry.Time))
    b.WriteByte(' ')

    // then, our error level
    b.WriteString(f.level(entry.Level))
    b.WriteByte(' ')

    // next, the message
    b.WriteString(f.message(entry.Message))
    if f.fancy {
        b.WriteByte('\t')
    } else {
        b.WriteByte(' ')
    }

    // and our sorted keys
    keys := make([]string, 0, len(entry.Data))
    for k := range entry.Data {
        keys = append(keys, k)
    }
    sort.Strings(keys)

    last := len(keys) - 1
    for i, k := range keys {
        b.WriteString(f.kv(k, entry.Data[k]))
        if i != last {
            b.WriteByte(' ')
        }
    }

    // end with a newline
    b.WriteByte('\n')

    return b.Bytes(), nil
}

func (f *Formatter) timestamp(t time.Time) string {
    formatted := t.Format(time.RFC3339)

    if f.fancy {
        return f.colors["dim"](formatted)
    }

    return f.kv("timestamp", formatted)
}

func (f *Formatter) level(lvl logrus.Level) string {
    var (
        level    string
        colorize func(string) string
    )

    switch lvl {
    case logrus.DebugLevel:
        level = "DEBUG"
        colorize = f.colors["white"]

    case logrus.InfoLevel:
        level = "INFO"
        colorize = f.colors["blue"]

    case logrus.WarnLevel:
        level = "WARN"
        colorize = f.colors["yellow"]

    case logrus.ErrorLevel:
        level = "ERROR"
        colorize = f.colors["red"]

    case logrus.FatalLevel:
        level = "FATAL"
        colorize = f.colors["red"]

    default:
        level = "UNKNOWN"
        colorize = f.colors["red"]
    }

    if !f.fancy {
        return f.kv("level", level)
    }

    return fmt.Sprintf("|%s|", colorize(level))
}

func (f *Formatter) message(msg string) string {
    if !f.fancy {
        return f.kv("msg", msg)
    }

    return f.colors["white"](msg)
}

func (f *Formatter) kv(key string, value interface{}) string {
    if !f.fancy {
        return fmt.Sprintf("%s=\"%v\"", key, f.replacer.Replace(fmt.Sprint(value)))
    }

    var color func(string) string
    if key == "error" {
        color = f.colors["red"]
    } else {
        color = f.colors["blue"]
    }

    return fmt.Sprintf("%s=%v", color(key), value)
}