From b7b2ee34de4dacda499f0e1c542a21610c3fb940 Mon Sep 17 00:00:00 2001
From: Daniel Smith <dbsmith@google.com>
Date: Fri, 2 Jan 2015 13:57:10 -0800
Subject: [PATCH] add flag helper func

---
 pkg/api/resource/quantity.go      | 32 +++++++++++++++++++++++++++++++
 pkg/api/resource/quantity_test.go | 12 ++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/pkg/api/resource/quantity.go b/pkg/api/resource/quantity.go
index d7e2ab595d3..e0020e1d991 100644
--- a/pkg/api/resource/quantity.go
+++ b/pkg/api/resource/quantity.go
@@ -18,6 +18,8 @@ package resource
 
 import (
 	"errors"
+	"flag"
+	"fmt"
 	"math/big"
 	"regexp"
 	"strings"
@@ -372,3 +374,33 @@ func (q *Quantity) Copy() *Quantity {
 		Format: q.Format,
 	}
 }
+
+// qFlag is a helper type for the Flag function
+type qFlag struct {
+	dest *Quantity
+}
+
+func (qf qFlag) Set(val string) error {
+	q, err := ParseQuantity(val)
+	if err != nil {
+		return err
+	}
+	// This copy is OK because q will not be referenced again.
+	*qf.dest = *q
+	return nil
+}
+
+func (qf qFlag) String() string {
+	return qf.dest.String()
+}
+
+// QuantityFlag is a helper that makes a quantity flag (using standard flag package).
+// Will panic if defaultValue is not a valid quantity.
+func QuantityFlag(flagName, defaultValue, description string) *Quantity {
+	q, err := ParseQuantity(defaultValue)
+	if err != nil {
+		panic(fmt.Errorf("can't use %v as a quantity: %v", defaultValue, err))
+	}
+	flag.Var(qFlag{q}, flagName, description)
+	return q
+}
diff --git a/pkg/api/resource/quantity_test.go b/pkg/api/resource/quantity_test.go
index ee1c1669cc1..5a4d71bb3cb 100644
--- a/pkg/api/resource/quantity_test.go
+++ b/pkg/api/resource/quantity_test.go
@@ -25,6 +25,10 @@ import (
 	"speter.net/go/exp/math/dec/inf"
 )
 
+var (
+	testQuantityFlag = QuantityFlag("quantityFlag", "1M", "dummy flag for testing the quantity flag mechanism")
+)
+
 func dec(i int64, exponent int) *inf.Dec {
 	// See the below test-- scale is the negative of an exponent.
 	return inf.NewDec(i, inf.Scale(-exponent))
@@ -425,3 +429,11 @@ func TestCopy(t *testing.T) {
 		t.Errorf("Copy didn't")
 	}
 }
+
+func TestQFlagSet(t *testing.T) {
+	qf := qFlag{&Quantity{}}
+	qf.Set("1Ki")
+	if e, a := "1Ki", qf.String(); e != a {
+		t.Errorf("Unexpected result %v != %v", e, a)
+	}
+}