Allow specifying the generator class (#114)

Fixes #113

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

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'):

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
 ~~~~~~~~~
 

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)
 

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):