Constructing functions with variable number of output arguments

Question

In Mathematica, you can construct a function f to have different definitions based on its input arguments. For example, f[x_] := ... and f[x_, y_] := .... You can also get more specific and define it for specific heads as f[x_someHead] := ... or inputs matching arbitrary patterns as f[x_?somePatternQ] := ... (of course, one must be careful to keep the order from specific to general). These definitions are the DownValues of the function.

Now, I'd like to construct a function that has different definitions based on the number of output variables/symbols requested, but for the same set of inputs. Consider the following dummy example in pseudo Mathematica code

f[mu_, sigma2_] := Switch[OutputArgs[], 
                           0, Plot[PDF[NormalDistribution[mu, sigma2], x], ...]
                           1, RandomReal[NormalDistribution[mu, sigma2], 1000]
                           2, {Mean@#, Variance@#}&@ RandomReal[...]
                   ]

Here OutputArgs[] is a functionality that I'd like, which when called from inside a function, tells you how many outputs have been requested (loosely reminiscent of OptionValue which knows which function it's called from). With this, I can use the function call simply as f[0,1] to plot the normal distribution (perhaps inside a Manipulate to play with the parameters), then when I'm satisfied, I can use the same function call, but with an output argument as a = f[0,1] to get a random sample, and with two outputs to get an empirical estimate of the mean and variance.

I realize that this is a bit at odds with the classical notion of a function having a well defined output based on the input, rather than a well defined output based on the number of outputs for a given input.

---

My interest is purely academic and arose from working on porting some MATLAB code to Mathematica. Some might have recognized that right away, and this is a feature that that I do find useful — a sentiment also shared by a few others that I've spoken to here. I wouldn't write code this way and I know that it is also possible to achieve the same by other means such as using flags as f[0, 1, "Plot"] or options as f[0, 1, Options -> "RandomSample"], etc. for the above example. However, I'm interested in exploring if there are ways to imitate this behaviour in Mathematica.

Answer 1

As you have mentioned there are other more standard ways to provide some variaty in your return value so I don't think that this is something that you should do. If you provide code using something like it you should have good reasons for doing so. Just being used to it because matlab has it probably is just not convincing enough.

That said, it can be done along these lines:

ClearAll[f, ff]
f /: Set[arg_Symbol, f[x_]] := Block[{nargout = 1}, arg = ff[x]];
f /: Set[args_List, f[x_]] := Block[{nargout = Length[args]}, args = ff[x]];
f[arg_] /; FreeQ[Stack[], Set] := Block[{nargout = 1}, ff[arg]];
ff[x_] := (Table[x, {nargout}]);

you can now use it like this:

In[101]:= f[5]
Out[101]= {5}
In[102]:= a=f[5]
Out[102]= {5}
In[104]:= {a,b}=f[5]
Out[104]= {5,5}
In[105]:= {a,b,c}=f[5]
Out[105]= {5,5,5}

This is not tested well and could have many pitfalls. I see it rather as a proof of concept and probably something for your own convenience. So use with care, especially if you are uncertain about all the details (as I am: I can't remember I ever had the need to use Stack).

Brett already found one pitfall. Actually that is a case where one could be argueing what the expected return value should be. If we fall back to the one argument case for everything that is not an explicit Set it is easy to solve, but I can see other problems already :-):

f[arg_] /; Not[MatchQ[Stack[],{___,Set,f,___}]]:=Block[{nargout=1},ff[arg]];

Answer 2

I can offer an alternative approach based on introspection and run-time code generation. It is not as elegant as one by @Albert, to be sure, but may still represent some interest. Here is the code (I made an assumption that your function is a symbol which has DownValues, but this can be generalized to other ...Values:

ClearAll[withNArgOut];
SetAttributes[withNArgOut, HoldAll];
withNArgOut[code_] :=
  Hold[code] //.
   {
    HoldPattern[Set[vars_List, fcall : f_Symbol[___]]] /;
      FreeQ[Unevaluated[fcall], $nArgAout] && ! FreeQ[DownValues[f], $nArgAout] :>
            Set[vars, Block[{$nArgAout = Length[Unevaluated[vars]]}, fcall]],
    HoldPattern[Set[var_, fcall : _[___]]] /; 
       !MatchQ[Unevaluated[var], _List] &&
         Hold[var] =!= Hold[$nArgAout] && FreeQ[Unevaluated[fcall], $nArgAout] :>
            Set[{var}, fcall]
   };

I intentionally did not release the last Hold, so that we can see what is generated. Here is how you could use it:

ClearAll[f];
f[x_] := Table[x^i, {i, $nArgAout}]

In[14]:= {#,ReleaseHold@#}&@withNArgOut[Module[{a,b},b=1;a=f[x]]]
Out[14]= {Hold[Module[{a,b},b=1;{a}=Block[{$nArgAout=Length[Unevaluated[{a}]]},f[x]]]],{x}}

In[15]:= {#,ReleaseHold@#}&@withNArgOut[Module[{a,b,c}, c=1;{a,b}=f[x]]]
Out[15]= {Hold[Module[{a,b,c},c=1;
           {a,b}=Block[{$nArgAout=Length[Unevaluated[{a,b}]]},f[x]]]],{x,x^2}}

In[16]:= withNArgOut[Module[{a,b,c}, c=1;{a,b}=g[f[x]]]]
Out[16]= Hold[Module[{a,b,c},c=1;{a,b}=g[f[x]]]]

You can always add a call to releaseHold in the definition of withNArgOut. The usage is that you wrap the withNArgOut around the code where you want this to take effect, and the extra Block[{$nArgAout= ...},...] wrappers will be generated for that code (where appropriate), just prior to execution.