Traits do not work the way I thought

My original idea was to use traits to change the behavior of a Graph object. When I ran into trouble I came up with this script to simplify the use case.

#!/usr/bin/env groovy
class Edge {}
trait DirectedEdge {}
class Graph {
    def edge() {
        newEdge()
    }
    def newEdge() {
        return new Edge()
    }
}
trait DirectedGraph {
    def newEdge() {
        return new Edge().withTraits(DirectedEdge)
    }
}

def graph = new Graph()
assert graph.newEdge() instanceof Edge : 'newEdge returns Edge'
assert graph.edge() instanceof Edge : 'edge returns an Edge'

def dgraph = new Graph() as DirectedGraph
assert dgraph.newEdge() instanceof DirectedEdge : 'newEdge returns DirectedEdge'
assert dgraph.edge() instanceof DirectedEdge : 'edge returns DirectedEdge'

It turns out that traits do not work the way I expected. The final assert fails which is one of the keys to changing the behavior of Graph to work like a DirectedGraph.

When I got stuck I asked this question on stack overflow. The answer explains the problem. When a trait is applied a proxy is returned which adds those methods and allows other method calls to fall through to the original object. Since the original object still implements newEdge and it is not overridden, the original newEdge method is used. This is why the edge method which calls newEdge is not affected when the DirectedGraph trait is applied. To fix this I changed the design to use metaprogramming to add features to the method newEdge and replace it.

#!/usr/bin/env groovy
class Edge {}
trait DirectedEdge {}
class Graph {
    def apply(Class pluginClass) {
        def plugin = pluginClass.newInstance()
        plugin.apply(this)
    }
    def edge() {
        newEdge()
    }
    def newEdge() {
        return new Edge()
    }
}

class DirectedGraph {
    def apply(Graph graph) {
        graph.metaClass.newEdge = { ->
            return graph.&newEdge().withTraits(DirectedEdge)
        }
    }
}

def graph = new Graph()
assert graph.newEdge() instanceof Edge : 'newEdge returns Edge'
assert graph.edge() instanceof Edge : 'edge returns an Edge'

def dgraph = new Graph()
dgraph.apply DirectedGraph
assert dgraph.newEdge() instanceof DirectedEdge : 'newEdge returns DirectedEdge'
assert dgraph.edge() instanceof DirectedEdge : 'edge returns DirectedEdge'

Graph is now changed to a DirectedGraph using dgraph.apply DirectedGraph. The apply method creates an instance of the provided class and calls its apply method passing the graph. The apply method performs the metaprogramming. It overwrites newEdge with the results of the old newEdge with the DirectedEdge trait. This should also make it very easy to apply things like DirectedWeightedGraph or have edges with attributes.