Add the onPanic and onFatal handlers

logger.go

@@ -0,0 +1,126 @@
+package tracelog
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"os"
+)
+
+type LoggerType string
+
+const (
+	InfoLoggerType    LoggerType = "INFO"
+	WarningLoggerType LoggerType = "WARNING"
+	ErrorLoggerType   LoggerType = "ERROR"
+	DebugLoggerType   LoggerType = "DEBUG"
+)
+
+// BasicLogger represents the logger
+// without the panic and fatal methods
+type BasicLogger interface {
+	Output(calldepth int, s string) error
+	Writer() io.Writer
+	SetOutput(w io.Writer)
+	Flags() int
+	SetFlags(flag int)
+	Prefix() string
+	SetPrefix(prefix string)
+	Print(v ...interface{})
+	Printf(format string, v ...interface{})
+	Println(v ...interface{})
+}
+
+type logger struct {
+	BasicLogger
+	onFatal func(format string, v ...interface{})
+	onPanic func(format string, v ...interface{})
+}
+
+func NewLogger(out io.Writer, loggerType LoggerType) *logger {
+	logPrefix := fmt.Sprintf("%s: ", loggerType)
+	internalLogger := log.New(out, logPrefix, timeFlags)
+	onFatal := func(format string, v ...interface{}) {
+		os.Exit(1)
+	}
+	onPanic := func(format string, v ...interface{}) {
+		panic(fmt.Sprintf(format, v...))
+	}
+	return &logger{BasicLogger: internalLogger, onFatal: onFatal, onPanic: onPanic}
+}
+
+func (logger *logger) PanicError(err error) {
+	logger.Printf(GetErrorFormatter(), err)
+	logger.onPanic(GetErrorFormatter(), err)
+}
+
+func (logger *logger) PanicfOnError(format string, err error) {
+	if err != nil {
+		logger.Printf(format, err)
+		logger.onPanic(format, err)
+	}
+}
+
+func (logger *logger) PanicOnError(err error) {
+	if err != nil {
+		logger.PanicError(err)
+	}
+}
+
+func (logger *logger) FatalError(err error) {
+	logger.Printf(GetErrorFormatter(), err)
+	logger.onFatal(GetErrorFormatter(), err)
+}
+
+func (logger *logger) FatalfOnError(format string, err error) {
+	if err != nil {
+		logger.Printf(format, err)
+		logger.onFatal(format, err)
+	}
+}
+
+func (logger *logger) FatalOnError(err error) {
+	if err != nil {
+		logger.FatalError(err)
+	}
+}
+
+func (logger *logger) PrintError(err error) {
+	logger.Printf(GetErrorFormatter()+"\n", err)
+}
+
+func (logger *logger) PrintOnError(err error) {
+	if err != nil {
+		logger.PrintError(err)
+	}
+}
+
+func (logger *logger) Fatal(v ...interface{}) {
+	logger.Print(v...)
+	logger.onFatal(GetErrorFormatter(), v)
+}
+
+func (logger *logger) Fatalf(format string, v ...interface{}) {
+	logger.Printf(format, v...)
+	logger.onFatal(format, v)
+}
+
+func (logger *logger) Fatalln(v ...interface{}) {
+	logger.Println(v...)
+	logger.onFatal(GetErrorFormatter(), v)
+}
+
+func (logger *logger) Panic(v ...interface{}) {
+	logger.Print(v...)
+	logger.onPanic(GetErrorFormatter(), v)
+}
+
+func (logger *logger) Panicf(format string, v ...interface{}) {
+	logger.Printf(format, v...)
+	logger.onPanic(format, v)
+}
+
+func (logger *logger) Panicln(v ...interface{}) {
+	logger.Println(v...)
+	logger.onPanic(GetErrorFormatter(), v)
+}

logging.go

@@ -9,35 +9,38 @@ import (
 	"github.com/pkg/errors"
 )
 
