Repository: berkerpeksag/astor

Issue:
Allow registering new `visit_xxx` for `SourceGenerator`
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:

```py
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.

Discussion:
berkerpeksag: > Would it make sense if I added maxline to the function signature and forwarded it to split_line?

Could you open a new issue to discuss this one? There is an open PR to refactor `astor/source_repr.py` and its friends, so maybe we can come up with a better solution.
matham: I just wanted to update that I have not forgotten about the PR, it's just that I ended up using astor quite a bit in new code I wrote, so before I continued with the PR I wanted to get a clearer picture for myself how I was going to use astor to make sure my solution was good approach.

In the end, I just wrote my own version of `to_source` (https://github.com/matham/kivy/blob/ast_compile/kivy/lang/compiler/ast_parse.py#L320) because I needed further customization, so I'm wondering if I should continue with this PR, or if it's better that the user make their own `to_source`, like I did rather than making `to_source` more complex? Especially if the API will be changing as you mentioned?

If the underlying API used in `to_source` is available more clearly, maybe that's better than adding more stuff to `to_source`? But either way, I can add tests etc to this PR or drop it. Whichever you think is better.
berkerpeksag: I think this PR would still help users to customize source generator. I recently came across a blog post using astor (I can't read Chinese, but I can read Python :)): http://blog.soliloquize.org/2017/02/08/%E4%BD%BF%E7%94%A8ast%E7%A7%BB%E9%99%A4%E4%BB%A3%E7%A0%81%E4%B8%AD%E7%9A%84print%E8%AF%AD%E5%8F%A5/

With `source_generator_class`, the latest example can be written without needing to subclass `ast.NodeTransformer`:

```py
import astor

class RemovePrintGenerator(astor.code_gen.SourceGenerator):

    def visit_Call(self, node):
        if node.func.id != 'print':
            super().visit_Call(node)

print(astor.to_source(astor.parse_file('foo.py'), source_generator_class=RemovePrintGenerator))
```

Patch:
diff --git a/astor/__init__.py b/astor/__init__.py
index 3d30a94..481af3c 100644
--- a/astor/__init__.py
+++ b/astor/__init__.py
@@ -11,7 +11,7 @@ Copyright 2013 (c) Berker Peksag
 
 import warnings
 
-from .code_gen import to_source  # NOQA
+from .code_gen import SourceGenerator, to_source  # NOQA
 from .node_util import iter_node, strip_tree, dump_tree  # NOQA
 from .node_util import ExplicitNodeVisitor  # NOQA
 from .file_util import CodeToAst, code_to_ast  # NOQA
diff --git a/astor/code_gen.py b/astor/code_gen.py
index 1cebf59..6ae1ce2 100644
--- a/astor/code_gen.py
+++ b/astor/code_gen.py
@@ -28,7 +28,8 @@ from .source_repr import pretty_source
 
 
 def to_source(node, indent_with=' ' * 4, add_line_information=False,
-              pretty_string=pretty_string, pretty_source=pretty_source):
+              pretty_string=pretty_string, pretty_source=pretty_source,
+              source_generator_class=None):
     """This function can convert a node tree back into python sourcecode.
     This is useful for debugging purposes, especially if you're dealing with
     custom asts not generated by python itself.
@@ -46,9 +47,18 @@ def to_source(node, indent_with=' ' * 4, add_line_information=False,
     of the nodes are added to the output.  This can be used to spot wrong line
     number information of statement nodes.
 
+    `source_generator_class` defaults to `SourceGenerator`, and specifies the
+    class that will be instantiated and used to generate the source code.
+
     """
-    generator = SourceGenerator(indent_with, add_line_information,
-                                pretty_string)
+    if source_generator_class is None:
+        source_generator_class = SourceGenerator
+    elif not isinstance(source_generator_class, SourceGenerator):
+        raise TypeError('source_generator_class should be a subclass of SourceGenerator')
+    elif not callable(source_generator_class):
+        raise TypeError('source_generator_class should be a callable')
+    generator = source_generator_class(
+        indent_with, add_line_information, pretty_string)
     generator.visit(node)
     generator.result.append('\n')
     if set(generator.result[0]) == set('\n'):
diff --git a/docs/changelog.rst b/docs/changelog.rst
index ab11d2a..8058759 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -27,6 +27,17 @@ New features
 .. _`Issue 138`: https://github.com/berkerpeksag/astor/issues/138
 .. _`PR 139`: https://github.com/berkerpeksag/astor/pull/139
 
+* :func:`astor.to_source` now has a *source_generator_class* parameter to
+  customize source code generation.
+  (Reported and fixed by matham in `Issue 113`_ and `PR 114`_.)
+
+.. _`Issue 113`: https://github.com/berkerpeksag/astor/issues/113
+.. _`PR 114`: https://github.com/berkerpeksag/astor/pull/114
+
+* The :class:`~SourceGenerator` class can now be imported from the
+  :mod:`astor` package directly. Previously, the ``astor.code_gen``
+  submodule was needed to be imported.
+
 Bug fixes
 ~~~~~~~~~
 
diff --git a/docs/index.rst b/docs/index.rst
index b009769..f4d06d5 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -129,7 +129,8 @@ Functions
 *********
 
 .. function:: to_source(source, indent_with=' ' * 4, \
-                        add_line_information=False)
+                        add_line_information=False,
+                        source_generator_class=astor.SourceGenerator)
 
     Convert a node tree back into Python source code.
 
@@ -140,6 +141,13 @@ Functions
     of the nodes are added to the output. This can be used to spot wrong line
     number information of statement nodes.
 
+    *source_generator_class* defaults to :class:`astor.SourceGenerator`, and
+    specifies the class that will be instantiated and used to generate the
+    source code.
+
+    .. versionchanged:: 0.8
+       *source_generator_class* was added.
+
 .. function:: codetoast
 .. function:: code_to_ast(codeobj)
 


Test patch:
diff --git a/tests/test_misc.py b/tests/test_misc.py
index 762e54b..36aefa0 100644
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -46,6 +46,29 @@ class PublicAPITestCase(unittest.TestCase):
             'astor.code_gen module instead.'
         )
 
+    def test_to_source_invalid_customize_generator(self):
+        class InvalidGenerator:
+            pass
+
+        node = ast.parse('spam = 42')
+
+        with self.assertRaises(TypeError) as cm:
+            astor.to_source(node, source_generator_class=InvalidGenerator)
+        self.assertEqual(
+            str(cm.exception),
+            'source_generator_class should be a subclass of SourceGenerator',
+        )
+
+        with self.assertRaises(TypeError) as cm:
+            astor.to_source(
+                node,
+                source_generator_class=astor.SourceGenerator(indent_with=' ' * 4),
+            )
+        self.assertEqual(
+            str(cm.exception),
+            'source_generator_class should be a callable',
+        )
+
 
 class FastCompareTestCase(unittest.TestCase):