add ∘ and ⨟ function between varname and optic

TuringLang · Apr 9, 2024 · ac5de77 · ac5de77
1 parent 9d35b63
commit ac5de77
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/src/varname.jl b/src/varname.jl
@@ -1,5 +1,6 @@
 using Accessors
 using Accessors: ComposedOptic, PropertyLens, IndexLens, DynamicIndexLens
+import Accessors: opcompose, ⨟
 
 const ALLOWED_OPTICS = Union{typeof(identity),PropertyLens,IndexLens,ComposedOptic}
 
@@ -148,6 +149,13 @@ function Accessors.set(obj, vn::VarName{sym}, value) where {sym}
     return Accessors.set(obj, PropertyLens{sym}() ⨟ getoptic(vn), value)
 end
 
+# Allow compositions with optic.
+function Base.:∘(optic::ALLOWED_OPTICS, vn::VarName{sym,<:ALLOWED_OPTICS}) where {sym}
+    return VarName{sym}(optic ∘ getoptic(vn))
+end
+function Accessors.opcompose(vn::VarName{sym,<:ALLOWED_OPTICS}, optic::ALLOWED_OPTICS) where {sym}
+    return VarName{sym}(getoptic(vn) ⨟ optic)
+end
 
 Base.hash(vn::VarName, h::UInt) = hash((getsym(vn), getoptic(vn)), h)
 function Base.:(==)(x::VarName, y::VarName)
@@ -631,8 +639,8 @@ function varname(expr::Expr, concretize=Accessors.need_dynamic_optic(expr))
         if concretize
             return :(
                 $(AbstractPPL.VarName){$sym}(
-                    $(AbstractPPL.concretize)($optics, $sym_escaped)
-                )
+                $(AbstractPPL.concretize)($optics, $sym_escaped)
+            )
             )
         elseif Accessors.need_dynamic_optic(expr)
             error("Variable name `$(expr)` is dynamic and requires concretization!")
@@ -687,7 +695,7 @@ function _parse_obj_optics(ex)
         else
             throw(ArgumentError(
                 string("Error while parsing :($ex). Second argument to `getproperty` can only be",
-                       "a `Symbol` or `String` literal, received `$property` instead.")
+                    "a `Symbol` or `String` literal, received `$property` instead.")
             ))
         end
     else

diff --git a/test/varname.jl b/test/varname.jl
@@ -5,7 +5,7 @@ using OffsetArrays
 using AbstractPPL: ⊑, ⊒, ⋢, ⋣, ≍
 
 using AbstractPPL: Accessors
-using AbstractPPL.Accessors: IndexLens, PropertyLens
+using AbstractPPL.Accessors: IndexLens, PropertyLens, ⨟
 
 macro test_strict_subsumption(x, y)
     quote
@@ -50,6 +50,11 @@ end
         @test test_equal(@varname(x.a[1:end, end][:], true), @varname(x.a[1:3,2][1:3]))
     end
 
+    @testset "compose and opcompose" begin
+        @test IndexLens(1) ∘ @varname(x.a) == @varname(x.a[1])
+        @test @varname(x.a) ⨟ IndexLens(1) == @varname(x.a[1])
+    end
+
     @testset "get & set" begin
         x = (a = [1.0 2.0; 3.0 4.0; 5.0 6.0], b = 1.0);
         @test get(x, @varname(a[1, 2])) == 2.0