Refactor to allow easier code generator subclassing

 - Add a test for subclassing that shows how to add custom
   nodes for comments, as requested in issue #50.

astor/code_gen.py

@@ -19,6 +19,10 @@
 better looking output by removing extraneous parentheses and
 wrapping long lines.
 
+The main entry point is the `SourceGenerator.to_source` class
+method.  For convenience and legacy compatibility, this may
+also be accessed via the module `to_source` variable.
+
 """
 
 import ast
@@ -28,41 +32,6 @@
 from .prettifier import Formatter
 
 
-###################################################################
-#               Main interface
-###################################################################
-
-
-def to_source(node, indent_with=' ' * 4, add_line_information=False,
-              formatter=Formatter):
-    """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.
-
-    It could be that the sourcecode is evaluable when the AST itself is not
-    compilable / evaluable.  The reason for this is that the AST contains
-    some more data than regular sourcecode does, which is dropped during
-    conversion.
-
-    Each level of indentation is replaced with `indent_with`.  Per default
-    this parameter is equal to four spaces as suggested by PEP 8, but it
-    might be adjusted to match the application's styleguide.
-
-    If `add_line_information` is set to `True` comments for the line
-    numbers of the nodes are added to the output.  This can be used
-    to spot wrong line number information of statement nodes.
-
-    The `formatter` parameter should be an object with an `s_lit()`
-    method that will build a representation of a literal string,
-    and an `out_format()` method that will perform any required
-    pretty-printing output transformations on the object.
-    """
-    generator = SourceGenerator(indent_with, add_line_information, formatter)
-    generator.visit(node)
-    generator.out_format(generator.statements)
-    return ''.join(generator.result)
-
-
 ###################################################################
 #               Main class
 ###################################################################
@@ -80,6 +49,40 @@ class SourceGenerator(dict):
     # Goofy from __future__ string handling for 2.7
     using_unicode_literals = False
 
+    ###################################################################
+    #               Main interface
+    ###################################################################
+
+    @classmethod
+    def to_source(cls, node, indent_with=' ' * 4, add_line_information=False,
+                  formatter=Formatter):
+        """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.
+
+        It could be that the sourcecode is evaluable when the AST itself is
+        not compilable / evaluable.  The reason for this is that the AST
+        contains some more data than regular sourcecode does, which is
+        dropped during conversion.
+
+        Each level of indentation is replaced with `indent_with`.  Per
+        default this parameter is equal to four spaces as suggested by
+        PEP 8, but it might be adjusted to match the application's styleguide.
+
+        If `add_line_information` is set to `True` comments for the line
+        numbers of the nodes are added to the output.  This can be used
+        to spot wrong line number information of statement nodes.
+
+        The `formatter` parameter should be an object with an `s_lit()`
+        method that will build a representation of a literal string,
+        and an `out_format()` method that will perform any required
+        pretty-printing output transformations on the object.
+        """
+        self = cls(indent_with, add_line_information, formatter)
+        self.visit(node)
+        self.out_format(self.statements)
+        return ''.join(self.result)
+
     ###################################################################
     #               Top-level tree constructs
     ###################################################################
@@ -745,6 +748,13 @@ def visit(self, node):
         self.write(node)
 
 
+###################################################################
+#               Convenience function
+###################################################################
+
+to_source = SourceGenerator.to_source
+
+
 ###################################################################
 #               Utility functions, classes, and instances
 ###################################################################

astor/prettifier.py

@@ -216,7 +216,7 @@ class Formatter(object):
     end_delim.add('):')
     begin_end_delim = begin_delim | end_delim
 
-    all_statements = set(('@|assert |async for |async def |async with |'
+    all_statements = set(('# |@|assert |async for |async def |async with |'
                           'break|continue|class |del |except|exec |'
                           'elif |else:|for |def |global |if |import |'
                           'from |nonlocal |pass|print |raise|return|'

tests/test_subclass_code_gen.py

@@ -0,0 +1,111 @@
+"""
+Part of the astor library for Python AST manipulation
+
+License: 3-clause BSD
+
+Copyright (c) 2014 Berker Peksag
+Copyright (c) 2015, 2017 Patrick Maupin
+
+Shows an example of subclassing of SourceGenerator
+to insert comment nodes.
+
+"""
+
+import ast
+
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest
+
+try:
+    from test_code_gen import canonical
+except ImportError:
+    from .test_code_gen import canonical
+
+from astor.code_gen import SourceGenerator
+
+
+class CommentCode(object):
+    """  Represents commented out code.
+    """
+    def __init__(self, subnode):
+        self.subnode = subnode
+
+
+class BlockComment(object):
+    """  Represents a block comment.
+    """
+    def __init__(self, text):
+        self.text = text
+
+
+class SourceWithComments(SourceGenerator):
+    """ Subclass the SourceGenerator and add our node.
+        When our node is visited, write the underlying node,
+        and then go back and comment it.
+    """
+
+    def visit_CommentCode(self, node):
+        statements = self.statements
+        index = len(statements)
+        self.write(node.subnode)
+        for index in range(index, len(statements)):
+            mylist = statements[index]
+            if mylist[0].startswith('\n'):
+                continue
+            mylist[0] = '#' + mylist[0]
+
+    def visit_BlockComment(self, node):
+        """ Print a block comment.  This currently
+            handles a single line, but it could be beefed
+            up to handle more, or even consolidated into
+            a single visit_Comment class with the CommentCode
+            visitor, and make decisions based on the type
+            of the node.
+        """
+        self.statement(node, '# ', node.text)
+
+
+class SubclassCodegenCase(unittest.TestCase):
+
+    def test_comment_node(self):
+        """ Strip the comments out of this source, then
+            try to regenerate them.
+        """
+
+        source = canonical("""
+            if 1:
+
+                def sam(a, b, c):
+                    # This is a block comment
+                    x, y, z = a, b, c
+
+            #    def bill(a, b, c):
+            #        x, y, z = a, b, c
+
+                def mary(a, b, c):
+                    x, y, z = a, b, c
+        """)
+
+        # Strip the block comment
+        uncommented_src = source.replace('        #'
+                                         ' This is a block comment\n', '')
+
+        # Uncomment the bill function and generate the AST
+        uncommented_src = uncommented_src.replace('#', '')
+        ast1 = ast.parse(uncommented_src)
+
+        # Modify the AST to comment out the bill function
+        ast1.body[0].body[1] = CommentCode(ast1.body[0].body[1])
+
+        # Add a comment under sam
+        ast1.body[0].body[0].body.insert(0, BlockComment(
+                                         "This is a block comment"))
+        # Assert it round-trips OK
+        dest = canonical(SourceWithComments.to_source(ast1))
+        self.assertEqual(source, dest)
+
+
+if __name__ == '__main__':
+    unittest.main()