Flatten a list of lists in Python and Go

flatten.go

package main

import (
	"encoding/json"
	"fmt"
)

func flatten(s []interface{}) (r []int) {
	// type switch already discussed in the tour of go
	for _, item := range s {
		switch item := item.(type) {
		// json.Unmarshal reads numbers as float64
		case float64:
			r = append(r, int(item))
		case []interface{}:
			r = append(r, flatten(item)...)
		}
	}
	return
}

func main() {
	// declaring n dimensional lists (slices) in go isn't easy
	// stuffing them in json is a work around
	var parsedList []interface{}
	list2d := `[[111,222,333],[444,555,[666]], [777,888,999]]`
	err := json.Unmarshal([]byte(list2d), &parsedList)
	if err != nil {
		fmt.Println(err)
	}

	fmt.Println(flatten(parsedList))
}

flatten.py

# Python is still my first go to language for such a task
# There are many ways to reach the goal of this task.
# If I would not have to consider readability I would implement
# this task in a more functional style. However, I came to realize
# that my *personal* love of FP, isn't always compatible with what
# most people learn at university and practice daily (OOP).
# Hence, using Generators in Python is a good middle between achieving
# good readability of the code and good performance.
#
# i.e the following SO solution
#
# flatten = lambda lst: [lst] if type(lst) is int else reduce(operator.concat,
#                                                             [flatten(ele) for ele in lst])
#
# super cool oneline, most Python beginners and even experiecned Python
# programmers will find it hard to understand without first reading the
# documentation for operator.concat and reduce.
# Also, this one-liner  does not handle exceptions, see below a "safe version".



def flatten(_list):
    """
    A version of flatten which doesn't handle non-iterable types
    """
    for item in _list:
        if isinstance(item, (list, tuple)):
            for elem in flatten(item):
                yield elem
        else:
            yield item


def safe_flatten(_list):
    """
    A safe version of flatten which doesn't explode on weired inputs
    """
    try:
        for item in _list:
            if isinstance(item, (list, tuple)):
                for elem in safe_flatten(item):
                    yield elem
            else:
                yield item
    except TypeError:
        yield _list


# a full test suite with pytest or unittest is not required here
if __name__ == '__main__':  
    # some test data
    test_data = (
    ([[1, 2, [3]], 4], [1, 2, 3, 4]),
    ([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], list(range(1, 11))),
    ([[1], [[2, 3]], [[4, [5], [6, 7, [8, [9, [10], ]]]], ]], list(range(1, 11))),
    ([1, ], [1, ]),
    (1, [1, ]),
    (None, [None])  # a weired test case
    )

    for test_pair in test_data:
        try:
            result = list(flatten(test_pair[0]))
        except TypeError:
            continue

        try:
            assert result == test_pair[1]

        except AssertionError:
            print("Want: ", result, "Got: ", test_pair[1])            
            break

    for test_pair in test_data:
        result = list(safe_flatten(test_pair[0]))

        try:
            assert result == test_pair[1]
        except AssertionError:
            print("Want: ", result, "Got: ", test_pair[1])
            break

flatten_test.go

package main

import (
	"encoding/json"
	"fmt"
	"reflect"
	"testing"
)

func TestSimpleList(t *testing.T) {
	var parsedList []interface{}
	list := `[1,2,3,4,5]`
	//list2d := `[[111,222,333],[444,555,666], [777,888,999]]`
	//listNd := `[[111,222,333],[444,555,666], [777,888,999,[1,2,3,[4,[5,6,[7]]]]]]`
	err := json.Unmarshal([]byte(list), &parsedList)

	if err != nil {
		fmt.Println(err)
	}

	got := flatten(parsedList)
	want := []int{1, 2, 3, 4, 5}

	if !reflect.DeepEqual(want, got) {
		t.Fatalf("Wanted %v, got %v", want, got)
	}

}

func Test2dList(t *testing.T) {
	var parsedList []interface{}
	list2d := `[[111,222,333],[444,555,666], [777,888,999]]`
	//listNd := `[[111,222,333],[444,555,666], [777,888,999,[1,2,3,[4,[5,6,[7]]]]]]`
	err := json.Unmarshal([]byte(list2d), &parsedList)

	if err != nil {
		fmt.Println(err)
	}

	got := flatten(parsedList)
	want := []int{111, 222, 333, 444, 555, 666, 777, 888, 999}

	if !reflect.DeepEqual(want, got) {
		t.Fatalf("Wanted %v, got %v", want, got)
	}

}

func TestNdList(t *testing.T) {
	var parsedList []interface{}
	listNd := `[[1,2,3],[4,5,6], [7,8,9,[10,11,12,[13,[14,15,[16]]]]]]`
	err := json.Unmarshal([]byte(listNd), &parsedList)

	if err != nil {
		fmt.Println(err)
	}

	got := flatten(parsedList)
	want := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}

	if !reflect.DeepEqual(want, got) {
		t.Fatalf("Wanted %v, got %v", want, got)
	}
}