Skip to content

Commit

Permalink
Support shuffle for iter
Browse files Browse the repository at this point in the history
  • Loading branch information
ringsaturn committed Aug 25, 2024
1 parent 56c2927 commit 5b89987
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 14 deletions.
15 changes: 10 additions & 5 deletions cities.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,14 @@ import (
"strconv"
)

//go:embed upstream/cities.json
var citiesbytes []byte
var (
//go:embed upstream/cities.json
citiesbytes []byte

Cities []*City

idx []int
)

type City struct {
Country string `json:"country"`
Expand All @@ -19,8 +25,6 @@ type City struct {
Admin2 string `json:"admin2"`
}

var Cities []*City

func init() {
type RawCity struct {
Country string `json:"country"`
Expand All @@ -35,7 +39,7 @@ func init() {
if err := json.Unmarshal(citiesbytes, &rawcities); err != nil {
panic(err)
}
for _, rawcity := range rawcities {
for i, rawcity := range rawcities {
lat, err := strconv.ParseFloat(rawcity.Lat, 64)
if err != nil {
panic(err)
Expand All @@ -52,5 +56,6 @@ func init() {
Admin1: rawcity.Admin1,
Admin2: rawcity.Admin2,
})
idx = append(idx, i)
}
}
25 changes: 18 additions & 7 deletions iter.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,33 @@

package gocitiesjson

import "iter"
import (
"iter"
"math/rand"
)

func All() iter.Seq[*City] {
func All(shuffle bool) iter.Seq[*City] {
return func(yield func(*City) bool) {
for _, city := range Cities {
if !yield(city) {
var idxes []int
if shuffle {
idxes = make([]int, len(idx))
copy(idxes, idx)
rand.Shuffle(len(idxes), func(i, j int) { idxes[i], idxes[j] = idxes[j], idxes[i] })
} else {
idxes = idx
}
for _, i := range idxes {
if !yield(Cities[i]) {
return
}
}
}
}

func Filter(filter func(*City) bool) iter.Seq[*City] {
func Filter(fn func(*City) bool, shuffle bool) iter.Seq[*City] {
return func(yield func(*City) bool) {
for _, city := range Cities {
if filter(city) {
for city := range All(shuffle) {
if fn(city) {
if !yield(city) {
return
}
Expand Down
4 changes: 2 additions & 2 deletions iter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (

func ExampleAll() {
c := 0
for city := range gocitiesjson.All() {
for city := range gocitiesjson.All(false) {
_ = city
c++
}
Expand All @@ -25,7 +25,7 @@ func ExampleFilter() {
bboxFilter := func(city *gocitiesjson.City) bool {
return city.Lng > 0 && city.Lng < 10 && city.Lat > 0 && city.Lat < 10
}
for city := range gocitiesjson.Filter(bboxFilter) {
for city := range gocitiesjson.Filter(bboxFilter, false) {
_ = city
c++
}
Expand Down

0 comments on commit 5b89987

Please # to comment.