Skip to content

Commit bd86b00

Browse files
thoasthoas
authored
Merge pull request thoas#112 from vellotis/feature-flatmap-and-flatten
Add `FlatMap` and `Flatten` transformation operations
2 parents 5035611 + 88055cb commit bd86b00

File tree

8 files changed

+238
-15
lines changed

8 files changed

+238
-15
lines changed

‎README.rst

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,29 @@ Manipulates an iteratee (map, slice) and transforms it to another type:
305305
return fmt.Sprintf("%d", k), v
306306
}) // map[string]string{"1": "Florent", "2": "Gilles"}
307307
308+
funk.FlatMap
309+
........
310+
311+
Manipulates an iteratee (map, slice) and transforms it to to a flattened collection of another type:
312+
313+
* map -> slice
314+
* slice -> slice
315+
316+
.. code-block:: go
317+
318+
r := funk.FlatMap([][]int{{1, 2}, {3, 4}}, func(x []int) []int {
319+
return append(x, 0)
320+
}) // []int{1, 2, 0, 3, 4, 0}
321+
322+
mapping := map[string][]int{
323+
"Florent": {1, 2},
324+
"Gilles": {3, 4},
325+
}
326+
327+
r = funk.FlatMap(mapping, func(k string, v []int) []int {
328+
return v
329+
}) // []int{1, 2, 3, 4}
330+
308331
funk.Get
309332
........
310333

