heapq added

3 files changed, 268 insertions, 17 deletions

diff --git a/modules/language/python/module/collections.scm b/modules/language/python/module/collections.scm
index bfb75fa..6e004a3 100644
--- a/modules/language/python/module/collections.scm
+++ b/modules/language/python/module/collections.scm
@@ -1,8 +1,19 @@
 (define-module (language python module collections)
   #:use-module (oop pf-objects)
   #:use-module (language python module collections abc)
-  #:export (ChainMap))
+  #:export (abc ChainMap))
 
+(define-python-class class-from-dict ()
+  (define __init__
+    (lambda (self d)
+      (for ((k v : d))
+	   (set self k v)))))
+
+(define abc
+  (class-from-dict
+   (resolve-module (module-public-interface
+		    '(language python module collections abc)))))
+  
 (define (u self)
   (let ((s (set)))
     (apply (ref s 'union) (ref self 'maps)))
@@ -143,14 +154,12 @@
             (sorted ((ref self 'items) #:key (_itemgetter 1) #:reverse #t)
 	    _heapq.nlargest(n ((ref self 'items)) #:key (_itemgetter 1))
 
-
+	    
 	       
 
 
   
 
-(define* (namedtuple typename field-names #key (verbose #f) (rename=#f)
-(module None))
   
   
 
diff --git a/modules/language/python/module/heapq.scm b/modules/language/python/module/heapq.scm
new file mode 100644
index 0000000..768aee2
--- /dev/null
+++ b/modules/language/python/module/heapq.scm
@@ -0,0 +1,236 @@
+(define-module (language python module heapq)
+  #:use-module (ice-9 control)
+  #:use-module (language python for)
+  #:use-module (language python yield)
+  #:use-module (language python try)
+  #:use-module (language python exceptions)
+  #:use-module (language python list)
+  #:use-module (language python def)
+  #:use-module (oop pf-objects)
+  #:use-module ((language python module python)
+		#:select (zip sorted range enumerate iter max min))
+  #:export (heappush heappop heapify heapreplace merge nlargest nsmallest
+		     heappushpop))
+
+(define-syntax-rule (aif it p a b) (let ((it p)) (if it a b)))
+
+(define (heappush heap item)
+  (pylist-append! heap item)
+  (siftdown heap 0 (- (len heap) 1)))
+
+(define (heappop heap)
+  (let ((lastelt (pylist-pop! heap)))
+    (if (> (len heap) 0)
+	(let ((ret (pylist-ref heap 0)))
+	  (pylist-set! heap 0 lastelt)
+	  (siftup heap 0)
+	  ret)
+	lastelt)))
+
+(define (heapreplace heap item)
+  (let ((ret (pylist-ref heap 0)))
+    (pylist-set! heap 0 item)
+    (siftup heap 0)
+    ret))
+
+(define (heapify x)
+  (let lp ((i (- (floor-quotient (len x) 2) 1)))
+    (if (>= i 0)
+	(begin
+	  (siftup x i)
+	  (lp (- i 1))))))
+
+(define (heappop_max heap)
+  (let ((lastelt (pylist-pop! heap)))
+    (if (> (len heap) 0)
+	(let ((ret (pylist-ref heap 0)))
+	  (pylist-set! heap 0 lastelt)
+	  (siftup_max heap 0)
+	  ret)
+	lastelt)))
+
+(define (heapreplace_max heap item)
+  (let ((ret (pylist-ref heap 0)))
+    (pylist-set! heap 0 item)
+    (siftup_max heap 0)
+    ret))	
+
+(define (heapify_max x)
+  (let lp ((i (- (floor-quotient (len x) 2) 1)))
+    (if (>= i 0)
+	(begin
+	  (siftup_max x i)
+	  (lp (- i 1))))))
+
+(define (siftdown heap startpos pos)
+  (define newitem (pylist-ref heap pos))
+  
+  (let lp ((pos pos))
+    (if (> pos startpos)
+	(let ()
+	  (define parentpos (ash (- pos 1) 1))
+	  (define parent    (pylist-ref heap parentpos))
+	  (if (< newitem parent)
+	      (begin
+		(pylist-set! heap pos parent)
+		(lp parentpos))
+	      (pylist-set! heap pos newitem)))
+	(pylist-set! heap pos newitem))))
+
+(define (siftup heap pos)
+  (define endpos   (len heap))
+  (define startpos pos)
+  (define newitem  (pylist-ref heap pos))
+  (let lp ((pos pos) (childpos (+ (* 2 pos) 1)))
+    (if (< childpos endpos)
+	(let* ((rightpos (+ childpos 1))
+	       (childpos
+		(if (and (< rightpos endpos)
+			 (not (< (pylist-ref heap childpos)
+				 (pylist-ref heap rightpos))))
+		    rightpos
+		    childpos)))
+	  (pylist-set! heap pos (pylist-ref heap childpos))
+	  (lp childpos (+ (* 2 childpos) 1)))
+	(begin
+	  (pylist-set! heap pos newitem)
+	  (siftdown heap startpos pos)))))
+
+(define (siftdown_max heap startpos pos)
+  (define newitem (pylist-ref heap pos))
+  
+  (let lp ((pos pos))
+    (if (> pos startpos)
+	(let ()
+	  (define parentpos (ash (- pos 1) 1))
+	  (define parent    (pylist-ref heap parentpos))
+	  (if (< parent newitem)
+	      (begin
+		(pylist-set! heap pos parent)
+		(lp parentpos))
+	      (pylist-set! heap pos newitem)))
+	(pylist-set! heap pos newitem))))
+
+(define (siftup_max heap pos)
+  (define endpos   (len heap))
+  (define startpos pos)
+  (define newitem  (pylist-ref heap pos))
+  (let lp ((pos pos) (childpos (+ (* 2 pos) 1)))
+    (if (< childpos endpos)
+	(let* ((rightpos (+ childpos 1))
+	       (childpos
+		(if (and (< rightpos endpos)
+			 (not (< (pylist-ref heap rightpos)
+				 (pylist-ref heap childpos))))
+		    rightpos
+		    childpos)))
+	  (pylist-set! heap pos (pylist-ref heap childpos))
+	  (lp childpos (+ (* 2 childpos) 1)))
+	(begin
+	  (pylist-set! heap pos newitem)
+	  (siftdown_max heap startpos pos)))))
+
+(def (merge (* iterables) (= key (lambda (x) x)) (= reverse #f))
+    ((make-generator ()
+       (lambda (yield)
+	 (define h (py-list))
+	 (define (h_append x) (pylist-append! h x))
+	 (define-values (-heapify -heappop -heapreplace direction)
+	   (if reverse
+	       (values heapify_max
+		       heappop_max
+		       heapreplace_max
+		       -1)
+	       (values heapify
+		       heappop
+		       heapreplace
+		       1)))
+	   
+       	 (for ((order it : (enumerate (map iter iterables)))) ()
+	      (define next- (aif x (ref it '__next__)
+				 x
+				 (lambda () (next it))))
+	      (try
+	       (lambda ()
+		 (h_append (vector (key (next-))
+				   (* order direction)
+				   next-)))
+	       (#:except StopIteration =>
+			 (lambda x (-heappop h)))))
+
+	 (-heapify h)
+	 (let lp ()
+	   (if (> (len h) 1)
+	       (try
+		(lambda ()
+		  (let lp ()
+		    (let ((s (pylist-ref h 0)))
+		      (yield (vector-ref s 0))
+		      (vector-set! s 0 (key ((vector-ref s 2))))
+		      (-heapreplace h s)
+		      (lp))))
+		(#:except StopIteration =>
+			  (lambda x (-heappop h) (lp))))))
+
+	 (if (> (len h) 0)
+	     (let* ((x (pylist-ref h 0))
+		    (n (pylist-ref x 2)))
+	       (let lp ((i (vector-ref x 0)))
+		 (yield i)
+		 (lp (n)))))))))
+		 
+
+
+	 
+(define-syntax-rule (mkn nsmallest heapify_max heapreplace_max q< ran o ++)
+  (define* (nsmallest n iterable #:key (key (lambda (x) x)))
+    (let/ec return
+      (if (equal? n 1)
+	  (return
+	   (let* ((sent (list 1))
+		  (res  (min (iter iterable) #:default sent #:key key)))
+	     (if (eq? sent res)
+		 '()
+		 (list res)))))
+
+      (let ((size (try
+		   (lambda () (len iterable))
+		   (#:except (or TypeError AttributeError) =>
+			     (lambda x #f)))))
+	(if (and size (>= n size))
+	    (return (pylist-slice (sorted iterable #:key key) 0
+				  ((@ (guile) min) n size)
+				  1))))
+  
+      (let ((it     (iter iterable))
+	    (result (py-list)))
+    
+	(for ((i elem : (zip (ran n) it))) ()
+	     (pylist-append! result (vector (key elem) i elem)))
+
+	(if (= 0 (len result))
+	    (return result))
+
+	(heapify_max result)
+
+	(for ((elem : it)) ((top   (vector-ref (pylist-ref result 0) 0))
+			    (order (o n)))
+	     (let ((k (key elem)))
+	       (if (q< k top)
+		   (begin
+		     (heapreplace_max result (vector k order elem))
+		     (values (vector-ref (pylist-ref result 0) 0)
+			     (++ order)))
+		   (values top order))))
+	(pylist-sort! result)
+	(return
+	 (let ((r (py-list)))
+	   (for ((x : result)) ()
+		(pylist-append! r (vector-ref x 2)))
+	   r))))))
+
+(mkn nsmallest heapify_max heapreplace_max < (lambda (n) (range n))
+     (lambda (n) n) (lambda (x) (+ x 1)))
+
+(mkn nlargest  heapify     heapreplace     > (lambda (n) (range 0 (- n) -1))
+     (lambda (n) (- n)) (lambda (x) (- x 1)))
diff --git a/modules/language/python/module/python.scm b/modules/language/python/module/python.scm
index 2ea57eb..8264fee 100644
--- a/modules/language/python/module/python.scm
+++ b/modules/language/python/module/python.scm
@@ -85,18 +85,18 @@
 (define staticmethod static-method)
 
 (define (enumerate l)
-  (make-generator enumerate
-     (lambda (yield)
+  ((make-generator ()
+    (lambda (yield)
        (for ((x : l)) ((i 0))
             (yield i x)
-           (+ i 1)))))
+           (+ i 1))))))
 
 (define (filter f l)
-  (make-generator enumerate
-    (lambda (yield)
-      (for ((x : l)) ()
-           (if (f x)
-               (yield x))))))
+  ((make-generator ()
+     (lambda (yield)
+       (for ((x : l)) ()
+            (if (f x)
+                (yield x)))))))
 
 (define miss ((@ (guile) list) 'miss))
 
@@ -174,7 +174,7 @@
      (let lp ((l l))
        (match l
          ((it)
-          (for ((x : it)) ((s default) (b default))
+          (for ((x : it)) ((s miss) (b miss))
                (if (eq? s miss)
                    (values (key x) x)
                    (let ((k (key x)))
@@ -183,7 +183,10 @@
                          (values s b))))
                #:final
                (if (eq? b miss)
-                   (raise ValueError "min does not work for zero length list")
+		   (if (eq? default miss)
+		       (raise ValueError
+			      "min does not work for zero length list")
+		       default)
                    b)))
          (_ (lp ((@ (guile) list) l))))))
 
@@ -191,8 +194,8 @@
      (let lp ((l l))
        (match l
          ((it)
-          (for ((x : it)) ((s default) (b default))
-               (if (eq? default miss)
+          (for ((x : it)) ((s miss) (b miss))
+               (if (eq? s miss)
                    (values (key x) x)
                    (let ((k (key x)))
                      (if (> k s)
@@ -200,7 +203,10 @@
                          (values s b))))
                #:final
                (if (eq? b miss)
-                   (raise ValueError "min does not work for zero length list")
+		   (if (eq? default miss)
+		       (raise ValueError
+			      "min does not work for zero length list")
+		       default)
                    b)))
          (_ (lp ((@ (guile) list) l))))))