Allow registering new visit_xxx for SourceGenerator

[matham]

I ran into an issue where I defined a custom AST node, and when I tried rendering the source with astor.to_source(node), it crashed saying No defined handler for node of type xxx.

I understand that astor would not know how to render some unknown AST node, but it would be nice if we could register a visit function, such that when astor encounters a unknown AST node, instead of crashing it checks the registry for the visit function.

There are at least two reasons for having a custom AST class:

• A user defined a truly new AST class.
• The user created a derived class of some known AST class, e.g. a Attribute. Because obj.__class__.__name__ is used, even though it inherits from some known AST class, astor does not know what to do with it.

I understand why the first use case would not be convincing. But number two should be common enough that allowing to register one AST class name as an alias for another, or alternatively allowing to register an explicit visit function would be useful (or both).

My use case is as follows: I do the tree walk and then replace some nodes with a proxy node. The proxy node refers to the original node. However, later at my convenience, I replace the node the proxy refers to with another node. This allows me to switch out nodes without having to walk the tree. Currently, I monkey patched astor to be able to handle this as follows:

def visit_ASTNodeRef(self, node, *largs, **kwargs):
    return self.visit(node.ref_node, *largs, **kwargs)
SourceGenerator.visit_ASTNodeRef = visit_ASTNodeRef

But, it would be nice if there was an official way of doing this. I can work on a PR if that sounds acceptable.

[berkerpeksag]

Thanks for the report. What do you think about adding a source_generator_class (or generator_class to make it shorter) parameter to astor.to_source() which accepts a subclass of SourceGenerator?

API example:

class CustomSourceGenerator(SourceGenerator):

    def visit_Spam(self, node):
        ...

astor.to_source(node, source_generator_class=CustomSourceGenerator)

[matham]

Your idea sounds perfect - I added a PR for that: #114.