Programmatic formatting for Mathematica code - possible?

Question

It seems to be both an interesting programming challenge and a very useful practical application to have a Mathematica program which would allow one to pretty-print Mathematica code, so that it is easy to read and understand. Some of the desirable features of such a program:

• Works both on string and box level
• Covers many or most common cases for Mathematica code
• Is robust
• Is customizable
• Is extensible

Is this possible?

Answer 1

Update November 3, 2013

Finally, the formatter has been made much more robust by adding a custom function like MakeBoxes, names CodeFormatterMakeBoxes, to construct simplified box representation. This solves the main problem that the formatter does not currently support many boxes, since CodeFormatterMakeBoxes constructs the pure RowBox-based representation.

Also, added functions CodeFormatterPrint which prints definitions for a given function / symbol, and CodeFormatterSpelunk to make system spelunking easier.

Some things to try:

Import["https://raw.github.com/lshifr/CodeFormatter/master/CodeFormatter.m"]

and then

CodeFormatterPrint[RunThrough]

CodeFormatterSpelunk[RunThrough]

Basically, the difference is only that CodeFormatterPrint does retain long (fully-qualified) symbol names, while CodeFormatterSpelunk does not. Some more details on spelunking in my answer here.

Update 19th of June, 2013

Code formatter has been extended to support spaces in place of tabs, variable - width tabs and an overall offset.

This made it possible finally to start using it for code formatting here on SE - check this out!

Also, several other improvements and several bugs fixed.

---

Short answer

Yes.

Code formatter

While I will be most happy to see other answers to this question, and will hope that there will be better answers than this one, I am also glad to announce the alpha version of the Mathematica code formatter, which I was working on for some time.

The project

The code formatter lives here, and the specific file (package) can be downloaded using this link. The code fomatter resides in a package CodeFormatter.m, and has currently two public functions: FullCodeFormat and FullCodeFormatCompact. Both take a piece of code converted to boxes, and return the box form of formatted code. The README file in the project contains the brief description of how to use them, and the notebook which is included in the project contains many more examples.

Stealing from README, the typical way to use this is to define a helper function like this one:

