Wrap function implementations returning a specific type into another function programatically

问题

I would like to wrap all the user defined functions in a scala project that return a certain type T, into a function that accepts a T and the function name as parameters.

eg.

given this function is in scope:

 def withMetrics[T](functionName: String)(f: => Try[T]): Try[T] = {
    f match {
        case _: Success[T] => println(s"send metric: success for $functionName")
        case _: Failure[T] => println(s"send metric: failure for $functionName")
    }

    f
}

the user can send metrics for their functions which return Try by doing

def userDefinedFunction: Try[_] =
    withMetrics("userDefinedFunction"){
        somethingRisky: Try[_]
    }

but I would like the user to only have to define

def userDefinedFunction: Try[_] =
    somethingRisky: Try[_]

and have his business logic that returns Try wrapped into withMetrics implicitly.

Note that the user should not have to annotate code, as that could lead to him forgetting about it. Instead all the user functions defined in his project should be wrapped into withMetrics automatically.

How can I achieve this by using Scala 2 or dotty macros? Or can this be achieved in another way?

回答1:

You can create macro annotation and annotate with it all classes, objects and traits where you want to instrument your methods.

import scala.annotation.{StaticAnnotation, compileTimeOnly}
import scala.language.experimental.macros
import scala.reflect.macros.whitebox

object Macros {
  @compileTimeOnly("enable macro paradise (or -Ymacro-annotations in 2.13) to expand macro annotations")
  class withMetrics extends StaticAnnotation {
    def macroTransform(annottees: Any*): Any = macro WithMetricsMacro.impl
  }

  object WithMetricsMacro {
    def impl(c: whitebox.Context)(annottees: c.Tree*): c.Tree = {
      import c.universe._

      def modify(stats: Seq[Tree]): Seq[Tree] = stats.map {
        case q"$mods def $tname[..$tparams](...$paramss): $tpt = $expr" =>
          q"$mods def $tname[..$tparams](...$paramss): $tpt = withMetrics(${tname.toString}){ $expr }"
      }

      annottees match {
        case q"$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends { ..$earlydefns } with ..$parents { $self => ..$stats }" :: tail =>
          q"""
             $mods class $tpname[..$tparams] $ctorMods(...$paramss) extends { ..$earlydefns } with ..$parents { $self => ..${modify(stats)} }
             ..$tail
            """
        case q"$mods trait $tpname[..$tparams] extends { ..$earlydefns } with ..$parents { $self => ..$stats }" :: tail =>
          q"""
              $mods trait $tpname[..$tparams] extends { ..$earlydefns } with ..$parents { $self => ..${modify(stats)} }
              ..$tail
             """
        case q"$mods object $tname extends { ..$earlydefns } with ..$parents { $self => ..$stats }" :: Nil =>
          q"$mods object $tname extends { ..$earlydefns } with ..$parents { $self => ..${modify(stats)} }"

        case _ =>
          c.abort(c.enclosingPosition, "Not a class, object or trait ")
      }
    }
  }
}

import Macros._
import scala.util.{Failure, Success, Try}

object App {

  @withMetrics
  class A {
    def userDefinedFunction: Try[String] = Try("aaa")
  }

  def withMetrics[T](functionName: String)(f: => Try[T]): Try[T] = {
    f match {
      case _: Success[T] => println(s"send metric: success for $functionName")
      case _: Failure[T] => println(s"send metric: failure for $functionName")
    }

    f
  }

  def main(args: Array[String]): Unit = {
    (new A).userDefinedFunction // send metric: success for userDefinedFunction
  }
}

This doesn't modify nested methods and methods in inner classes, objects, traits. If necessary this can be done too with scala.reflect.api.Trees.Traverser/Transformer. Or you can just annotate inner classes, objects, traits when necessary.

来源：https://stackoverflow.com/questions/56917025/wrap-function-implementations-returning-a-specific-type-into-another-function-pr