From c13eae60c163cd30de682099598d114fb3767325 Mon Sep 17 00:00:00 2001
From: Joshua T Corbin <jcorbin@wunjo.org>
Date: Tue, 20 Dec 2016 11:49:31 -0800
Subject: [PATCH] Implement flag.{Value,Getter} for zap.Level (#214)

This makes it easy to add a command line log level flag.
---
 level.go      | 28 ++++++++++++++++++++++++++++
 level_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/level.go b/level.go
index 311263ef7..8aa0b0892 100644
--- a/level.go
+++ b/level.go
@@ -141,6 +141,34 @@ func (l *Level) UnmarshalText(text []byte) error {
 	return nil
 }
 
+// Set sets the level for the flag.Value interface.
+func (l *Level) Set(s string) error {
+	switch s {
+	case "debug":
+		*l = DebugLevel
+	case "info":
+		*l = InfoLevel
+	case "warn":
+		*l = WarnLevel
+	case "error":
+		*l = ErrorLevel
+	case "dpanic":
+		*l = DPanicLevel
+	case "panic":
+		*l = PanicLevel
+	case "fatal":
+		*l = FatalLevel
+	default:
+		return fmt.Errorf("unrecognized level: %q", s)
+	}
+	return nil
+}
+
+// Get gets the level for the flag.Getter interface.
+func (l *Level) Get() interface{} {
+	return *l
+}
+
 // Enabled returns true if the given level is at or above this level.
 func (l Level) Enabled(lvl Level) bool {
 	return lvl >= l
diff --git a/level_test.go b/level_test.go
index 70241f19c..f899b8547 100644
--- a/level_test.go
+++ b/level_test.go
@@ -21,6 +21,9 @@
 package zap
 
 import (
+	"bytes"
+	"flag"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -102,3 +105,29 @@ func TestLevelUnmarshalUnknownText(t *testing.T) {
 	err := l.UnmarshalText([]byte("foo"))
 	assert.Contains(t, err.Error(), "unrecognized level", "Expected unmarshaling arbitrary text to fail.")
 }
+
+func TestLevelAsFlagValue(t *testing.T) {
+	var (
+		lvl Level
+		buf bytes.Buffer
+	)
+
+	fs := flag.NewFlagSet("levelTest", flag.ContinueOnError)
+	fs.SetOutput(&buf)
+	fs.Var(&lvl, "level", "a log level")
+
+	// changing works
+	assert.Equal(t, InfoLevel, lvl)
+	assert.NoError(t, fs.Parse([]string{"-level", "warn"}))
+	assert.Equal(t, WarnLevel, lvl)
+	assert.NoError(t, fs.Parse([]string{"-level", "debug"}))
+	assert.Equal(t, DebugLevel, lvl)
+
+	// errors work
+	assert.Error(t, fs.Parse([]string{"-level", "nope"}))
+	assert.Equal(t,
+		`invalid value "nope" for flag -level: unrecognized level: "nope"`,
+		strings.Split(buf.String(), "\n")[0],
+		"expected error output")
+	buf.Reset()
+}