prn = CellPrint[Cell[BoxData[#], "Input"]] &

and then, use it like:

prn@FullCodeFormat@MakeBoxes@
   Module[{a, b}, a = 1; Block[{c, d}, c = a + 1; d = b + 2]; b]

Screenshots

These are some screenshots of code pieces processed by FullCodeFormat.

Further plans/development

There are a number of things I plan to add to the formatter and /or develop based on it, such as

• Develop the palette to paste code to SE, based on the formatter (this will come out very soon)
• Support more boxes
• Refactor the code to eliminate the code duplication and better separate the DSL layer
• Add more ways to format code

I will start accepting pull requests soon, perhaps after I do the main code refactoring I currently plan. Meanwhile, please do fork me on GitHub if you are interested in playing with this.

Comments, suggestions, bug reports

All welcome. For bug reports, I have not decided yet what would be the best place to put them, but the Issues section on GitHub project repository seems appropriate.

A simplified bare-bones code formatting engine

Here, I will try to explain my approach, and provide a minimal functioning code-formatting "engine", which is a simplified version of the one I referred to above. The motivation for this section (and in fact, for placing the entire post here rather than on meta for example), is to make the code of the formatter accessible, and explain it in simple terms. This would allow you to fork the project and modify the formatter easier to suit your needs, should you wish to do so.

Design problems and choices

The main idea here is that while parsed expression is a bit too high-level for the formatting purposes, plus tries to evaluate all the time, the box-level is a bit too low-level, and the formatter based on that has a danger of not being robust.

Therefore, I take the box input, and create an intermediate inert Mathematica expression representation with preprocess and preformat (see below). The formatting procedure itself has two stages: the format proper - which only decides where to put new lines and tabs, and the tabification stage (tabify), which "executes" the tabification instructions given by format at the previous stage.

This architecture is because there is a certain impedance mismatch between the statement "I want to move this block of several lines of code one or more tabs to the right", and the actual way to achieve that with boxes (I suspect, this was one of the major obstacles for implementing code formatter - since I am sure, many people tried that). By separation of these two stages, it was possible to make this tabification abstraction reasonably robust. The final stage is post-formatting (postformat). It takes the result produced by format and tabify, and converts that back to boxes.

By using this 3-layer architecture, I am able to make high-level description of the actual formatting rules, in the definitions of format only, and the rest is taken care of by other layers. This makes the formatter both more robust (because, for example, the tabification engine is rather general and does not depend on specific formatting rules, including those I may wish to add later), and more easily extensible. In a sense, I implemented a very small DSL for code formatting.

Settings and preprocessing

Now, the code. First comes the only setting we will have here:

$maxLineLength = 70;

this defines the maximal length of line of code, and is used by the formatter to decide which long lines need dissection.

Next comes the pre-processing function, which serves to remove spaces and tabs possibly existing in the box expression:

ClearAll[preprocess];
preprocess[boxes_] :=
  boxes //.
      {RowBox[{("\t" | "\n") .., expr___}] :> expr} //.
      {
        s_String /; StringMatchQ[s, Whitespace] :> Sequence[],
            RowBox[{r_RowBox}] :> r
      };

Converting boxes to intermediate inert representation

Now, we will define the heads of our inert intermediate representation, to which we want to transform the original box expression:

ClearAll[$blocks, blockQ];
$blocks = {
   CompoundExpressionBlock, GeneralHeadBlock, GeneralBlock, 
   StatementBlock, NewlineBlock, FinalTabBlock, GeneralSplitHeadBlock,   
   SuppressedCompoundExpressionBlock, CommaSeparatedGeneralBlock     
};

blockQ[block_Symbol] :=    MemberQ[$blocks, block];

The following function will translate the box expression into this intermediate language. It is very simplistic and misses many important cases - a more comprehensive one is in the code of the CodeFormatter` - but it illustrates the general structure. Note that is is recursive, moving from outside to inside.

ClearAll[preformat];

preformat[RowBox[elems : {PatternSequence[_, ";"] ..}]] :=
  SuppressedCompoundExpressionBlock @@ Map[
      Map[preformat, StatementBlock @@ DeleteCases[#, ";"]] &,
      Split[elems, # =!= ";" &]];

preformat[RowBox[elems : {PatternSequence[_, ";"] .., _}]] :=
  CompoundExpressionBlock @@ Map[
      Map[preformat, StatementBlock @@ DeleteCases[#, ";"]] &,
      Split[elems, # =!= ";" &]];

preformat[RowBox[elems : {PatternSequence[_, ","] .., _}]] :=
  CommaSeparatedGeneralBlock @@ 
    Map[preformat, DeleteCases[elems, ","]];

preformat[RowBox[elems_List]] /; ! FreeQ[elems, "\n" | "\t", 1] :=
  preformat[RowBox[DeleteCases[elems, "\n" | "\t"]]];

preformat[RowBox[{head_, "[", elems___, "]"}]] :=
  GeneralHeadBlock[preformat@head, 
    Sequence @@ Map[preformat, {elems}]];

preformat[RowBox[elems_List]] :=
  GeneralBlock @@ Map[preformat, elems];

preformat[block_?blockQ[args_]] :=
  block @@ Map[preformat, {args}];

preformat[a_?AtomQ] := a;

preformat[expr_] :=
    Throw[{$Failed, expr}, preformat];

You can see that it treats only few selected heads like CompoundExpression separately, plus has rules for general heads.

Formatting

Next will come two helper functions, used in formatting to determine whether or not a given line of code is too long and needs to be split. The first one, maxLen, determines the maximal length of the code line in an expression, accounting for the fact that it may already have been split into several lines.

Clear[maxLen];
maxLen[boxes : _RowBox ] :=
  Max@Replace[
      Split[Append[Cases[boxes, s_String, Infinity], "\n"], # =!= "\n" &],
      {s___, ("\t" | " ") ..., "\n"} :> 
        Total[{s} /. {"\t" -> 4, ss_ :> StringLength[ss]}],
      {1}];

maxLen[expr_] :=
  With[ {boxes = postformat@expr},
       maxLen[boxes] /; MatchQ[boxes, _RowBox ]
   ];

maxLen[expr_] :=
  Throw[{$Failed, expr}, maxLen];

Note that maxLen uses not yet defined postformat, which is perhaps a bit stronger coupling between components than desirable, and is a design short-cut to be removed. The next one is a simple convenience function:

ClearAll[needSplitQ];
needSplitQ[expr_, currentTab_] :=
  maxLen[expr] > $maxLineLength - currentTab;

Now comes the main formatting function, format. All specific formatting rules are included here. It takes an intermediate inert expression as a first argument, and the current number of tabs inserted, as a second one. It is also essentially recursive, and processing an expression from outside to inside.

ClearAll[format];
format[expr_] :=  format[expr, 0];    

format[TabBlock[expr_], currentTab_] :=
  TabBlock[format[expr, currentTab + 4]];

format[NewlineBlock[expr_, flag_], currentTab_] :=
  NewlineBlock[format[expr, currentTab], flag];    

format[(ce : (CompoundExpressionBlock | 
    SuppressedCompoundExpressionBlock))[elems__], 
    currentTab_] :=
  With[ {formatted = Map[format[#, currentTab] &, {elems}]},
       (ce @@ Map[NewlineBlock[#, False] &, formatted]) /; 
         !FreeQ[formatted, NewlineBlock]
  ];

format[StatementBlock[el_], currentTab_] :=
    StatementBlock[format[el, currentTab]];

format[expr : GeneralHeadBlock[head_, elems___], currentTab_] :=
  With[ {splitQ = needSplitQ[expr, currentTab]},
       GeneralSplitHeadBlock[
           format[head, currentTab],
           Sequence @@ Map[
               format[If[ splitQ,
                             TabBlock@NewlineBlock[#, False],
                             #
                         ], 
                   currentTab] &,
               {elems}]] /; splitQ
   ];

(* For a generic block, it is not obvious that we have to tab, so we don't*)
format[expr : (block_?blockQ[elems___]), currentTab_] :=
  With[ {splitQ = needSplitQ[expr, currentTab]},
       block @@ Map[
           format[If[ splitQ,
                         NewlineBlock[#, False],
                         #
                     ], currentTab] &,
           {elems}]
   ];

format[a_?AtomQ, _] := a;

You can see that format uses two new block types: NewlineBlock and TabBlock, and the former also accepts a flag which can be True or False. This flag, when being set to True, forces the formatter to create a new line, while when being set to False, tells the formatter to propagate the new line request deeper into the expression. The TabBlock directive also accepts a similar flag. The reason that the flags are needed in this approach is that it is not straightforward to implement the abstraction such as "move this piece of code one tab to the right" on the box level, for example because each new line in the boxes must be tabbed separately.

In any case, format does only part of the job, because it only instructs what must be done. It has a companion, tabify, which actually executes the instructions of format:

ClearAll[tabify];
tabify[expr_] /; ! FreeQ[expr, TabBlock[_]] :=
    tabify[expr //. TabBlock[sub_] :> TabBlock[sub, True]];

tabify[(block_?blockQ /; ! MemberQ[{TabBlock, FinalTabBlock}, block])[
elems___]] :=
    block @@ Map[tabify, {elems}];

tabify[TabBlock[FinalTabBlock[el_, flag_], tflag_]] :=
  FinalTabBlock[tabify[TabBlock[el, tflag]], flag];

tabify[TabBlock[NewlineBlock[el_, flag_], _]] :=
  tabify[NewlineBlock[TabBlock[el, True], flag]];

tabify[TabBlock[t_TabBlock, flag_]] :=
  tabify[TabBlock[tabify[t], flag]];

tabify[TabBlock[(block_?blockQ /; ! MemberQ[{TabBlock}, block])[ 
     elems___], flag_]] :=
  FinalTabBlock[
    block @@ Map[tabify@TabBlock[#, False] &, {elems}],
    flag];

tabify[TabBlock[a_?AtomQ, flag_]] :=
  FinalTabBlock[a, flag];

tabify[expr_] :=  expr;

You can see that it introduces another tab-related block, FinalTabBlock - which is a block that signifies the need to tab a particular line by one tab, and is inert in the sense that once TabBlock is converted to FinalTabBlock, it does not any more actively influence the work of tabify.

Post-formatting

The final stage of the formatting procedure is to take the expression processed with format and tabify. We need one helper function which serves to prevent the addition of several new lines ("gaps" in the formatted code) , by determining whether or not the next line of code starts with a new line (if so, the NewlineBlock directive around it is ignored):

ClearAll[isNextNewline];
isNextNewline[_NewlineBlock] := True;

isNextNewline[block : (_?blockQ | TabBlock)[fst_, ___]] :=
  isNextNewline[fst];

isNextNewline[_] := False;

Here is finally the code for the inverse converter from the inert intermediate representation to boxes, postformat:

ClearAll[postformat];
postformat[GeneralBlock[elems__]] :=
  RowBox[postformat /@ {elems}];

postformat[CompoundExpressionBlock[elems__]] :=
  RowBox[Riffle[postformat /@ {elems}, ";"]];

postformat[SuppressedCompoundExpressionBlock[elems__]] :=
  RowBox[Append[Riffle[postformat /@ {elems}, ";"], ";"]];

postformat[GeneralHeadBlock[head_, elems___]] :=
  RowBox[{postformat@head, "[", 
      Sequence @@ Riffle[postformat /@ {elems}, ","], "]"}];

postformat[GeneralSplitHeadBlock[head_, elems___]] :=
  With[ {formattedElems = postformat /@ {elems}},
       RowBox[{postformat@head, "[",
           Sequence @@ Riffle[Most[formattedElems], ","],
           Last[formattedElems], "]"}]
   ];

postformat[GeneralBlock[elems___]] :=
  RowBox[Riffle[postformat /@ {elems}, ","]];

postformat[StatementBlock[elem_]] :=
  postformat[elem];

postformat[NewlineBlock[elem_?isNextNewline, False]] :=
  postformat@elem;

postformat[CommaSeparatedGeneralBlock[elems__]] :=
  RowBox[Riffle[postformat /@ {elems}, ","]];

postformat[NewlineBlock[elem_, _]] :=
  RowBox[{"\n", postformat@elem}];

postformat[FinalTabBlock[expr_, True]] :=
  RowBox[{"\t", postformat@expr}];

postformat[FinalTabBlock[expr_, False]] :=
  postformat@expr;

postformat[a_?AtomQ] :=  a;

postformat[arg_] :=
  Throw[{$Failed, arg}, postformat];

It is also necessarily recursive, and the code should be pretty much self-documenting.

The final function

The function which brings it all together is very simple:

ClearAll[fullCodeFormat];
fullCodeFormat[boxes_] :=
  postformat@tabify@format@preformat@preprocess@boxes;

Examples and limitations

This very simplified version of the formatter is quite limited. However, it can handle a few not so trivial examples. Here is a rather non-trivial one to try:

prn@fullCodeFormat@MakeBoxes[
   Compile[{{data, _Real, 2}}, 
     Module[{means = Table[0., {maxIndex}], num = Table[0, {maxIndex}], 
        ctr = 0, i = 0, index = 0, resultIndices = Table[0, {maxIndex}], 
          indexHash = Table[0, {maxIndex}]}, 
      Do[index = IntegerPart[data[[i, 2]]];
       means[[index]] += data[[i, 1]];
       num[[index]]++;
       If[indexHash[[index]] == 0, indexHash[[index]] = 1;
       resultIndices[[++ctr]] = index];, {i, Length[data]}];
       resultIndices = Take[resultIndices, ctr];
       Transpose[{resultIndices, 
       means[[resultIndices]] + num[[resultIndices]]}]],(*Module*)
       CompilationTarget -> "C", RuntimeOptions -> "Speed"]]

here is a screenshot of what you should see as a result (prn is defined as prn = CellPrint[Cell[BoxData[#], "Input"]] &):

Because of the way the final function is written, it is very easy to inspect what is happening in the intermediate stages. For example, you can only apply

format@preformat@preprocess@boxes

to see what format is doing.

This simplified formatter has a number of limitations, so don't expect it to work nicely on code involving e.g. function definitions through SetDelayed, some complex patterns, etc. The real code formatter in the CodeFormatter` package, while having the same core, has a number of additional rules to handle more cases.

Answer 2

It is worth mentioning GeneralUtilities`PrintDefinitions.

Use it like <<GeneralUtilities` then GeneralUtilities`PrintDefinitions@f to get a pretty-printed, formatted version of ??f in a new notebook (with pink background).

The function returns a handle to the created notebook. One should be able to programmatically extract the plain formatted code from there.

Answer 3

t = CreateTemporary[];
Save[t, mySymbol];
Import[t, "Text"]

applies some basic formatting to the definitions immediately and indirectly associated with mySymbol.

Answer 4

Now here is it:

https://github.com/WolframResearch/codeformatter

before

f[x :_a |_b] :=0

after

f[x : _a | _b] :=
    0

---

a larger example:

With CodeFormatter`$DefaultLineWidth=Infinity;

Manipulate[
Block[{a=1.,n=8,θ=Pi/(2 n),e=Cos[θ],b=a Sin[θ]},
Graphics[{
Circle[{0,0},{a,b}],
Table[{If[Sin@#>Sin[θ],Circle[{-a Cos[θ] Cos[#],0},a Sin[θ] Sin[#]]]&/@
{2 i θ+t,t-2 (i-1) θ}},{i,n-1}]
},PlotRange->1.1a]
],{t,0,Pi/2}]

to

ref https://github.com/WolframResearch/codeformatter/issues/3

Manipulate[
    Block[{a = 1., n = 8, θ = Pi / (2 n), e = Cos[θ], b = a Sin[θ]},
        Graphics[
            {
                Circle[{0, 0}, {a, b}]
                ,
                Table[
                    {
                        If[Sin @ # > Sin[θ],
                                Circle[{-a Cos[θ] Cos[#], 0}, a Sin[θ] Sin[#]]
                            ]& /@ {2 i θ + t, t - 2 (i - 1) θ}
                    }
                    ,
                    {i, n - 1}
                ]
            }
            ,
            PlotRange -> 1.1 a
        ]
    ]
    ,
    {t, 0, Pi / 2}
]