Permanently extending the behaviour of functions (like decorators)

Preamble

I've used similar constructs a number of times. This answer assumes that you want to decorate at definition time (which is what happens in python), rather than doing that after the definition has been created. If you need to add definitions to existing function, this and this discussions can be helpful. If you have several definitions, you will need to decorate all of them (where you want the decorated behavior) separately.

The main idea and simplest examples

The idea

Here is the idea:

ClearAll[decorator];
decorator /: SetDelayed[decorator[dec_][f_[args___]], rhs_]:=
  f[a:PatternSequence[args]]:=
     dec[Unevaluated @ f, Unevaluated @ {a}, Unevaluated @ rhs];

Links to similar usage examples

Here are a few links to answers where I previously used this method in one form or another - they can serve as additional examples of how this can be applied:

  • How to find the name of the current function?
  • Generic templating system
  • Resource management in Mathematica (the lex function there)
  • How to implement a discrete data representation?
  • How to protect functions in a Mathematica Package?

Tests

Here is how it can be used. First, we define our decorator

ClearAll[mydec];
mydec[f_, args_, rhs_] :=
  Module[{res},
    Print["Before..."];
    res = rhs;
    Print["After..."];
    res
  ]

Now, we define f:

decorator[mydec]@
 f[x_] := x^2;

We now test:

f[5]   

During evaluation of In[18]:= Before...

During evaluation of In[18]:= After...

(* 25  *)

Here is your sample app:

decorator[
  Function[{f, args, rhs}, 
  Block[{$ContextPath = {}, $Context = "myEmpty`"}, rhs]]
] @ g[] := ToString[sym]

then

g[]

