Mixin to wrap every method of a Scala trait

Question

Suppose I have a trait Foo with several methods. I want to create a new trait which extends Foo but "wraps" each method call, for example with some print statement (in reality this will be something more complicated / I have a couple of distinct use cases in mind).

trait Foo {
  def bar(x: Int) = 2 * x
  def baz(y: Int) = 3 * y
}

I can do this manually, by overriding each method. But this seems unnecessarily verbose (and all too easy to call the wrong super method):

object FooWrapped extends FooWrapped
trait FooWrapped extends Foo {
  override def bar(x: Int) ={
    println("call")
    super.bar(x)
  }
  override def baz(y: Int) ={
    println("call")
    super.baz(y)
  }
}

scala> FooWrapped.bar(3)
call
res3: Int = 6

I was hoping to write a mixin trait, that I would be able to reuse with other traits, and use as:

trait FooWrapped extends Foo with PrintCall

That way I don't have to manually override each method (the mixin would do this for me).

Is it possible to write such a mixin trait in Scala? What would it look like?

It should be doable using macro annotations (docs.scala-lang.org/overviews/macros/annotations), I think. This is a comment, not an answer, because I can't write the macro at the moment. — Alexey Romanov
You can also use the stackable trait pattern to append certain behavior to existing traits, such that your invocations look like: val a = new SomeClass with PrintCall with Logger with Whatever. You simply inject the behaviors in the order you wish them to be performed (in case of print and log it doesn't matter, but if your traits are AddTwo and MultiplyByThree then the order matters). I'm writing in comment since it's not really what you asked for, but if you're not familiar with this pattern do check it out, it can be pretty cool (however you do need to explicitly override each method). — slouc
@slouc I'm hoping for a solution that doesn't mean I have to write override def bar manually for every method. Mainly as doing that means I can't write this in general way (without knowing all the methods and signatures in advance). — Andy Hayden

Archeg Archeg · Accepted Answer · 2016-03-27T20:41:12

Update Here is the macro. It was much less painful than I thought it will be because of quasiquotes. They are awesome. This code does only a little and you probably will have to improve it. It may not account some special situations. Also it assumes that neither parent class nor it's method has type params, it wraps only the methods of the given class or trait, but not it's parents methods, it may not work if you have auxilary constructors etc. Still I hope it will give you an idea of how to do that for your specific needs, making it working for all of the situations unfortunately is too big job for me right now.

object MacrosLogging {
  import scala.language.experimental.macros
  import scala.reflect.macros.blackbox


  def log_wrap[T](): T = macro log_impl[T]
  def log_impl[T : c.WeakTypeTag](c: blackbox.Context)(): c.Expr[T] = {
    import c.universe._

    val baseType = implicitly[c.WeakTypeTag[T]].tpe
    val body = for {
       member <- baseType.declarations if member.isMethod && member.name.decodedName.toString != "$init$"
      method = member.asMethod
      params = for {sym <- method.paramLists.flatten} yield q"""${sym.asTerm.name}: ${sym.typeSignature}"""
      paramsCall = for {sym <- method.paramLists.flatten} yield sym.name
      methodName = member.asTerm.name.toString
    } yield  {
      q"""override def ${method.name}(..$params): ${method.returnType} = { println("Method " + $methodName + " was called"); super.${method.name}(..$paramsCall); }"""
    }

    c.Expr[T] {q""" { class A extends $baseType { ..$body }; new A } """}
  }
}

If you do not want to create an instance, but you do want to add logging only for your trait so you could mixin further, you can do this with relatively the same code, but using macro paradise type annotations: http://docs.scala-lang.org/overviews/macros/annotations These allow you to tag your class definitions and perform modifications right inside the definitions

You could do something like you want with Dynamic, but there is a catch - you can't make it of original type, so it's not a mixin. Dynamic starts to work only if type checks fails, so you can't mixin real type (or I do not know how to do that). The real answer would probably require macros (as @AlexeyRomanov suggested in comments), but I am not sure how to write one, maybe I'll come up with it later. Still Dynamic might work for you if you are not looking for DI here

  trait Foo {
    def bar(x: Int) = 2 * x
    def baz(y: Int) = 3 * y
  }
  import scala.reflect.runtime.{universe => ru}
  import scala.language.dynamics

  trait Wrapper[T] extends Dynamic {
    val inner: T
    def applyDynamic(name: String)(args: Any*)(implicit tt: ru.TypeTag[T], ct: ClassTag[T]) = {
      val im = tt.mirror.reflect(inner)
      val method = tt.tpe.decl(ru.TermName(name)).asMethod
      println(method)
      val mm = im.reflectMethod(method)
      println(s"$name was called with $args")
      mm.apply(args:_*)
    }
  }

  class W extends Wrapper[Foo] {
    override val inner: Foo = new Foo() {}
  }

  val w = new W // Cannot be casted to Foo
  println(w.bar(5)) // Logs a call and then returns 10

You can read more about Dynamic here: https://github.com/scala/scala/blob/2.12.x/src/library/scala/Dynamic.scala

Mixin to wrap every method of a Scala trait

1 Answers