github.com/songzhibin97/go-baseutils@v0.0.2-0.20240302024150-487d8ce9c082/app/bconcurrent/stream.go (about)

     1  package bconcurrent
     2  
     3  import "context"
     4  
     5  func Stream[T any](ctx context.Context, values ...T) <-chan T {
     6  	out := make(chan T)
     7  	go func() {
     8  		defer close(out)
     9  		for _, value := range values {
    10  			select {
    11  			case <-ctx.Done():
    12  				return
    13  			case out <- value:
    14  			}
    15  		}
    16  	}()
    17  	return out
    18  }
    19  
    20  // TaskN 只取流中的前N个数据
    21  func TaskN[T any](ctx context.Context, valueStream <-chan T, num int) <-chan T {
    22  	outStream := make(chan T)
    23  	go func() {
    24  		defer close(outStream)
    25  		for i := 0; i < num; i++ {
    26  			select {
    27  			case <-ctx.Done():
    28  				return
    29  			case v, ok := <-valueStream:
    30  				if !ok {
    31  					return
    32  				}
    33  				select {
    34  				case <-ctx.Done():
    35  					return
    36  				case outStream <- v:
    37  
    38  				}
    39  			}
    40  		}
    41  	}()
    42  	return outStream
    43  }
    44  
    45  // TaskFn 筛选流中的数据,只保留满足条件的数据
    46  func TaskFn[T any](ctx context.Context, valueStream <-chan T, fn func(v T) bool) <-chan T {
    47  	outStream := make(chan T)
    48  	go func() {
    49  		defer close(outStream)
    50  		for {
    51  			select {
    52  			case <-ctx.Done():
    53  				return
    54  			case v, ok := <-valueStream:
    55  				if !ok {
    56  					return
    57  				}
    58  				if fn(v) {
    59  					select {
    60  					case <-ctx.Done():
    61  						return
    62  					case outStream <- v:
    63  					}
    64  				}
    65  			}
    66  		}
    67  	}()
    68  	return outStream
    69  }
    70  
    71  // TaskWhile 只取满足条件的数据,一旦不满足就不再取
    72  func TaskWhile[T any](ctx context.Context, valueStream <-chan T, fn func(v T) bool) <-chan T {
    73  	outStream := make(chan T)
    74  	go func() {
    75  		defer close(outStream)
    76  		for {
    77  			select {
    78  			case <-ctx.Done():
    79  				return
    80  			case v, ok := <-valueStream:
    81  				if !ok {
    82  					return
    83  				}
    84  				if fn(v) {
    85  					select {
    86  					case <-ctx.Done():
    87  						return
    88  					case outStream <- v:
    89  					}
    90  					return
    91  				}
    92  			}
    93  		}
    94  	}()
    95  	return outStream
    96  }
    97  
    98  // SkipN 跳过流中的前N个数据
    99  func SkipN[T any](ctx context.Context, valueStream <-chan T, num int) <-chan T {
   100  	outStream := make(chan T)
   101  	go func() {
   102  		defer close(outStream)
   103  		for i := 0; i < num; i++ {
   104  			select {
   105  			case <-ctx.Done():
   106  				return
   107  			case _, ok := <-valueStream:
   108  				if !ok {
   109  					return
   110  				}
   111  			}
   112  		}
   113  		for {
   114  			select {
   115  			case <-ctx.Done():
   116  				return
   117  			case v, ok := <-valueStream:
   118  				if !ok {
   119  					return
   120  				}
   121  				select {
   122  				case <-ctx.Done():
   123  					return
   124  				case outStream <- v:
   125  				}
   126  			}
   127  		}
   128  	}()
   129  	return outStream
   130  }
   131  
   132  // SkipFn 跳过满足条件的数据
   133  func SkipFn[T any](ctx context.Context, valueStream <-chan T, fn func(v T) bool) <-chan T {
   134  	outStream := make(chan T)
   135  	go func() {
   136  		defer close(outStream)
   137  		for {
   138  			select {
   139  			case <-ctx.Done():
   140  				return
   141  			case v, ok := <-valueStream:
   142  				if !ok {
   143  					return
   144  				}
   145  				if !fn(v) {
   146  					select {
   147  					case <-ctx.Done():
   148  						return
   149  					case outStream <- v:
   150  					}
   151  				}
   152  			}
   153  		}
   154  	}()
   155  	return outStream
   156  }
   157  
   158  // SkipWhile 跳过满足条件的数据,一旦不满足,当前这个元素以后的元素都会输出
   159  func SkipWhile[T any](ctx context.Context, valueStream <-chan T, fn func(v T) bool) <-chan T {
   160  	outStream := make(chan T)
   161  	go func() {
   162  		defer close(outStream)
   163  		for {
   164  			select {
   165  			case <-ctx.Done():
   166  				return
   167  			case v, ok := <-valueStream:
   168  				if !ok {
   169  					return
   170  				}
   171  				if fn(v) {
   172  					select {
   173  					case <-ctx.Done():
   174  						return
   175  					default:
   176  
   177  					}
   178  				} else {
   179  					select {
   180  					case <-ctx.Done():
   181  						return
   182  					case outStream <- v:
   183  					}
   184  					for {
   185  						select {
   186  						case <-ctx.Done():
   187  							return
   188  						case v, ok = <-valueStream:
   189  							if !ok {
   190  								return
   191  							}
   192  							select {
   193  							case <-ctx.Done():
   194  								return
   195  							case outStream <- v:
   196  							}
   197  						}
   198  					}
   199  				}
   200  			}
   201  		}
   202  	}()
   203  	return outStream
   204  }