Repository: elliotchance/pie
Pull request #57 (2019-04-27 02:15:06)

Added Keys() and Values()
This means pie now supports functions for maps!

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/elliotchance/pie/57)
<!-- Reviewable:end -->



Patch:
diff --git a/README.md b/README.md
index 32e7550..280854b 100644
--- a/README.md
+++ b/README.md
@@ -118,33 +118,35 @@ This will only generate `myInts.Average`, `myInts.Sum` and `myStrings.Select`.
 
 # Functions
 
-| Function     | String | Number | Struct | Big-O    | Description |
-| ------------ | :----: | :----: | :----: | :------: | ----------- |
-| `All`        | ✓      | ✓      | ✓      | n        | All will return true if all callbacks return true. If the list is empty then true is always returned. |
-| `Any`        | ✓      | ✓      | ✓      | n        | Any will return true if any callbacks return true. If the list is empty then false is always returned. |
-| `Append`     | ✓      | ✓      | ✓      | n        | A new slice with the elements appended to the end. |
-| `AreSorted`  | ✓      | ✓      |        | n        | Check if the slice is already sorted. |
-| `AreUnique`  | ✓      | ✓      |        | n        | Check if the slice contains only unique elements. |
-| `Average`    |        | ✓      |        | n        | The average (mean) value, or a zeroed value. |
-| `Contains`   | ✓      | ✓      | ✓      | n        | Check if the value exists in the slice. |
-| `Extend`     | ✓      | ✓      | ✓      | n        | A new slice with the elements from each slice appended to the end. |
-| `First`      | ✓      | ✓      | ✓      | 1        | The first element, or a zeroed value. |
-| `FirstOr`    | ✓      | ✓      | ✓      | 1        | The first element, or a default value. |
-| `Join`       | ✓      |        |        | n        | A string from joining each of the elements. |
-| `JSONString` | ✓      | ✓      | ✓      | n        | The JSON encoded string. |
-| `Last`       | ✓      | ✓      | ✓      | 1        | The last element, or a zeroed value. |
-| `LastOr`     | ✓      | ✓      | ✓      | 1        | The last element, or a default value. |
-| `Len`        | ✓      | ✓      | ✓      | 1        | Number of elements. |
-| `Max`        | ✓      | ✓      |        | n        | The maximum value, or a zeroes value. |
-| `Min`        | ✓      | ✓      |        | n        | The minimum value, or a zeroed value. |
-| `Reverse`    | ✓      | ✓      | ✓      | n        | Reverse elements. |
-| `Select`     | ✓      | ✓      | ✓      | n        | A new slice containing only the elements that returned true from the condition. |
-| `Sort`       | ✓      | ✓      |        | n⋅log(n) | Return a new sorted slice. |
-| `Sum`        |        | ✓      |        | n        | Sum (total) of all elements. |
-| `ToStrings`  | ✓      | ✓      | ✓      | n        | Transforms each element to a string. |
-| `Transform`  | ✓      | ✓      | ✓      | n        | A new slice where each element has been transformed. |
-| `Unique`     | ✓      | ✓      |        | n⋅log(n) | Return a new slice with only unique elements. |
-| `Unselect`   | ✓      | ✓      | ✓      | n        | A new slice containing only the elements that returned false from the condition. |
+| Function     | String | Number | Struct| Maps | Big-O    | Description |
+| ------------ | :----: | :----: | :----:| :--: | :------: | ----------- |
+| `All`        | ✓      | ✓      | ✓     |      | n        | All will return true if all callbacks return true. If the list is empty then true is always returned. |
+| `Any`        | ✓      | ✓      | ✓     |      | n        | Any will return true if any callbacks return true. If the list is empty then false is always returned. |
+| `Append`     | ✓      | ✓      | ✓     |      | n        | A new slice with the elements appended to the end. |
+| `AreSorted`  | ✓      | ✓      |       |      | n        | Check if the slice is already sorted. |
+| `AreUnique`  | ✓      | ✓      |       |      | n        | Check if the slice contains only unique elements. |
+| `Average`    |        | ✓      |       |      | n        | The average (mean) value, or a zeroed value. |
+| `Contains`   | ✓      | ✓      | ✓     |      | n        | Check if the value exists in the slice. |
+| `Extend`     | ✓      | ✓      | ✓     |      | n        | A new slice with the elements from each slice appended to the end. |
+| `First`      | ✓      | ✓      | ✓     |      | 1        | The first element, or a zeroed value. |
+| `FirstOr`    | ✓      | ✓      | ✓     |      | 1        | The first element, or a default value. |
+| `Join`       | ✓      |        |       |      | n        | A string from joining each of the elements. |
+| `JSONString` | ✓      | ✓      | ✓     |      | n        | The JSON encoded string. |
+| `Keys`       |        |        |       | ✓    | n        | Returns all keys in the map (in random order). |
+| `Last`       | ✓      | ✓      | ✓     |      | 1        | The last element, or a zeroed value. |
+| `LastOr`     | ✓      | ✓      | ✓     |      | 1        | The last element, or a default value. |
+| `Len`        | ✓      | ✓      | ✓     |      | 1        | Number of elements. |
+| `Max`        | ✓      | ✓      |       |      | n        | The maximum value, or a zeroes value. |
+| `Min`        | ✓      | ✓      |       |      | n        | The minimum value, or a zeroed value. |
+| `Reverse`    | ✓      | ✓      | ✓     |      | n        | Reverse elements. |
+| `Select`     | ✓      | ✓      | ✓     |      | n        | A new slice containing only the elements that returned true from the condition. |
+| `Sort`       | ✓      | ✓      |       |      | n⋅log(n) | Return a new sorted slice. |
+| `Sum`        |        | ✓      |       |      | n        | Sum (total) of all elements. |
+| `ToStrings`  | ✓      | ✓      | ✓     |      | n        | Transforms each element to a string. |
+| `Transform`  | ✓      | ✓      | ✓     |      | n        | A new slice where each element has been transformed. |
+| `Unique`     | ✓      | ✓      |       |      | n⋅log(n) | Return a new slice with only unique elements. |
+| `Unselect`   | ✓      | ✓      | ✓     |      | n        | A new slice containing only the elements that returned false from the condition. |
+| `Values`     |        |        |       | ✓    | n        | Returns all values in the map (in random order). |
 
 # FAQ
 
