hy/tests/native_tests/comprehensions.hy

107 lines
3.2 KiB
Hy
Raw Normal View History

(import
types
pytest)
(defn test-comprehension-types []
; Forms that get compiled to real comprehensions
(assert (is (type (lfor x "abc" x)) list))
(assert (is (type (sfor x "abc" x)) set))
(assert (is (type (dfor x "abc" [x x])) dict))
(assert (is (type (gfor x "abc" x)) types.GeneratorType))
; Forms that get compiled to loops
(assert (is (type (lfor x "abc" :do (setv y 1) x)) list))
(assert (is (type (sfor x "abc" :do (setv y 1) x)) set))
(assert (is (type (dfor x "abc" :do (setv y 1) [x x])) dict))
(assert (is (type (gfor x "abc" :do (setv y 1) x)) types.GeneratorType)))
#@ ((pytest.mark.parametrize "specialop" ["lfor" "sfor" "gfor" "dfor"])
(defn test-comprehensions [specialop]
(setv cases [
['(f x [] x)
[]]
['(f j [1 2 3] j)
[1 2 3]]
['(f x (range 3) (* x 2))
[0 2 4]]
['(f x (range 2) y (range 2) (, x y))
[(, 0 0) (, 0 1) (, 1 0) (, 1 1)]]
['(f (, x y) (.items {"1" 1 "2" 2}) (* y 2))
[2 4]]
['(f x (do (setv s "x") "ab") y (do (+= s "y") "def") (+ x y s))
["adxy" "aexy" "afxy" "bdxyy" "bexyy" "bfxyy"]]
['(f x (range 4) :if (% x 2) (* x 2))
[2 6]]
['(f x "abc" :setv y (.upper x) (+ x y))
["aA" "bB" "cC"]]
['(f x "abc" :do (setv y (.upper x)) (+ x y))
["aA" "bB" "cC"]]
['(f
x (range 3)
y (range 3)
:if (> y x)
z [7 8 9]
:setv s (+ x y z)
:if (!= z 8)
(, x y z s))
[(, 0 1 7 8) (, 0 1 9 10) (, 0 2 7 9) (, 0 2 9 11)
(, 1 2 7 10) (, 1 2 9 12)]]
['(f
x [0 1]
:setv l []
y (range 4)
:do (.append l (, x y))
:if (>= y 2)
z [7 8 9]
:if (!= z 8)
(, x y (tuple l) z))
[(, 0 2 (, (, 0 0) (, 0 1) (, 0 2)) 7)
(, 0 2 (, (, 0 0) (, 0 1) (, 0 2)) 9)
(, 0 3 (, (, 0 0) (, 0 1) (, 0 2) (, 0 3)) 7)
(, 0 3 (, (, 0 0) (, 0 1) (, 0 2) (, 0 3)) 9)
(, 1 2 (, (, 1 0) (, 1 1) (, 1 2)) 7)
(, 1 2 (, (, 1 0) (, 1 1) (, 1 2)) 9)
(, 1 3 (, (, 1 0) (, 1 1) (, 1 2) (, 1 3)) 7)
(, 1 3 (, (, 1 0) (, 1 1) (, 1 2) (, 1 3)) 9)]]
['(f x (range 4) :do (unless (% x 2) (continue)) (* x 2))
[2 6]]
['(f x (range 4) :setv p 9 :do (unless (% x 2) (continue)) (* x 2))
[2 6]]
['(f x (range 20) :do (when (= x 3) (break)) (* x 2))
[0 2 4]]
['(f x (range 20) :setv p 9 :do (when (= x 3) (break)) (* x 2))
[0 2 4]]
['(f x [4 5] y (range 20) :do (when (> y 1) (break)) z [8 9] (, x y z))
[(, 4 0 8) (, 4 0 9) (, 4 1 8) (, 4 1 9)
(, 5 0 8) (, 5 0 9) (, 5 1 8) (, 5 1 9)]]])
(for [[expr answer] cases]
; Mutate the case as appropriate for the operator before
; evaluating it.
(setv expr (+ (HyExpression [(HySymbol specialop)]) (cut expr 1)))
(when (= specialop "dfor")
(setv expr (+ (cut expr 0 -1) `([~(get expr -1) 1]))))
(setv result (eval expr))
(when (= specialop "dfor")
(setv result (.keys result)))
(assert (= (sorted result) answer) (str expr)))))
(defn test-raise-in-comp []
(defclass E [Exception] [])
(setv l [])
(import pytest)
(with [(pytest.raises E)]
(lfor
x (range 10)
:do (.append l x)
:do (when (= x 5)
(raise (E)))
x))
(assert (= l [0 1 2 3 4 5])))