‎builder.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,13 @@ type Builder interface {
1111
Compact() Builder
1212
Drop(n int) Builder
1313
Filter(predicate interface{}) Builder
14+
Flatten() Builder
1415
FlattenDeep() Builder
1516
Initial() Builder
1617
Intersect(y interface{}) Builder
1718
Join(rarr interface{}, fnc JoinFnc) Builder
1819
Map(mapFunc interface{}) Builder
20+
FlatMap(mapFunc interface{}) Builder
1921
Reverse() Builder
2022
Shuffle() Builder
2123
Tail() Builder

‎chain_builder.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ func (b *chainBuilder) Drop(n int) Builder {
2121
func (b *chainBuilder) Filter(predicate interface{}) Builder {
2222
return &chainBuilder{Filter(b.collection, predicate)}
2323
}
24+
func (b *chainBuilder) Flatten() Builder {
25+
return &chainBuilder{Flatten(b.collection)}
26+
}
2427
func (b *chainBuilder) FlattenDeep() Builder {
2528
return &chainBuilder{FlattenDeep(b.collection)}
2629
}
@@ -36,6 +39,9 @@ func (b *chainBuilder) Join(rarr interface{}, fnc JoinFnc) Builder {
3639
func (b *chainBuilder) Map(mapFunc interface{}) Builder {
3740
return &chainBuilder{Map(b.collection, mapFunc)}
3841
}
42+
func (b *chainBuilder) FlatMap(mapFunc interface{}) Builder {
43+
return &chainBuilder{FlatMap(b.collection, mapFunc)}
44+
}
3945
func (b *chainBuilder) Reverse() Builder {
4046
return &chainBuilder{Reverse(b.collection)}
4147
}

‎chain_builder_test.go

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,30 @@ func TestChainFilter_SideEffect(t *testing.T) {
144144
is.NotEqual([]*foo{&foo{"foo"}, &foo{"bar"}}, in)
145145
}
146146

147+
func TestChainFlatten(t *testing.T) {
148+
testCases := []struct {
149+
In interface{}
150+
}{
151+
{
152+
In: [][]int{{1, 2}, {3, 4}},
153+
},
154+
{
155+
In: [][][]int{{{1, 2}, {3, 4}}, {{5, 6}}},
156+
},
157+
}
158+
159+
for idx, tc := range testCases {
160+
t.Run(fmt.Sprintf("test case #%d", idx+1), func(t *testing.T) {
161+
is := assert.New(t)
162+
163+
expected := Flatten(tc.In)
164+
actual := Chain(tc.In).Flatten().Value()
165+
166+
is.Equal(expected, actual)
167+
})
168+
}
169+
}
170+
147171
func TestChainFlattenDeep(t *testing.T) {
148172
testCases := []struct {
149173
In interface{}
@@ -261,6 +285,33 @@ func TestChainMap(t *testing.T) {
261285
}
262286
}
263287

288+
func TestChainFlatMap(t *testing.T) {
289+
testCases := []struct {
290+
In interface{}
291+
FlatMapFnc interface{}
292+
}{
293+
{
294+
In: [][]int{{1}, {2}, {3}, {4}},
295+
FlatMapFnc: func(x []int) []int { return x },
296+
},
297+
{
298+
In: map[string][]int{"Florent": {1}, "Gilles": {2}},
299+
FlatMapFnc: func(k string, v []int) []int { return v },
300+
},
301+
}
302+
303+
for idx, tc := range testCases {
304+
t.Run(fmt.Sprintf("test case #%d", idx+1), func(t *testing.T) {
305+
is := assert.New(t)
306+
307+
expected := FlatMap(tc.In, tc.FlatMapFnc)
308+
actual := Chain(tc.In).FlatMap(tc.FlatMapFnc).Value()
309+
310+
is.ElementsMatch(expected, actual)
311+
})
312+
}
313+
}
314+
264315
func TestChainMap_SideEffect(t *testing.T) {
265316
is := assert.New(t)
266317

‎lazy_builder.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ func (b *lazyBuilder) Drop(n int) Builder {
1818
func (b *lazyBuilder) Filter(predicate interface{}) Builder {
1919
return &lazyBuilder{func() interface{} { return Filter(b.exec(), predicate) }}
2020
}
21+
func (b *lazyBuilder) Flatten() Builder {
22+
return &lazyBuilder{func() interface{} { return Flatten(b.exec()) }}
23+
}
2124
func (b *lazyBuilder) FlattenDeep() Builder {
2225
return &lazyBuilder{func() interface{} { return FlattenDeep(b.exec()) }}
2326
}
@@ -33,6 +36,9 @@ func (b *lazyBuilder) Join(rarr interface{}, fnc JoinFnc) Builder {
3336
func (b *lazyBuilder) Map(mapFunc interface{}) Builder {
3437
return &lazyBuilder{func() interface{} { return Map(b.exec(), mapFunc) }}
3538
}
39+
func (b *lazyBuilder) FlatMap(mapFunc interface{}) Builder {
40+
return &lazyBuilder{func() interface{} { return FlatMap(b.exec(), mapFunc) }}
41+
}
3642
func (b *lazyBuilder) Reverse() Builder {
3743
return &lazyBuilder{func() interface{} { return Reverse(b.exec()) }}
3844
}

‎lazy_builder_test.go

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,30 @@ func TestLazyFilter_SideEffect(t *testing.T) {
144144
is.NotEqual([]*foo{&foo{"foo"}, &foo{"bar"}}, in)
145145
}
146146

147+
func TestLazyFlatten(t *testing.T) {
148+
testCases := []struct {
149+
In interface{}
150+
}{
151+
{
152+
In: [][]int{{1, 2}, {3, 4}},
153+
},
154+
{
155+
In: [][][]int{{{1, 2}, {3, 4}}, {{5, 6}}},
156+
},
157+
}
158+
159+
for idx, tc := range testCases {
160+
t.Run(fmt.Sprintf("test case #%d", idx+1), func(t *testing.T) {
161+
is := assert.New(t)
162+
163+
expected := Flatten(tc.In)
164+
actual := LazyChain(tc.In).Flatten().Value()
165+
166+
is.Equal(expected, actual)
167+
})
168+
}
169+
}
170+
147171
func TestLazyFlattenDeep(t *testing.T) {
148172
testCases := []struct {
149173
In interface{}
@@ -261,6 +285,33 @@ func TestLazyMap(t *testing.T) {
261285
}
262286
}
263287

288+
func TestLazyFlatMap(t *testing.T) {
289+
testCases := []struct {
290+
In interface{}
291+
FlatMapFnc interface{}
292+
}{
293+
{
294+
In: [][]int{{1}, {2}, {3}, {4}},
295+
FlatMapFnc: func(x []int) []int { return x },
296+
},
297+
{
298+
In: map[string][]int{"Florent": {1}, "Gilles": {2}},
299+
FlatMapFnc: func(k string, v []int) []int { return v },
300+
},
301+
}
302+
303+
for idx, tc := range testCases {
304+
t.Run(fmt.Sprintf("test case #%d", idx+1), func(t *testing.T) {
305+
is := assert.New(t)
306+
307+
expected := Map(tc.In, tc.FlatMapFnc)
308+
actual := LazyChain(tc.In).Map(tc.FlatMapFnc).Value()
309+
310+
is.ElementsMatch(expected, actual)
311+
})
312+
}
313+
}
314+
264315
func TestLazyMap_SideEffect(t *testing.T) {
265316
is := assert.New(t)
266317

‎transform.go

Lines changed: 61 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ func ToMap(in interface{}, pivot string) interface{} {
9898
return collection.Interface()
9999
}
100100

101-
func mapSlice(arrValue reflect.Value, funcValue reflect.Value) interface{} {
101+
func mapSlice(arrValue reflect.Value, funcValue reflect.Value) reflect.Value {
102102
funcType := funcValue.Type()
103103

104104
if funcType.NumIn() != 1 || funcType.NumOut() == 0 || funcType.NumOut() > 2 {
@@ -125,7 +125,7 @@ func mapSlice(arrValue reflect.Value, funcValue reflect.Value) interface{} {
125125
resultSlice = reflect.Append(resultSlice, result)
126126
}
127127

128-
return resultSlice.Interface()
128+
return resultSlice
129129
}
130130

131131
if funcType.NumOut() == 2 {
@@ -141,13 +141,13 @@ func mapSlice(arrValue reflect.Value, funcValue reflect.Value) interface{} {
141141
collection.SetMapIndex(results[0], results[1])
142142
}
143143

144-
return collection.Interface()
144+
return collection
145145
}
146146

147-
return nil
147+
return reflect.Value{}
148148
}
149149

150-
func mapMap(arrValue reflect.Value, funcValue reflect.Value) interface{} {
150+
func mapMap(arrValue reflect.Value, funcValue reflect.Value) reflect.Value {
151151
funcType := funcValue.Type()
152152

153153
if funcType.NumIn() != 2 || funcType.NumOut() == 0 || funcType.NumOut() > 2 {
@@ -170,7 +170,7 @@ func mapMap(arrValue reflect.Value, funcValue reflect.Value) interface{} {
170170
resultSlice = reflect.Append(resultSlice, result)
171171
}
172172

173-
return resultSlice.Interface()
173+
return resultSlice
174174
}
175175

176176
// two parameters, should be a map
@@ -188,14 +188,24 @@ func mapMap(arrValue reflect.Value, funcValue reflect.Value) interface{} {
188188

189189
}
190190

191-
return collection.Interface()
191+
return collection
192192
}
193193

194-
return nil
194+
return reflect.Value{}
195195
}
196196

197197
// Map manipulates an iteratee and transforms it to another type.
198198
func Map(arr interface{}, mapFunc interface{}) interface{} {
199+
result := mapFn(arr, mapFunc, "Map")
200+
201+
if result.IsValid() {
202+
return result.Interface()
203+
}
204+
205+
return nil
206+
}
207+
208+
func mapFn(arr interface{}, mapFunc interface{}, funcName string) reflect.Value {
199209
if !IsIteratee(arr) {
200210
panic("First parameter must be an iteratee")
201211
}
@@ -214,13 +224,50 @@ func Map(arr interface{}, mapFunc interface{}) interface{} {
214224

215225
if kind == reflect.Slice || kind == reflect.Array {
216226
return mapSlice(arrValue, funcValue)
227+
} else if kind == reflect.Map {
228+
return mapMap(arrValue, funcValue)
217229
}
218230

219-
if kind == reflect.Map {
220-
return mapMap(arrValue, funcValue)
231+
panic(fmt.Sprintf("Type %s is not supported by "+funcName, arrType.String()))
232+
}
233+
234+
// FlatMap manipulates an iteratee and transforms it to a flattened collection of another type.
235+
func FlatMap(arr interface{}, mapFunc interface{}) interface{} {
236+
result := mapFn(arr, mapFunc, "FlatMap")
237+
238+
if result.IsValid() {
239+
return flatten(result).Interface()
240+
}
241+
242+
return nil
243+
}
244+
245+
// Flatten flattens a two-dimensional array.
246+
func Flatten(out interface{}) interface{} {
247+
return flatten(reflect.ValueOf(out)).Interface()
248+
}
249+
250+
func flatten(value reflect.Value) reflect.Value {
251+
sliceType := value.Type()
252+
253+
if (value.Kind() != reflect.Slice && value.Kind() != reflect.Array) ||
254+
(sliceType.Elem().Kind() != reflect.Slice && sliceType.Elem().Kind() != reflect.Array) {
255+
panic("Argument must be an array or slice of at least two dimensions")
256+
}
257+
258+
resultSliceType := sliceType.Elem().Elem()
259+
260+
resultSlice := reflect.MakeSlice(reflect.SliceOf(resultSliceType), 0, 0)
261+
262+
length := value.Len()
263+
264+
for i := 0; i < length; i++ {
265+
item := value.Index(i)
266+
267+
resultSlice = reflect.AppendSlice(resultSlice, item)
221268
}
222269

223-
panic(fmt.Sprintf("Type %s is not supported by Map", arrType.String()))
270+
return resultSlice
224271
}
225272

226273
// FlattenDeep recursively flattens array.
@@ -233,18 +280,18 @@ func flattenDeep(value reflect.Value) reflect.Value {
233280

234281
resultSlice := reflect.MakeSlice(reflect.SliceOf(sliceType), 0, 0)
235282

236-
return flatten(value, resultSlice)
283+
return flattenRecursive(value, resultSlice)
237284
}
238285

239-
func flatten(value reflect.Value, result reflect.Value) reflect.Value {
286+
func flattenRecursive(value reflect.Value, result reflect.Value) reflect.Value {
240287
length := value.Len()
241288

242289
for i := 0; i < length; i++ {
243290
item := value.Index(i)
244291
kind := item.Kind()
245292

246293
if kind == reflect.Slice || kind == reflect.Array {
247-
result = flatten(item, result)
294+
result = flattenRecursive(item, result)
248295
} else {
249296
result = reflect.Append(result, item)
250297
}

0 commit comments

Comments
 (0)