@@ -194,7 +196,9 @@ add unit tests to `pie/float64s_test.go` and `pie/ints_test.go`.
 6. If you chose `ForAll` or `ForStrings`, then you must add unit tests to
 `pie/strings_test.go`.
 
-7. Update the README to list the new functions.
+7. If you chose `ForMaps`, then you must add unit tests to `pie/currencies.go`.
+
+8. Update the README to list the new functions.
 
 ## Why is the emoji a slice of pizza instead of a pie?
 
diff --git a/functions/keys.go b/functions/keys.go
new file mode 100644
index 0000000..1770f85
--- /dev/null
+++ b/functions/keys.go
@@ -0,0 +1,21 @@
+package functions
+
+// Keys returns the keys in the map. All of the items will be unique.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m MapType) Keys() KeySliceType {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make(KeySliceType, len(m))
+	for key := range m {
+		keys[i] = key
+		i++
+	}
+
+	return keys
+}
diff --git a/functions/main.go b/functions/main.go
index 1fdb18a..35bdce3 100644
--- a/functions/main.go
+++ b/functions/main.go
@@ -4,6 +4,7 @@ const (
 	ForNumbers = 1 << iota
 	ForStrings
 	ForStructs
+	ForMaps
 
 	ForAll               = ForNumbers | ForStrings | ForStructs
 	ForNumbersAndStrings = ForNumbers | ForStrings
@@ -29,6 +30,7 @@ var Functions = []struct {
 	{"FirstOr", "first_or.go", ForAll},
 	{"Join", "join.go", ForStrings},
 	{"JSONString", "json_string.go", ForAll},
+	{"Keys", "keys.go", ForMaps},
 	{"Last", "last.go", ForAll},
 	{"LastOr", "last_or.go", ForAll},
 	{"Len", "len.go", ForAll},
@@ -42,11 +44,15 @@ var Functions = []struct {
 	{"Transform", "transform.go", ForAll},
 	{"Unique", "unique.go", ForNumbersAndStrings},
 	{"Unselect", "unselect.go", ForAll},
+	{"Values", "values.go", ForMaps},
 }
 
 type ElementType float64
 type SliceType []ElementType
 type StringElementType string
 type StringSliceType []StringElementType
+type KeyType string
+type KeySliceType []KeyType
+type MapType map[KeyType]ElementType
 
 var ElementZeroValue ElementType
diff --git a/functions/values.go b/functions/values.go
new file mode 100644
index 0000000..5f92925
--- /dev/null
+++ b/functions/values.go
@@ -0,0 +1,21 @@
+package functions
+
+// Values returns the values in the map.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m MapType) Values() []ElementType {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make([]ElementType, len(m))
+	for _, value := range m {
+		keys[i] = value
+		i++
+	}
+
+	return keys
+}
diff --git a/main.go b/main.go
index fb21a12..4fa0689 100644
--- a/main.go
+++ b/main.go
@@ -34,7 +34,19 @@ func getIdentName(e ast.Expr) string {
 	}
 }
 
