summaryrefslogtreecommitdiff
path: root/modules/language/python/module/itertools.scm
blob: d333ee99d9bfd7ebaf8acc9c8309daff1526568f (about) (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
(define-module (language python module itertools)
  #:use-module (language python for)
  #:use-module (language python yield)
  #:use-module (language python def)
  #:use-module (language python module copy)
  #:use-module (language python module python)
  #:use-module (language python module copy)
  #:export (count cycle repeat accumulate chain compress dropwhile
		  filterfalse groupby isslice starmap takewhile
		  tee zip_longest))

(define count
  (make-generator (start #:optional (step 1))
    (lambda* (yield start #:optional (step 1))
      (let lp ((i start))
	(yield i)
	(lp (+ i step))))))

(define cycle
  (make-generator
   (lambda (yield p)
     (let lp ()
       (for ((x : p)) () (yield x))
       (lp)))))

(define repeat
  (make-generator
   (lambda* (yield e #:optional (n -1))
     (let lp ((i 0))
       (if (not (= i n))
	   (begin
	     (yield e)
	     (lp (+ i 1))))))))

(define accumulate
  (make-generator
   (lambda* (yield p #:optional (f +))
     (for ((x : p)) ((s 0) (first? #t))
	  (if first?
	      (begin
		(yield x)
		(values x #f))
	      (let ((s (f x s)))
		(yield s)
		(values s #f)))))))

(define-python-class chain ()
  (define __call__
    (make-generator
     (lambda (yield . l)
       (let lp ((l l))
	 (if (pair? l)
	     (begin
	       (for ((x : (car l))) ()
		    (yield x))
	       (lp (cdr l))))))))
  (define from_iterable
    (make-generator
     (lambda (yield i)
       (for ((ii : i)) ()
	    (for ((x : ii)) ()
		 (yield x)))))))

(define compress
  (make-generator
   (lambda (yield data selectors)
     (for ((d : data) (s : selectors)) ()
	  (if s (yield d))))))

(define dropwhile
  (make-generator
   (lambda (yield pred seq)
     (for ((x : seq)) ((start? #f))
	  (if start?
	      (begin
		(yield x)
		#t)
	      (if (pred x)
		  #f
		  (begin
		    (yield x)
		    #t)))))))

(define filterfalse
  (make-generator
   (lambda (yield pred seq)
     (for ((x : seq))
	  (if (not (f x)) (yield x))))))

(define none (list 'none))
(define groupby
  (make-generator
   (lambda* (yield seq #:optional (key (lambda (x) x)))
     (for ((x : seq)) ((k none)) ((l '()))
	  (if (eq? k none)
	      (values (key x) (list x))
	      (let ((kk (key x)))
		(if (equal? k kk)
		    (values k (cons x l))
		    (begin
		      (yield k (reverse l))
		      (values kk (list x))))))
	  #:final
	  (if (not (eq? k none))
	      (yield l (reverse l)))))))
     
     
(define isslice
  (make-generator
   (lambda* (yield seq #:optional (start 0) (stop -1) (step 1))
     (for ((x : seq) (i : (count 0)))
	  (if (= i stop) (break))
	  (if (and (>= i start)
		   (= (modulo (- i start) step) 0))
	      (yield x))))))

(define starmap
  (make-generator
   (lambda (yield f seq)
     (for ((x : seq)) () (yield (f x))))))

(define takewhile
  (make-generator
   (lambda (yield pred seq)
     (for ((x : seq)) ()
	  (if (not (pred x)) (break))
	  (yield x)))))
  
(define tee
  (make-generator
   (lambda (yield it n)
     (let lp ((i 0))
       (if (< i n)
	   (cons (deepcopy it)
		 (lp (+ i 1)))
	   '())))))

(define zip_longest
  (make-generator
   (lam (yield (* l) (= fillvalue None))
	(define mkit
	  (make-generator
	   (lambda (yield it)
	     (for ((x : it)) ()
		  (yield (cons 1 x))
	     (let lp ()
	       (yield (cons 0 0))
	       (lp)))))
	  (for ((x : (apply zip (map mkit l))))
	       (if (= (apply + (map car x)) 0)
		   (break)
		   (yield (map (lambda (x) (if (= (car x) 0) fillvalue (cdr x)))
			       x))))))))