(* Global`sym *)

A more complete framework

The code

Here is the code for a more complete framework, which would, in particular, allow for decorator chaining:

ClearAll[decorator, customHold, foldRHS];
SetAttributes[{customHold, foldRHS}, HoldAll];

foldRHS[f_, args_, rhs_, decs_List]:=
  ReplaceRepeated[
    Hold @ Evaluate @ Fold[
      customHold[#2[f, args, #]]&, 
      customHold[Unevaluated[rhs]], 
      Reverse @ decs
    ],
    customHold[x___]:>x
  ]; 

decorator /: (h:(Set | SetDelayed))[decorator[decs__][decorator[d_][arg_]], rhs_]:=
  h[decorator[decs, d][arg], rhs];

decorator /: (h:(Set | SetDelayed))[
  decorator[decs__][f_[args___, o:Verbatim[OptionsPattern[]] | Verbatim[Pattern][_, Verbatim[OptionsPattern[]]]]],
  rhs_
]:=
  Module[{a, opts},
    h @@ Join[
      Hold[f[a:PatternSequence[args], opts: o]], 
      foldRHS[f, {a, opts}, rhs, {decs}]    
    ]
  ];    

decorator /: (h:(Set | SetDelayed))[decorator[decs__][f_[args___]], rhs_]:=
  Module[{a},
    h @@ Join[
      Hold[f[a:PatternSequence[args]]], 
      foldRHS[f, {a}, rhs, {decs}]  
    ]
  ];

decorator /: Set[decorator[decs__][lhs_], HoldPattern @ Function[args_, body_, attrs_:{}]]:=
  lhs = Function @@ Join[
    Hold[args],
    foldRHS[None, args, body, {decs}],
    Hold[attrs]
  ];

Simple examples

Let us define decorator generator, which can generate simple wrapping decorators:

ClearAll[dec];  
dec[name_, hold_:False]:=
  With[{atts = If[TrueQ @ hold, {HoldAll}, {}]},
    Function[{fun, args, rhs},
      Module[{res},
        Print["In ", name, " before..."];
        Print["Result of ", HoldForm[rhs], " : ", res = rhs];
        Print["In ", name, " after..."];
        res
      ],
      atts
    ]
  ];    

And use them to define two functions:

ClearAll[f, ff];

decorator[dec["first"]] @
decorator[dec["second"]] @ 
f[x_]:=x^2

decorator[dec["first", True]] @
decorator[dec["second", True]] @ 
ff[x_]:=x^2     

We can now test. In the first case, we see that the inner decorator executes before the outer one:

f[5]   

During evaluation of In[378]:= In second before...

During evaluation of In[378]:= Result of 5^2 : 25

During evaluation of In[378]:= In second after...

During evaluation of In[378]:= In first before...

During evaluation of In[378]:= Result of 25 : 25

During evaluation of In[378]:= In first after...

(* 25 *)

because neither of them holds the internal code. But in the second case, the execution is happening from outside to inside:

ff[5]   

During evaluation of In[380]:= In first before...

During evaluation of In[380]:= In second before...

During evaluation of In[380]:= Result of 5^2 : 25

During evaluation of In[380]:= In second after...

  During evaluation of In[380]:= Result of 
    Function[{fun$,args$,rhs$},Module[{res$},
     Print[In ,second, before...];
     Print[Result of ,rhs$, : ,res$=rhs$];
     Print[In ,second, after...];res$],
     {HoldAll}][ff,{5},Unevaluated[5^2]
    ] : 25

During evaluation of In[380]:= In first after...

(* 25 *)

Using pure functions

The decorator defined above also allows one to decorate definitions like this:

someVariable = Function[args, body]

So for example, the following pf function will behave like the f above:

decorator[dec["first"]] @
decorator[dec["second"]] @
pf = Function[x, x^2];

which can be easily checked.

An example: option validation

To add one less trivial example, I would like to show how option validation can be carried out using decorators. I will use this answer as a source of examples, although the validation process described below will be different (will check option values immediately as they are passed, unlike in my answer there, which was checking them once they were computed / used).

The following two functions will implement the option-checking logic:

ClearAll[checkOptions];
checkOptions[testF_][f_, args___, opts:OptionsPattern[]]:=
  False =!= Replace[
    Flatten[{opts}], 
    (name_ -> val_) :>  If[!TrueQ @ testF[f, name, val], 
      Message[OptionCheck::invldopt, name, f, val];
      Return[False, Replace]    
    ],
    {1}
  ];

ClearAll[OptionCheck];
OptionCheck::invldopt = "Option `1` for function `2` received invalid value `3`";
OptionCheck[testFunction_]:=
  Function[
    {f, args, code},
    If[!checkOptions[testFunction][f, Sequence @@ args],
      Return[$Failed],
      (* else *)
      code
    ],
    HoldAll
  ];

Basically, OptionCheck takes a test function and generates a decorator, which we can use to validate options. The test function should take the function name, option name, and option value, and return True if the option value is valid and False otherwise.

Let us now define some function:

ClearAll[g];

Options[g]={"a"->1,"b"->False};

decorator[Function[{f, args, code}, 
  Print["The passed arguments were: ", args];code, 
  HoldAll]
] @
decorator[
  OptionCheck[optcheck]
] @ 
g[x_,y_,OptionsPattern[]]:=
  Module[{z=x+y,q,a},a=OptionValue["a"];
    q = If[OptionValue["b"],a,a+1];
    {x,y,q}
  ]

and we also need to define optcheck:

ClearAll[optcheck];
optcheck[g, "a",val_]:=IntegerQ[val];
optcheck[g, "b",val_]:=MatchQ[val,True|False];

Now we can test. First, testing the valid inputs:

g[1, 2]
g[1, 2, "a" -> 10]
g[1, 2, "a" -> 10, "b" -> False]

During evaluation of In[392]:= The passed arguments were: {1,2}

(* {1, 2, 2} *)

During evaluation of In[392]:= The passed arguments were: {1,2,a->10}

(* {1, 2, 11} *)

During evaluation of In[392]:= The passed arguments were: {1,2,a->10,b->False}

(* {1, 2, 11} *)

where you can see that the outer decorator was printing the passed arguments.

Now, the input where some options are given values of inappropriate type:

g[1, 2, "a" -> 1.5]

During evaluation of In[395]:= The passed arguments were: {1,2,a->1.5}

During evaluation of In[395]:= OptionCheck::invldopt: Option a for function g received invalid value 1.5`

(* $Failed *)

and

g[1, 2, "a" -> 1, "b" -> 2]

During evaluation of In[396]:= The passed arguments were: {1,2,a->1,b->2}

During evaluation of In[396]:= OptionCheck::invldopt: Option b for function g received invalid value 2

(* $Failed *)

where you can see both decorators at play - the outer one was still printing the passed arguments, while the inner one detected invalid options, issued error message and returned $Failed without executing the body of the function.

An aside: the role of decorators

In python, decorators are very useful. But not because you can modify the behavior of some existing function (because if the function already exists, chances are that you don't own it, but rather import it from some other package). The main usefulness of decorators is that they factor out the logic, which doesn't belong to the function you write, but rather should augment it in some way. For example, perform certain argument checks or checks of the result. Or even add some functionality to your function or class (for example, class decorators can add some methods to the class you define, and can be cheaper alternatives to metaclasses).

So, they are useful to add certain behavior and factorize some logic which can be implemented separately, from your main function. The main point is, you typically use decorators on the code you control / write, to have it more concise, and also when certain checks / etc have been already implemented elsewhere and you simply need to decorate your function to enable them.

Changing existing definitions is also a well-known practice in python, but it is called monkey-patching. This typically involves changes performed on existing functions or classes, where for example you can import a class and add or modify some of its methods at run-time. This is usually reserved to cases where you have to work around some bug or limitation in code which you use but don't own, and when other softer methods (such as subclassing) are not an option (for example, the authors of the package did not provide necessary hooks).

So in Mathematica too, changing definitions of already existing function generally doesn't strike me as a good idea. This is a much more intruding method, and the necessity of doing this means that you perform it on the code that you don't own / control. This can be much more error-prone, you can't easily chain this approach, and I would consider this being the last resort, just like monkey - patching in python.

Notes

The code presented above can surely be further improved and extended. I did not intend to provide a fully bulletproof production - quality implementation, this has rather a status of a proof - of - concept prototype. For example, conditional definitions such as f[arg_]/;condition := rhs won't work with the simple approach above, and there are a few other limitations. Perhaps, at some point the code above could become more robust, and then could actually be used for real applications.


Based on your requirements I would move the original symbol to a new context.

wrapper[fn_, mod_] :=
  Module[{ctxt = Context[fn], name = SymbolName[fn], x},
    Context[fn] = "wrapper`" <> ctxt <> ToString[x] <> "`";
    Symbol[ctxt <> name][arg___] := fn[arg] // Unevaluated // mod
  ]

Your example:

sym;
g[] := ToString[sym]

wrapper[g, Block[{$ContextPath = {}, $Context = "myEmpty`"}, #] &]

g[]
"Global`sym"

Recursion:

fib[1 | 2] = 1;
fib[n_Integer?Positive] := fib[n - 1] + fib[n - 2]

wrapper[fib, #2 &[Print["Before"], #, Print["After"]] &]

fib[7]

Before

After

13

Nested wrappers:

wrapper[fib, #2 &[Print["Primordial"], #, Print["Final"]] &]

fib[12]

Primordial

Before

After

Final

144

Related:

  • Same name for functions in package and Global context

A seemingly robust approach seems to be to modify the RHS of DownValues of the symbol you want to decorate. So lets define a decorator function:

decorate[f_Symbol] := (
 DownValues[f] = 
  Cases[DownValues[f], RuleDelayed[A_, expr_] :> RuleDelayed[A,
   Module[{result},
    Print["Before..."];
    result = expr;
    Print["After..."];
    result
   ]
  ]
 ]
)

This will modify all DownValue definitions of the symbol with the appropriate wrapper. (It will not work for built in symbols, and functions that are defined through other means e.g. UpValues)

We use the Cases statement for the replacement to ensure two things:

  1. The RHS of the DownValues are not evaluated during decoration.
  2. Only DelayedRule statements at the top level of the DownValues are replaced.

Example:

Lets define a function f:

Clear[f];
f[x_]:=x^2;
f[1] = 0; 
f[]  = Null;

We can then add the "decoration" by:

decorate@f;

evaluation of f then yields:

f[5]
During evaluation of In[]:= Before...

During evaluation of In[]:= After...
(* 25 *)

and

f[1]
During evaluation of In[]:= Before...

During evaluation of In[]:= After...
(* 0 *)

and

f[]
During evaluation of In[]:= Before...

During evaluation of In[]:= After...
(* Null *)