Customize dataclass representations

[cristipp]

• PrettyPrinter version: 0.14.0
• Python version: 3.6.6
• Operating System: MacOs

Description

I work on a compiler for a minilanguage. The worflow depends on printing out the AST at various points in the compilation pipeline, and on copy/pasting said AST as test-case expected values. For grokking purposes, the AST representation should be properly indented. I use a home-baked pprint scheme, but I discovered prettyprinter and I would love to migrate. However, the default representation of dataclasses is too verbose. Is there a good way to customize dataclass representations:

A. A way to pprint short names for dataclasses, Plus(...) instead of python3.tests.playground.Plus(...).

B. A way to pprint the arguments of some dataclasses as args instead of kwargs, Const(1) instead of Const(value=1)

C. A way to keep the indentation of a given dataclass consistent, always multiline or always single line, instead of depending on contextual line length.

What I Did

from dataclasses import dataclass
from prettyprinter import pprint, install_extras

install_extras(
    # Comment out any packages you are not using.
    include=[
        'dataclasses',
    ],
    warn_on_error=True
)

@dataclass
class Expr:
    pass

@dataclass
class Plus(Expr):
    e0: Expr
    e1: Expr

@dataclass
class Times(Expr):
    e0: Expr
    e1: Expr

@dataclass
class Const(Expr):
    value: int

def test_prettyprinter():
    expr = Plus(Const(1), Times(Const(2), Const(3)))
    pprint(expr)

Output:

python3.tests.playground.Plus(
    e0=python3.tests.playground.Const(value=1),
    e1=python3.tests.playground.Times(
        e0=python3.tests.playground.Const(value=2),
        e1=python3.tests.playground.Const(value=3)
    )
)

Desired output:

Plus(
    e0=Const(1),
    e1=Times(
        e0=Const(2),
        e1=Const(3)
    )
)

[cristipp]

To work around A, I did a bit of monkeypatching:

from prettyprinter.prettyprinter import IMPLICIT_MODULES
IMPLICIT_MODULES.add(
    'python3.tests.playground'
)

Output:

Plus(e0=Const(value=1), e1=Times(e0=Const(value=2), e1=Const(value=3)))

This is almost right, though for consistency [easier reading, easier copy/paste in different contexts] I'd prefer:

Plus(
    e0=Const(value=1), 
    e1=Times(
        e0=Const(value=2), 
        e1=Const(value=3)))

Thanks for a great library!

[tommikaikkonen]

You can implement and register your custom prettyprinter for each dataclass you've defined. For example, if the automatic prettyprinter for Const gives python3.tests.playground.Const(value=3), you should be able to get Const(3) in the output after registering the following printer:

from prettyprinter import register_pretty, pretty_call

@register_pretty(Const)
def pretty_const(value, ctx):
    return pretty_call(
        ctx,
        'Const',  # this is documented as accepting only a callable, but a string also works
        value.value
    )

Since this registered printer is more specific in the subclass hierarchy, it will be used instead of the generic dataclass printer.

For more information, see https://prettyprinter.readthedocs.io/en/latest/usage.html#pretty-printing-your-own-types and the API reference for pretty_call and register_pretty.

[tommikaikkonen]

If you want different kinds of indentation than what the library currently offers, I think you'll have to drop down to the layout primitives used internally rather than the much easier to use public functions like pretty_call. If you want to go down this route, you'd probably implement your own version of this function build_fncall:

https://github.com/tommikaikkonen/prettyprinter/blob/master/prettyprinter/prettyprinter.py#L847

If the layout primitives are not enough to express your wanted indentation style, you'd have to modify the layout algorithm in layout.py, probably by supplying it with a different "fitting predicate" function which decides whether a subtree of layout primitives (that is shaped somewhat like the Python data you're rendering) is rendered on the same line or broken up. Right now the predicates look at the line length, like you mentioned.

The layout system and primitives are not the easiest to grok at first, but prettyprinter uses a very similar system as the Prettier library for JavaScript. There should be a lot of practical examples of using the primitives in their codebase in addition to this library. There are links to some whitepapers at the top of layout.py that might also help understand them.