+type LogLevel string
+
 const (
-	NormalLogLevel = "NORMAL"
-	DevelLogLevel  = "DEVEL"
-	ErrorLogLevel  = "ERROR"
-	timeFlags      = log.LstdFlags | log.Lmicroseconds
+	NormalLogLevel LogLevel = "NORMAL"
+	DevelLogLevel  LogLevel = "DEVEL"
+	ErrorLogLevel  LogLevel = "ERROR"
+	timeFlags               = log.LstdFlags | log.Lmicroseconds
 )
 
-var InfoLogger = NewErrorLogger(os.Stdout, "INFO: ")
-var WarningLogger = NewErrorLogger(os.Stdout, "WARNING: ")
-var ErrorLogger = NewErrorLogger(os.Stderr, "ERROR: ")
-var DebugLogger = NewErrorLogger(ioutil.Discard, "DEBUG: ")
+var InfoLogger = NewLogger(os.Stderr, InfoLoggerType)
+var WarningLogger = NewLogger(os.Stderr, WarningLoggerType)
+var ErrorLogger = NewLogger(os.Stderr, ErrorLoggerType)
+var DebugLogger = NewLogger(ioutil.Discard, DebugLoggerType)
 
-var LogLevels = []string{NormalLogLevel, DevelLogLevel, ErrorLogLevel}
+var LogLevels = []LogLevel{NormalLogLevel, DevelLogLevel, ErrorLogLevel}
 var logLevel = NormalLogLevel
-var logLevelFormatters = map[string]string{
+var logLevelFormatters = map[LogLevel]string{
 	NormalLogLevel: "%v",
 	ErrorLogLevel:  "%v",
 	DevelLogLevel:  "%+v",
 }
 
-func setupLoggers() {
-	if logLevel == NormalLogLevel {
-		DebugLogger = NewErrorLogger(ioutil.Discard, "DEBUG: ")
-	} else if logLevel == ErrorLogLevel {
-		DebugLogger = NewErrorLogger(ioutil.Discard, "DEBUG: ")
-		InfoLogger = NewErrorLogger(ioutil.Discard, "INFO: ")
-		WarningLogger = NewErrorLogger(ioutil.Discard, "WARNING: ")
-	} else {
-		DebugLogger = NewErrorLogger(os.Stdout, "DEBUG: ")
+func syncLogLevel() {
+	switch logLevel {
+	case NormalLogLevel:
+		DebugLogger.SetOutput(ioutil.Discard)
+	case ErrorLogLevel:
+		DebugLogger.SetOutput(ioutil.Discard)
+		InfoLogger.SetOutput(ioutil.Discard)
+		WarningLogger.SetOutput(ioutil.Discard)
+	default: // assume DevelLogLevel
+		DebugLogger.SetOutput(os.Stderr)
 	}
 }
 
@@ -57,7 +60,7 @@ func GetErrorFormatter() string {
 	return logLevelFormatters[logLevel]
 }
 
-func UpdateLogLevel(newLevel string) error {
+func UpdateLogLevel(newLevel LogLevel) error {
 	isCorrect := false
 	for _, level := range LogLevels {
 		if newLevel == level {
@@ -69,6 +72,24 @@ func UpdateLogLevel(newLevel string) error {
 	}
 
 	logLevel = newLevel
-	setupLoggers()
+	syncLogLevel()
 	return nil
 }
+
+func SetOnPanicFunc(onPanic func(format string, v ...interface{})) {
+	updateLoggers(func(l *logger) {
+		l.onPanic = onPanic
+	})
+}
+
+func SetOnFatalFunc(onFatal func(format string, v ...interface{})) {
+	updateLoggers(func(l *logger) {
+		l.onFatal = onFatal
+	})
+}
+
+func updateLoggers(updateFunc func(*logger)) {
+	for _, l := range []*logger{InfoLogger, WarningLogger, DebugLogger, ErrorLogger} {
+		updateFunc(l)
+	}
+}