path: root/modules/language/python/module/collections.scm

(define-module (language python module collections)
  #:use-module (oop pf-objects)
  #:use-module (language python module collections abc)
  #:use-module (language python module heapq)
  #:use-module (language python for)
  #:use-module (language python yield)
  #:use-module (language python def)
  #:use-module (language python list)
  #:use-module (language python string)
  #:use-module (language python dict)
  #:use-module (language python exceptions)
  #:use-module ((language python module python)
		#:select ((map . pymap)))
 
  #:export (abc OrderedDict ChainMap Counter UserDict UserString UserList))

(define-python-class class-from-dict ()
  (define __init__
    (lambda (self d)
      (set self '__dict__ d))))

(define abc
  (class-from-dict
   (resolve-module (module-public-interface
		    '(language python module collections abc)))))

(define-python-class _OrderedDictKeysView (KeysView)
  (define __reversed__
    (lambda (self)
      ((make-generator ()
         (lambda (yield)
	   (for ((k v : (reversed (reg self '_mapping))))
		(yield k))))))))

(define-python-class _OrderedDictValuesView (ValuesView)
  (define __reversed__
    (lambda (self)
      ((make-generator ()
         (lambda (yield)
	   (for ((k v : (reversed (reg self '_mapping))))
		(yield v))))))))

(define-python-class _OrderedDictItemsView (ItemsView)
  (define __reversed__
    (lambda (self)
      ((make-generator ()
         (lambda (yield)
	   (for ((k v : (reversed (reg self '_mapping))))
		(yield (list k v)))))))))

(define-inlinable (link) (vector 0 0 0))
(define-inlinable (get-prev l)    (vector-ref  l 0))
(define-inlinable (get-next l)    (vector-ref  l 1))
(define-inlinable (get-key  l)    (vector-ref  l 2))
(define-inlinable (set-prev! l v) (vector-set! l 0 v))
(define-inlinable (set-next! l v) (vector-set! l 1 v))
(define-inlinable (set-key!  l v) (vector-set! l 2 v))

(define proxy (list 'mu))

(define <dict> (cons <py-hashtable> '_))
(define dict-set!   (resolve-method-c pylist-set!    <dict>))
(define dict-ref    (resolve-method-c pylist-ref     <dict>))
(define dict-del!   (resolve-method-c pylist-delete! <dict>))
(define dict-pop!   (resolve-method-c pylist-pop!    <dict>))
(define dict-clear! (resolve-method-c pylist-clear!  <dict>))

(define-python-class OrderedDict (dict)
  (define __init__
    (lam (self (* args) (** kwds))
	 (if (> (len args)  1)
	     (raise TypeError
		    (format #f
			    "expected at most 1 arguments, got %d"
			    (len args))))
	
	 (try
	  (lambda () (ref self '__root))
	  (#:except AttributeError =>
		    (lambda x
		      (let* ((l (link)))
			(set self '__root l)
			(set-next! l l)
			(set-prev! l l)))))
	
	 (set self '__map (dict))
	 (py-apply py-update self (* args) (** kwds))))

  (define __setitem__
    (lambda* (self key value #:key
		   (dict_setitem dict-set!)
		   (proxy        #f)
		   (link         link))
     (if (in key self)
	 (let* ((link (link))
		(root (ref self '__root))
		(last (get-last root)))
	   (pylist-set! (ref self '__map) key link)
	   (set-prev! link last)
	   (set-next! link root)
	   (set-key!  link key)
           (set-next! last link)
           (set-prev! root link)
	   (dict_setitem self key value)))))

  (define __delitem__
    (lambda* (self key #:key
		   (dict_delitem dict-set!)
		   (dict_delitem dict-del!))
      (let ((link (pylist-pop! (ref self '__map) key))
	    (link_prev (get-prev link))
	    (link_next (get-next link)))
	(set-next! link_prev link_next)
	(set-prev! link_next link_prev)
	(set-next! link None)
	(set-prev! link None))))
  
  (define __iter__
    (lambda* (self #:key (get-next get-next))
      ((make-generator ()
        (lambda (yield)
	  (let ((root (ref self '__root)))
	    (let lp ((curr (get-next root)))
	      (if ((not (eq? curr root)))
		  (let ((key (get-key curr)))
		    (yield key (pylist-ref self key))
		    (lp (get-next curr)))))))))))

  (define __reversed__
    (lambda (self)
      (__iter__ self #:get-next get-prev)))
  
  (define clear
    (lambda (self)
      (let ((root (ref self '__root)))
	(set-prev! root root)
	(set-next! root root)
	
        (py-clear (ref self '__map))
        (dict-clear! self))))

  (define popitem
    (lambda* (self #:key (last #t))
      (if (= (len self) 0)
	  (raise KeyError "dictionary is empty"))
      (let ((root (ref self '__root)))	
        (let* ((link
		(if last
		    (let* ((link (get-prev root))
			   (prev (get-prev link)))	      
		      (set-next! prev root)
		      (set-prev! root prev)
		      link)
		    (let* ((link (get-next root))
			   (prev (get-next link)))	      
		      (set-prev! prev root)
		      (set-next! root prev)
		      link)))
	       (key (get-key link)))
	  (dict-del! (ref self '__map) key)
	  (values key
		  (dict-pop! self key))))))
  
  (define move_to_end
    (lambda* (self key #:key (last #t))
      (let* ((link (pylist-ref (ref self '__map) key))
	     (prev (get-prev link))
	     (next (get-next link)))
        (set-next! prev next)
        (set-prev! next prev)
        (let ((root (ref self '__root)))
	  (if last
	      (let ((last (get-prev root)))
		(set-prev! link last)
		(set-next! link root)
		(set-prev! root link)
		(set-next! last link))
	      (let ((first (get-next root)))
		(set-prev! link  root)
		(set-next! link  first)
		(set-prev! first link)
		(set-next! root  link)))))))

  (define update   (ref MutableMapping 'update))
  (define __update update)

  (define keys
    (lambda (self) _OrderedDictKeysView(self)))
  (define items
    (lambda (self) _OrderedDictItemsView(self)))
  (define values
    (lambda (self) _OrderedDictValuesView(self)))

  (define __ne__ (ref MutableMapping '__ne__))

  (define __marker (object))

  (define pop
    (lambda* (self key #:key (default __marker))
      (if (in key self)
	  (let ((result (dict-ref self key)))
	    (__delitem__ self key)
	    result)  
	  (if (eq? default __marker)
	      (raise KeyError key)
	      default))))

  (define setdefault
    (lambda* (self key #:key (default None))
       (if (in key self)
	   (dict-ref self key)
	   (begin
	     (__setitem__ self key default)
	     default))))
  
  (define copy
    (lambda (self)
      ((ref self '__class__) self)))

  (define fromkeys
    (classmethod 
     (lambda* (cls iterable #key (value None))
      (let ((self (cls)))
	(for ((key : iterable)) ()
	     (__setitem__ self key value))
	self))))
  
  (define __eq__
    (lambda (self other)
      (if (isinstance other OrderedDict)
	  (and ((dict-equal? self other)
	       (all (map _equal self other)))
	  ((ref dict '__eq__) self other))))


(define (u self)
  (let ((s (set)))
    (apply (ref s 'union) (ref self 'maps)))
  s)

(define-python-class ChainMap (MutableMapping)
  (define __init__
    (lambda (self . l)
      (set self 'maps
	   (if (null? l)
	       (list (dict))
	       l))))

  (define __getitem__
    (lambda (self key)
      (let/ec ret
	 (let lp ((l (ref self 'maps)))
	   (if (pair? l)
	       (let ((m (car l)))
		 (if (in key m)
		     (pylist-ref m key)
		     (lp (cdr l))))
	       (raise (KeyError key)))))))
  

  (define get
    (lambda* (self key #:key (default None))
      (if (in key self)
	  (pylist-ref self key)
	  default)))

  (define __len__
    (lambda (self)
      (len (u self))))


  (define __iter__
    (lambda (self)
      (iter (u self))))

  (define __contains__
    (lambda (self, key)
      (for ((m : (ref self 'maps)))
	   (if (in key m)
	       (break #t))
	   #:final #f)))
  
  (define __bool__
    (lambda (self)
      (any (ref self 'maps))))

  (define __repr__
    (lambda (self)
      (format #f "~a(~a,~{,~a~})"
	      (ref (ref self '__class__) '__name__)
	      (map (lambda (x) (repr x)) (ref self 'maps)))))

  (define fromkeys
    (class-method
     (lambda (cls iter . args)
       (apply cls (for ((x : iter)) ((l '()))
		       (cons (apply py-fromkeys x args) l)
		       
		       #:final
		       (reverse l))))))

  (define copy
    (lambda (self)
      (let ((maps (ref self 'maps)))
	(apply (ref self '__class__) (copy (car maps)) (cdr maps)))))

  (define __copy__
    (lambda (self)
      ((ref self 'copy))))

  (define new_child
    (lambda* (self #:optional (r #f))
	     (apply (ref self '__class__) (if r r (set)) (ref self 'maps))))

  (define parents
    (property
     (lambda (self)
       (apply (ref self '__class__) (cdr (ref self 'maps))))))
  
  (define __setitem__
    (lambda (self key value)
      (pylist-set! (car (ref self 'maps)) key value)))
    

  (define __delitem__
    (lambda (self, key)
      (try
       (lambda () (pylist-del! (car (ref self 'maps))))
       #:except KeyError =>
       (lambda x
	 (raise KeyError
		(format #f "Key not found in the first mapping: ~a" key))))))

  (define popitem
    (lambda (self)
      (try
       (lambda () (popitem (car (ref self 'maps))))	 
       #:except KeyError =>
       (lambda x
	 (raise KeyError "No keys found in the first mapping")))))

  (define pop
    (lambda (self key . args)
      (try
       (lambda () (apply py-pop (car (ref self 'maps)) args))
       #:except KeyError =>
       (lambda ()
	 (raise KeyError
		(format #f "Key not found in the first mapping: ~ a" key))))))

  (define clear
    (lambda (self)
      (py-clear (car (ref self 'maps))))))

(define-python-class Counter (dict)
  (define __init__
    (lam (self (* args) (** kwds))
        (if (> (length args) 1)
            (raise TypeError
		   (format
		    #f
		    "expected at most 1 arguments, got ~ a"
		    (length  args))))
        ((ref (super Counter self) '__init__))
        (py-apply py-update self (* args) (** kwds))))

  (define __missing__
    (lambda (self key) 0))

  (define most_common
    (lambda* (self #:key (n None)):
       (if (eq? n None)
            (sorted ((ref self 'items) #:key (_itemgetter 1) #:reverse #t))
	    (nlargest n (py-items self) #:key (_itemgetter 1)))))

  (define elements
    (lambda (self)
      ((make-generator ()
	 (for ((k v : self)) ()
	      (if (and (number? v) (integer? v) (> v 0))
		  (for ((i : range(v))) ()
		       (yield k))))))))

  (define fromkeys
    (lambda x
      (raise NotImplementedError)))
    
  
  (define update
    (lam (self (* args) (** kwds))
	 (if (> (len args) 1)
	     (raise TypeError
		    (format #f "expected at most 1 arguments, got %d"
			    (len args))))
	 
        (let ((iterable (pylist-get args o None)))
	  (if (not (eq? iterable None))
	      (if (is-a? iterable <py-dict>)
		  (for ((elem count : iterable)) ()
		       (pylist-set! self elem
				    (+ count (pylist-get self elem 0))))
		  (for ((k : iterable)) ()
		       (pylist-set! self elem
				    (+ 1 (pylist-get self elem 0)))))))
        
	(for ((k v : kwds)) ()
	     (pylist-set! self k
			  (+ count (pylist-get self k 0))))))
  (define subtracts
    (lam (self (* args) (** kwds))
	 (if (> (len args) 1)
	     (raise TypeError
		    (format #f "expected at most 1 arguments, got %d"
			    (len args))))

	 (let ((iterable (pylist-get args 0 None)))
	   (if (not (eq? iterable None))
	       (if (is-a? iterable <py-dict>)
		   (for ((elem count : iterable)) ()
			(pylist-set! self elem
				     (- (pylist-get self elem 0) count)))
		   (for ((elem : iterable)) ()
			(pylist-set! self elem
				     (- (pylist-get self elem 0) 1))))))
	        
	(for ((k v : kwds)) ()
	     (pylist-set! self k
			  (- (pylist-get self k 0) v)))))

  (define __delitem__
    (lambda (self k)
      (if (in k self)
	  ((ref dict '__delitem__) k))))


  (define __add__
    (lambda (self other)
      (if (not (isinstance other Counter))
	  NotImplemented
	  (let ((result (Counter)))	    
	    (for ((elem count : self)) ()
		 (let ((newcount (+ count (pylist-ref other elem))))
		   (if (> newcount 0)
		       (pylist-set! result elem newcount))))
	    
	    (for ((elem count : other))
		 (if (and (not (in elem self))
			  (> count 0))
		     (pylist-set! result elem count)))
	    
	    return))))

  (define __sub__
    (lambda (self other)
      (if (not (isinstance other Counter))
	  NotImplemented
	  (let ((result (Counter)))	    
	    (for ((elem count : self)) ()
		 (let ((newcount (- count (pylist-ref other elem))))
		   (if (> newcount 0)
		       (pylist-set! result elem newcount))))
	    
	    (for ((elem count : other))
		 (if (and (not (in elem self))
			  (> count 0))
		     (pylist-set! result elem (- count))))
	    
	    return))))

  (define __or__
    (lambda (self other)
      (if (not (isinstance other Counter))
	  NotImplemented
	  (let ((result (Counter)))	    
	    (for ((elem count : self)) ()
		 (let ((newcount (max count (pylist-ref other elem))))
		   (if (> newcount 0)
		       (pylist-set! result elem newcount))))
	    
	    (for ((elem count : other)) ()
		 (if (and (not (in elem self))
			  (> count 0))
		     (pylist-set! result elem count)))
	    
	    return))))

  (define __and__
    (lambda (self other)
      (if (not (isinstance other Counter))
	  NotImplemented
	  (let ((result (Counter)))	    
	    (for ((elem count : self)) ()
		 (let ((newcount (min count (pylist-ref other elem))))
		   (if (> newcount 0)
		       (pylist-set! result elem newcount))))	    
	    return))))

  (define __pos__
    (lambda (self)
      (let ((result (Counter)))       
        (for ((elem count : self))
	     (if (> count 0)
                (pylist-set! result elem count)))
        result)))


  (define __neg__
    (lambda (self)
      (let ((result (Counter)))
        (for ((elem count : self)) ()
	     (if (< count 0)
		 (pylist-set! result elem (- count))))
        result)))

  (define _keep_positive
    (lambda (self)
      (define ks
	(for ((k v : self)) (l '())
	     (if (<= v 0)
		 (cons k l)
		 l)))
      (let lp ((ks ks))
	(if (pair? ks)
	    (begin
	      (pylist-remove! self (car ks))
	      (lp (cdr ks)))))
      self))

  (define __iadd__
    (lambda (self, other)
      (for ((elem count : other))
	   (pylist-set! self elem
			(+ (pylist-get self elem 0) count)))
        ((ref self '_keep_positive))))

  (define __isub__
    (lambda (self, other)
      (for ((elem count : other))
	   (pylist-set! self elem
			(- (pylist-get self elem 0) count)))
      ((ref self '_keep_positive))))

  (define __ior__
    (lambda (self, other)
      (for ((elem count : other))
	   (pylist-set! self elem
			(max (pylist-get self elem 0) count)))
      ((ref self '_keep_positive))))

  (define __iand__
    (lambda (self, other)
      (for ((elem count : other))
	   (pylist-set! self elem
			(min (pylist-get self elem 0) count)))
      ((ref self '_keep_positive)))))

(define-python-class UserDict   (dict))
(define-python-class UserString (py-string))
(define-python-class UserList   (py-list))