-func findType(pkgs map[string]*ast.Package, name string) (packageName, elementType string) {
+func getKeyAndElementType(pkgName, name string, typeSpec *ast.TypeSpec) (string, string, string) {
+	if t, ok := typeSpec.Type.(*ast.ArrayType); ok {
+		return pkgName, "", getIdentName(t.Elt)
+	}
+
+	if t, ok := typeSpec.Type.(*ast.MapType); ok {
+		return pkgName, getIdentName(t.Key), getIdentName(t.Value)
+	}
+
+	panic(fmt.Sprintf("type %s must be a slice or map", name))
+}
+
+func findType(pkgs map[string]*ast.Package, name string) (packageName, keyType, elementType string) {
 	for pkgName, pkg := range pkgs {
 		for _, file := range pkg.Files {
 			for _, decl := range file.Decls {
@@ -42,11 +54,7 @@ func findType(pkgs map[string]*ast.Package, name string) (packageName, elementTy
 					for _, spec := range genDecl.Specs {
 						if typeSpec, ok := spec.(*ast.TypeSpec); ok {
 							if typeSpec.Name.String() == name {
-								if t, ok := typeSpec.Type.(*ast.ArrayType); ok {
-									return pkgName, getIdentName(t.Elt)
-								} else {
-									panic(fmt.Sprintf("type %s must be a slice", name))
-								}
+								return getKeyAndElementType(pkgName, name, typeSpec)
 							}
 						}
 					}
@@ -58,8 +66,12 @@ func findType(pkgs map[string]*ast.Package, name string) (packageName, elementTy
 	panic(fmt.Sprintf("type %s does not exist", name))
 }
 
-func getType(name string) int {
-	switch name {
+func getType(keyType, elementType string) int {
+	if keyType != "" {
+		return functions.ForMaps
+	}
+
+	switch elementType {
 	case "int8", "uint8", "byte", "int16", "uint16", "int32", "rune", "uint32",
 		"int64", "uint64", "int", "uint", "uintptr", "float32", "float64",
 		"complex64", "complex128":
@@ -123,11 +135,11 @@ func main() {
 	check(err)
 
 	for _, arg := range os.Args[1:] {
-		sliceType, fns := getFunctionsFromArg(arg)
-		packageName, elementType := findType(pkgs, sliceType)
-		kind := getType(elementType)
+		mapOrSliceType, fns := getFunctionsFromArg(arg)
+		packageName, keyType, elementType := findType(pkgs, mapOrSliceType)
+		kind := getType(keyType, elementType)
 
-		templates := []string{}
+		var templates []string
 		for _, function := range functions.Functions {
 			if len(fns) > 0 && !pie.Strings(fns).Contains(function.Name) {
 				continue
@@ -155,10 +167,13 @@ func main() {
 			t += tmpl[i:] + "\n"
 		}
 
-		t = strings.Replace(t, "StringSliceType", sliceType, -1)
+		t = strings.Replace(t, "StringSliceType", mapOrSliceType, -1)
 		t = strings.Replace(t, "StringElementType", elementType, -1)
-		t = strings.Replace(t, "SliceType", sliceType, -1)
 		t = strings.Replace(t, "ElementType", elementType, -1)
+		t = strings.Replace(t, "MapType", mapOrSliceType, -1)
+		t = strings.Replace(t, "KeyType", elementType, -1)
+		t = strings.Replace(t, "KeySliceType", "[]"+keyType, -1)
+		t = strings.Replace(t, "SliceType", mapOrSliceType, -1)
 
 		switch kind {
 		case functions.ForNumbers:
@@ -187,15 +202,12 @@ func main() {
 		// with go fmt.
 		t = strings.TrimRight(t, "\n") + "\n"
 
-		// Filter out any functions we dont want.
-		//t = filterFunctions(t, functions)
-
-		err := ioutil.WriteFile(strings.ToLower(sliceType)+"_pie.go", []byte(t), 0755)
+		err := ioutil.WriteFile(strings.ToLower(mapOrSliceType)+"_pie.go", []byte(t), 0755)
 		check(err)
 	}
 }
 
-func getFunctionsFromArg(arg string) (string, []string) {
+func getFunctionsFromArg(arg string) (mapOrSliceType string, fns []string) {
 	parts := strings.Split(arg, ".")
 
 	return parts[0], parts[1:]
diff --git a/pie/currencies.go b/pie/currencies.go
new file mode 100644
index 0000000..2432807
--- /dev/null
+++ b/pie/currencies.go
@@ -0,0 +1,13 @@
+package pie
+
+type currency struct {
+	NumericCode, Exponent int
+}
+
+//go:generate pie currencies
+type currencies map[string]currency
+
+var isoCurrencies = currencies{
+	"AUD": {36, -2},
+	"USD": {840, -2},
+}
diff --git a/pie/currencies_pie.go b/pie/currencies_pie.go
new file mode 100755
index 0000000..0f6eb69
--- /dev/null
+++ b/pie/currencies_pie.go
@@ -0,0 +1,41 @@
+package pie
+
+// Keys returns the keys in the map. All of the items will be unique.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m currencies) Keys() []string {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make([]string, len(m))
+	for key := range m {
+		keys[i] = key
+		i++
+	}
+
+	return keys
+}
+
+// Values returns the values in the map.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m currencies) Values() []currency {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make([]currency, len(m))
+	for _, value := range m {
+		keys[i] = value
+		i++
+	}
+
+	return keys
+}
diff --git a/template.go b/template.go
index 1c39a45..c41a70f 100644
--- a/template.go
+++ b/template.go
@@ -164,6 +164,28 @@ func (ss StringSliceType) Join(glue string) (s string) {
 
 	return s
 }
+`,
+	"Keys": `package functions
+
+// Keys returns the keys in the map. All of the items will be unique.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m MapType) Keys() KeySliceType {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make(KeySliceType, len(m))
+	for key := range m {
+		keys[i] = key
+		i++
+	}
+
+	return keys
+}
 `,
 	"Last": `package functions
 
@@ -392,5 +414,27 @@ func (ss SliceType) Unselect(condition func(ElementType) bool) (ss2 SliceType) {
 
 	return
 }
+`,
+	"Values": `package functions
+
+// Values returns the values in the map.
+//
+// Due to Go's randomization of iterating maps the order is not deterministic.
+func (m MapType) Values() []ElementType {
+	// Avoid allocation
+	l := len(m)
+	if l == 0 {
+		return nil
+	}
+
+	i := 0
+	keys := make([]ElementType, len(m))
+	for _, value := range m {
+		keys[i] = value
+		i++
+	}
+
+	return keys
+}
 `,
 }


Test patch:
diff --git a/pie/currencies_test.go b/pie/currencies_test.go
new file mode 100644
index 0000000..ba39c03
--- /dev/null
+++ b/pie/currencies_test.go
@@ -0,0 +1,30 @@
+package pie
+
+import (
+	"github.com/stretchr/testify/assert"
+	"sort"
+	"testing"
+)
+
+func TestCurrencies_Keys(t *testing.T) {
+	assert.Equal(t, []string(nil), currencies(nil).Keys())
+
+	assert.Equal(t, []string(nil), currencies{}.Keys())
+
+	keys := isoCurrencies.Keys()
+	sort.Strings(keys)
+	assert.Equal(t, []string{"AUD", "USD"}, keys)
+}
+
+func TestCurrencies_Values(t *testing.T) {
+	assert.Equal(t, []currency(nil), currencies(nil).Values())
+
+	assert.Equal(t, []currency(nil), currencies{}.Values())
+
+	values := isoCurrencies.Values()
+	sort.Slice(values, func(i, j int) bool {
+		return values[i].NumericCode < values[j].NumericCode
+	})
+
+	assert.Equal(t, []currency{{36, -2}, {840, -2}}, values)
+}