more function optics

Author: aplavin

src/functionlenses.jl

@@ -138,12 +138,14 @@ function set(x, ::typeof(abs), y)
     s = sign(x)
     iszero(s) ? y * one(x) : y * s
 end
+set(x, ::typeof(abs2), y) = set(x, abs, √y)
 
 set(x, ::typeof(mod2pi), y) = set(x, @optic(mod(_, 2π)), y)
 set(x, f::Base.Fix2{typeof(fld)}, y) = set(x, @optic(first(fldmod(_, f.x))), y)
 set(x, f::Base.Fix2{typeof(mod)}, y) = set(x, @optic(last(fldmod(_, f.x))), y)
 set(x, f::Base.Fix2{typeof(div)}, y) = set(x, @optic(first(divrem(_, f.x))), y)
 set(x, f::Base.Fix2{typeof(rem)}, y) = set(x, @optic(last(divrem(_, f.x))), y)
+set(x, f::Base.Fix2{typeof(mod),<:AbstractUnitRange}, y) = @set mod($x - first(f.x), length(f.x)) + first(f.x) = y
 
 set(x::AbstractString, f::Base.Fix1{typeof(parse), Type{T}}, y::T) where {T} = string(y)
 
@@ -171,6 +173,7 @@ delete(s::AbstractString, o::typeof(last)) = chop(s; head=0, tail=1)
 delete(s::AbstractString, o::Base.Fix2{typeof(first)}) = chop(s; head=o.x, tail=0)
 delete(s::AbstractString, o::Base.Fix2{typeof(last)}) = chop(s; head=0, tail=o.x)
 
+set(s::AbstractString, o::typeof(chomp), v) = endswith(s, '\n') ? v * '\n' : v
 if VERSION >= v"1.8"
     set(s::AbstractString, o::Base.Fix2{typeof(chopsuffix), <:AbstractString}, v) =
         endswith(s, o.x) ? v * o.x : v

test/test_functionlenses.jl

@@ -257,7 +257,7 @@ end
             # invertible lenses below: no need for extensive testing, simply forwarded to InverseFunctions
             inv, +, exp, sqrt, @optic(2 + _), @optic(_ * 3), @optic(log(2, _)),
             # non-invertible lenses, indirectly forwarded to InverseFunctions
-            @optic(mod(_, 21)), @optic(fld(_, 3)), @optic(rem(_, 21)), @optic(div(_, 3)),
+            @optic(mod(_, 21)), @optic(fld(_, 3)), @optic(rem(_, 21)), @optic(div(_, 3)), @optic(mod(_, 1:22)),
         ]
         x = 5
         test_getset_laws(o, x, 10, 20; cmp=isapprox)
@@ -272,6 +272,7 @@ end
     @test set(0+0im, abs, 10) == 10
     @test set(0+1e-100im, abs, 10) == 10im
     @test_throws DomainError @set(abs(x) = -10)
+    test_getset_laws(abs2, 1+2im, 3, 4, cmp=(≈))
 
     # composition
     o = @optic 1/(1 + exp(-_))
@@ -362,6 +363,9 @@ end
     test_getset_laws(@optic(lstrip(==(' '), _)), " abc  ", "def", "")
     test_getset_laws(@optic(rstrip(==(' '), _)), " abc  ", "def", "")
     test_getset_laws(@optic(strip(==(' '), _)), " abc  ", "def", "")
+    test_getset_laws(chomp, "abc", "def", "")
+    test_getset_laws(chomp, "abc\n", "def", "")
+    test_getset_laws(chomp, "abc\n\n", "def\n", "")
 
     if VERSION >= v"1.8"
         test_getset_laws(@optic(chopprefix(_, "def")), "def abc", "xyz", "")