Repository: elliotchance/pie
Pull request #34 (2019-04-11 13:49:44)

Add support for slices of struct pointers
Slices of pointers have all the same functionality but there is a slightly different syntax.

<!-- 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/34)
<!-- Reviewable:end -->



Patch:
diff --git a/README.md b/README.md
index 47d383c..a110196 100644
--- a/README.md
+++ b/README.md
@@ -118,25 +118,25 @@ This will only generate `myInts.Average`, `myInts.Sum` and `myStrings.Only`.
 
 # Functions
 
-| Function     | Description                                                                      | String | Number | Struct | Big-O      |
-| ------------ | -------------------------------------------------------------------------------- | :----: | :----: | :----: | :--------: |
-| `AreSorted`  | Check if the slice is already sorted.                                            | ✓      | ✓      |        | n          |
-| `Average`    | The average (mean) value, or a zeroed value.                                     |        | ✓      |        | n          |
-| `Contains`   | Check if the value exists in the slice.                                          | ✓      | ✓      | ✓      | n          |
-| `First`      | The first element, or a zeroed value.                                            | ✓      | ✓      | ✓      | 1          |
-| `FirstOr`    | The first element, or a default value.                                           | ✓      | ✓      | ✓      | 1          |
-| `JSONString` | The JSON encoded string.                                                         | ✓      | ✓      | ✓      | n          |
-| `Last`       | The last element, or a zeroed value.                                             | ✓      | ✓      | ✓      | 1          |
-| `LastOr`     | The last element, or a default value.                                            | ✓      | ✓      | ✓      | 1          |
-| `Len`        | Number of elements.                                                              | ✓      | ✓      | ✓      | 1          |
-| `Max`        | The maximum value, or a zeroes value.                                            | ✓      | ✓      |        | n          |
-| `Min`        | The minimum value, or a zeroed value.                                            | ✓      | ✓      |        | n          |
-| `Only`       | A new slice containing only the elements that returned true from the condition.  | ✓      | ✓      | ✓      | n          |
-| `Reverse`    | Reverse elements.                                                                | ✓      | ✓      | ✓      | n          |
-| `Sort`       | Return a new sorted slice.                                                       | ✓      | ✓      |        | *n⋅log(n)* |
-| `Sum`        | Sum (total) of all elements.                                                     |        | ✓      |        | n          |
-| `Transform`  | A new slice where each element has been transformed.                             | ✓      | ✓      | ✓      | n          |
-| `Without`    | A new slice containing only the elements that returned false from the condition. | ✓      | ✓      | ✓      | n          |
+| Function     | String | Number | Struct | Big-O      | Description                                                                      |
+| ------------ | :----: | :----: | :----: | :--------: | -------------------------------------------------------------------------------- |
+| `AreSorted`  | ✓      | ✓      |        | n          | Check if the slice is already sorted.                                            |
+| `Average`    |        | ✓      |        | n          | The average (mean) value, or a zeroed value.                                     |
+| `Contains`   | ✓      | ✓      | ✓      | n          | Check if the value exists in the slice.                                          |
+| `First`      | ✓      | ✓      | ✓      | 1          | The first element, or a zeroed value.                                            |
+| `FirstOr`    | ✓      | ✓      | ✓      | 1          | The first element, or a default value.                                           |
+| `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.                                            |
+| `Only`       | ✓      | ✓      | ✓      | n          | A new slice containing only the elements that returned true from the condition.  |
+| `Reverse`    | ✓      | ✓      | ✓      | n          | Reverse elements.                                                                |
+| `Sort`       | ✓      | ✓      |        | *n⋅log(n)* | Return a new sorted slice.                                                       |
+| `Sum`        |        | ✓      |        | n          | Sum (total) of all elements.                                                     |
+| `Transform`  | ✓      | ✓      | ✓      | n          | A new slice where each element has been transformed.                             |
+| `Without`    | ✓      | ✓      | ✓      | n          | A new slice containing only the elements that returned false from the condition. |
 
 # FAQ
 
diff --git a/main.go b/main.go
index 7bf2e1a..dc321c9 100644
--- a/main.go
+++ b/main.go
@@ -27,6 +27,19 @@ func check(err error) {
 	}
 }
 
+func getIdentName(e ast.Expr) string {
+	switch v := e.(type) {
+	case *ast.Ident:
+		return v.Name
+
+	case *ast.StarExpr:
+		return "*" + getIdentName(v.X)
+
+	default:
+		panic(fmt.Sprintf("cannot decode %T", e))
+	}
+}
+
 func findType(pkgs map[string]*ast.Package, name string) (packageName, elementType string) {
 	for pkgName, pkg := range pkgs {
 		for _, file := range pkg.Files {
@@ -36,7 +49,7 @@ func findType(pkgs map[string]*ast.Package, name string) (packageName, elementTy
 						if typeSpec, ok := spec.(*ast.TypeSpec); ok {
 							if typeSpec.Name.String() == name {
 								if t, ok := typeSpec.Type.(*ast.ArrayType); ok {
-									return pkgName, t.Elt.(*ast.Ident).Name
+									return pkgName, getIdentName(t.Elt)
 								} else {
 									panic(fmt.Sprintf("type %s must be a slice", name))
 								}
@@ -119,7 +132,6 @@ func main() {
 		}
 
 		// Aggregate imports.
-
 		t := fmt.Sprintf("package %s\n\nimport (", packageName)
 		for _, imp := range getAllImports(templates) {
 			t += fmt.Sprintf("\n\t%s", imp)
@@ -142,7 +154,15 @@ func main() {
 			t = strings.Replace(t, "ElementZeroValue", `""`, -1)
 
 		case "struct":
-			t = strings.Replace(t, "ElementZeroValue", fmt.Sprintf("%s{}", elementType), -1)
+			zeroValue := fmt.Sprintf("%s{}", elementType)
+
+			// If its a pointer we need to replace '*' -> '&' when
+			// instantiating.
+			if elementType[0] == '*' {
+				zeroValue = "&" + zeroValue[1:]
+			}
+
+			t = strings.Replace(t, "ElementZeroValue", zeroValue, -1)
 		}
 
 		// The TrimRight is important to remove an extra new line that conflicts
diff --git a/pie/car.go b/pie/car.go
index daedc2e..8181825 100644
--- a/pie/car.go
+++ b/pie/car.go
@@ -1,7 +1,8 @@
 package pie
 
-//go:generate pie cars
+//go:generate pie cars carPointers
 type cars []car
+type carPointers []*car
 
 type car struct {
 	Name, Color string
diff --git a/pie/carpointers_pie.go b/pie/carpointers_pie.go
new file mode 100755
index 0000000..2594c1b
--- /dev/null
+++ b/pie/carpointers_pie.go
@@ -0,0 +1,100 @@
+package pie
+
+import (
+	"encoding/json"
+)
+
+func (ss carPointers) Contains(lookingFor *car) bool {
+	for _, s := range ss {
+		if s == lookingFor {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (ss carPointers) Only(condition func(*car) bool) (ss2 carPointers) {
+	for _, s := range ss {
+		if condition(s) {
+			ss2 = append(ss2, s)
+		}
+	}
+
+	return
+}
+
+func (ss carPointers) Without(condition func(*car) bool) (ss2 carPointers) {
+	for _, s := range ss {
+		if !condition(s) {
+			ss2 = append(ss2, s)
+		}
+	}
+
+	return
+}
+
+func (ss carPointers) Transform(fn func(*car) *car) (ss2 carPointers) {
+	if ss == nil {
+		return nil
+	}
+
+	ss2 = make([]*car, len(ss))
+	for i, s := range ss {
+		ss2[i] = fn(s)
+	}
+
+	return
+}
+
+func (ss carPointers) FirstOr(defaultValue *car) *car {
+	if len(ss) == 0 {
+		return defaultValue
+	}
+
+	return ss[0]
+}
+
+func (ss carPointers) LastOr(defaultValue *car) *car {
+	if len(ss) == 0 {
+		return defaultValue
+	}
+
+	return ss[len(ss)-1]
+}
+
+func (ss carPointers) First() *car {
+	return ss.FirstOr(&car{})
+}
+
+func (ss carPointers) Last() *car {
+	return ss.LastOr(&car{})
+}
+
+func (ss carPointers) Len() int {
+	return len(ss)
+}
+
+func (ss carPointers) JSONString() string {
+	if ss == nil {
+		return "[]"
+	}
+
+	data, _ := json.Marshal(ss)
+
+	return string(data)
+}
+
+func (ss carPointers) Reverse() carPointers {
+
+	if len(ss) < 2 {
+		return ss
+	}
+
+	sorted := make([]*car, len(ss))
+	for i := 0; i < len(ss); i++ {
+		sorted[i] = ss[len(ss)-i-1]
+	}
+
+	return sorted
+}
diff --git a/pie_all.go b/pie_all.go
index 30b0bc3..a1c0ad7 100644
--- a/pie_all.go
+++ b/pie_all.go
@@ -7,6 +7,8 @@ import (
 // The functions in this file work for all slices types.
 
 // Contains returns true if the element exists in the slice.
+//
+// When using slices of pointers it will only compare by address, not value.
 func (ss SliceType) Contains(lookingFor ElementType) bool {
 	for _, s := range ss {
 		if s == lookingFor {
@@ -46,6 +48,10 @@ func (ss SliceType) Without(condition func(ElementType) bool) (ss2 SliceType) {
 
 // Transform will return a new slice where each element has been transformed.
 // The number of element returned will always be the same as the input.
+//
+// Be careful when using this with slices of pointers. If you modify the input
+// value it will affect the original slice. Be sure to return a new allocated
+// object or deep copy the existing one.
 func (ss SliceType) Transform(fn func(ElementType) ElementType) (ss2 SliceType) {
 	if ss == nil {
 		return nil
diff --git a/template.go b/template.go
index 3f70cc1..59acae4 100644
--- a/template.go
+++ b/template.go
@@ -10,6 +10,8 @@ import (
 // The functions in this file work for all slices types.
 
 // Contains returns true if the element exists in the slice.
+//
+// When using slices of pointers it will only compare by address, not value.
 func (ss SliceType) Contains(lookingFor ElementType) bool {
 	for _, s := range ss {
 		if s == lookingFor {
@@ -49,6 +51,10 @@ func (ss SliceType) Without(condition func(ElementType) bool) (ss2 SliceType) {
 
 // Transform will return a new slice where each element has been transformed.
 // The number of element returned will always be the same as the input.
+//
+// Be careful when using this with slices of pointers. If you modify the input
+// value it will affect the original slice. Be sure to return a new allocated
+// object or deep copy the existing one.
 func (ss SliceType) Transform(fn func(ElementType) ElementType) (ss2 SliceType) {
 	if ss == nil {
 		return nil


Test patch:
diff --git a/pie/carpointers_test.go b/pie/carpointers_test.go
new file mode 100644
index 0000000..01e03f4
--- /dev/null
+++ b/pie/carpointers_test.go
@@ -0,0 +1,295 @@
+package pie
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strings"
+	"testing"
+)
+
+var carPointerA = &car{"a", "green"}
+var carPointerB = &car{"b", "blue"}
+var carPointerC = &car{"c", "gray"}
+var carPointerEmpty = &car{}
+
+var carPointersContainsTests = []struct {
+	ss       carPointers
+	contains *car
+	expected bool
+}{
+	{nil, carPointerA, false},
+	{nil, carPointerEmpty, false},
+	{nil, nil, false},
+	{carPointers{carPointerA, carPointerB, carPointerC}, carPointerA, true},
+	{carPointers{carPointerA, carPointerB, carPointerC}, carPointerB, true},
+	{carPointers{carPointerA, carPointerB, carPointerC}, carPointerC, true},
+	{carPointers{carPointerA, carPointerB, carPointerC}, &car{"a", "green"}, false},
+	{carPointers{carPointerA, carPointerB, carPointerC}, &car{"A", ""}, false},
+	{carPointers{carPointerA, carPointerB, carPointerC}, &car{}, false},
+	{carPointers{carPointerA, carPointerB, carPointerC}, &car{"d", ""}, false},
+	{carPointers{carPointerA, carPointerEmpty, carPointerC}, carPointerEmpty, true},
+	{carPointers{carPointerA, nil, carPointerC}, nil, true},
+	{carPointers{carPointerA, carPointerEmpty, carPointerC}, nil, false},
+}
+
+func TestCarPointers_Contains(t *testing.T) {
+	for _, test := range carPointersContainsTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.expected, test.ss.Contains(test.contains))
+		})
+	}
+}
+
+var carPointersOnlyAndWithoutTests = []struct {
+	ss                carPointers
+	condition         func(*car) bool
+	expectedOnly      carPointers
+	expectedWithout   carPointers
+	expectedTransform carPointers
+}{
+	{
+		nil,
+		func(s *car) bool {
+			return s.Name == ""
+		},
+		nil,
+		nil,
+		nil,
+	},
+	{
+		carPointers{carPointerA, carPointerB, carPointerC},
+		func(s *car) bool {
+			return s.Name != "b"
+		},
+		carPointers{carPointerA, carPointerC},
+		carPointers{carPointerB},
+		carPointers{&car{"A", "green"}, &car{"B", "blue"}, &car{"C", "gray"}},
+	},
+}
+
+func TestCarPointers_Only(t *testing.T) {
+	for _, test := range carPointersOnlyAndWithoutTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.expectedOnly, test.ss.Only(test.condition))
+		})
+	}
+}
+
+func TestCarPointers_Without(t *testing.T) {
+	for _, test := range carPointersOnlyAndWithoutTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.expectedWithout, test.ss.Without(test.condition))
+		})
+	}
+}
+
+func TestCarPointers_Transform(t *testing.T) {
+	for _, test := range carPointersOnlyAndWithoutTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.expectedTransform, test.ss.Transform(func(c *car) *car {
+				return &car{
+					Name:  strings.ToUpper(c.Name),
+					Color: c.Color,
+				}
+			}))
+		})
+	}
+}
+
+var carPointersFirstAndLastTests = []struct {
+	ss             carPointers
+	first, firstOr *car
+	last, lastOr   *car
+}{
+	{
+		nil,
+		&car{},
+		&car{"default1", "unknown"},
+		&car{},
+		&car{"default2", "unknown"},
+	},
+	{
+		carPointers{&car{"foo", "red"}},
+		&car{"foo", "red"},
+		&car{"foo", "red"},
+		&car{"foo", "red"},
+		&car{"foo", "red"},
+	},
+	{
+		carPointers{carPointerA, carPointerB},
+		carPointerA,
+		carPointerA,
+		carPointerB,
+		carPointerB,
+	},
+	{
+		carPointers{carPointerA, carPointerB, carPointerC},
+		carPointerA,
+		carPointerA,
+		carPointerC,
+		carPointerC,
+	},
+}
+
+func TestCarPointers_FirstOr(t *testing.T) {
+	for _, test := range carPointersFirstAndLastTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.firstOr, test.ss.FirstOr(&car{"default1", "unknown"}))
+		})
+	}
+}
+
+func TestCarPointers_LastOr(t *testing.T) {
+	for _, test := range carPointersFirstAndLastTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.lastOr, test.ss.LastOr(&car{"default2", "unknown"}))
+		})
+	}
+}
+
+func TestCarPointers_First(t *testing.T) {
+	for _, test := range carPointersFirstAndLastTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.first, test.ss.First())
+		})
+	}
+}
+
+func TestCarPointers_Last(t *testing.T) {
+	for _, test := range carPointersFirstAndLastTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.last, test.ss.Last())
+		})
+	}
+}
+
+var carPointersStatsTests = []struct {
+	ss       carPointers
+	min, max *car
+	len      int
+}{
+	{
+		nil,
+		&car{},
+		&car{},
+		0,
+	},
+	{
+		carPointers{},
+		&car{},
+		&car{},
+		0,
+	},
+	{
+		carPointers{&car{"foo", "red"}},
+		&car{"foo", "red"},
+		&car{"foo", "red"},
+		1,
+	},
+	{
+		carPointers{&car{"bar", "yellow"}, &car{"Baz", "black"}, &car{"qux", "cyan"}, &car{"foo", "red"}},
+		&car{"Baz", "black"},
+		&car{"qux", "cyan"},
+		4,
+	},
+}
+
+func TestCarPointers_Len(t *testing.T) {
+	for _, test := range carPointersStatsTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.len, test.ss.Len())
+		})
+	}
+}
+
+var carPointersJSONTests = []struct {
+	ss         carPointers
+	jsonString string
+}{
+	{
+		nil,
+		`[]`, // Instead of null.
+	},
+	{
+		carPointers{},
+		`[]`,
+	},
+	{
+		carPointers{&car{"foo", "red"}},
+		`[{"Name":"foo","Color":"red"}]`,
+	},
+	{
+		carPointers{&car{"bar", "yellow"}, &car{"Baz", "black"}, &car{"qux", "cyan"}, &car{"foo", "red"}},
+		`[{"Name":"bar","Color":"yellow"},{"Name":"Baz","Color":"black"},{"Name":"qux","Color":"cyan"},{"Name":"foo","Color":"red"}]`,
+	},
+}
+
+func TestCarPointers_JSONString(t *testing.T) {
+	for _, test := range carPointersJSONTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.jsonString, test.ss.JSONString())
+		})
+	}
+}
+
+var carPointersSortTests = []struct {
+	ss        carPointers
+	sorted    carPointers
+	reversed  carPointers
+	areSorted bool
+}{
+	{
+		nil,
+		nil,
+		nil,
+		true,
+	},
+	{
+		carPointers{},
+		carPointers{},
+		carPointers{},
+		true,
+	},
+	{
+		carPointers{&car{"foo", "red"}},
+		carPointers{&car{"foo", "red"}},
+		carPointers{&car{"foo", "red"}},
+		true,
+	},
+	{
+		carPointers{&car{"bar", "yellow"}, &car{"Baz", "black"}, &car{"foo", "red"}},
+		carPointers{&car{"Baz", "black"}, &car{"bar", "yellow"}, &car{"foo", "red"}},
+		carPointers{&car{"foo", "red"}, &car{"Baz", "black"}, &car{"bar", "yellow"}},
+		false,
+	},
+	{
+		carPointers{&car{"bar", "yellow"}, &car{"Baz", "black"}, &car{"qux", "cyan"}, &car{"foo", "red"}},
+		carPointers{&car{"Baz", "black"}, &car{"bar", "yellow"}, &car{"foo", "red"}, &car{"qux", "cyan"}},
+		carPointers{&car{"foo", "red"}, &car{"qux", "cyan"}, &car{"Baz", "black"}, &car{"bar", "yellow"}},
+		false,
+	},
+	{
+		carPointers{&car{"Baz", "black"}, &car{"bar", "yellow"}},
+		carPointers{&car{"Baz", "black"}, &car{"bar", "yellow"}},
+		carPointers{&car{"bar", "yellow"}, &car{"Baz", "black"}},
+		true,
+	},
+}
+
+func TestCarPointers_Reverse(t *testing.T) {
+	for _, test := range carPointersSortTests {
+		t.Run("", func(t *testing.T) {
+			defer assertImmutableCarPointers(t, &test.ss)()
+			assert.Equal(t, test.reversed, test.ss.Reverse())
+		})
+	}
+}
diff --git a/pie/main_test.go b/pie/main_test.go
index 0f6dff0..f0899c7 100644
--- a/pie/main_test.go
+++ b/pie/main_test.go
@@ -44,3 +44,13 @@ func assertImmutableCars(t *testing.T, ss *cars) func() {
 		assert.True(t, before == after)
 	}
 }
+
+func assertImmutableCarPointers(t *testing.T, ss *carPointers) func() {
+	before := (*ss).JSONString()
+
+	return func() {
+		after := (*ss).JSONString()
+		assert.Equal(t, before, after)
+		assert.True(t, before == after)
+	}
+}