Add command line handling and diff function

astor/__init__.py

 # -*- coding: utf-8 -*-
 """
-Part of the astor library for Python AST manipulation.
+astor is a library and command line tool that is designed
+to manipulate, dump, pretty-print, and compare Python
+abstract syntax trees (ASTs).
 
-License: 3-clause BSD
+astor homepage: https://github.com/berkerpeksag/astor
+
+----
 
-Copyright 2012 (c) Patrick Maupin
-Copyright 2013 (c) Berker Peksag
+License: 3-clause BSD
 
+Copyright 2008      (c) Armin Ronacher
+Copyright 2012-2017 (c) Patrick Maupin
+Copyright 2013-2017 (c) Berker Peksag
 """
 
 import warnings
@@ -21,7 +27,6 @@ from .tree_walk import TreeWalk  # NOQA
 
 __version__ = '0.6'
 
-
 # DEPRECATED!!!
 
 # These aliases support old programs.  Please do not use in future.

astor/command_line.py

+# -*- coding: utf-8 -*-
+"""
+License: 3-clause BSD
+
+Copyright 2017 (c) Patrick Maupin
+
+This supports subcommands by importing every sibling
+submodule, and letting them call `add_command` to
+register themselves.
+
+The main document string should have a header and
+a footer, separated by '----\n'
+"""
+
+import sys
+import os
+
+# Each submodule can place its commands here, indexed
+# by command name
+subcommands = {}
+fixed_path = False
+
+
+class OneCommand(str):
+    """ OneCommand subclasses string for easy sorting.
+        The string should be something that indicates
+        how it should be sorted for display.  Perhaps
+        the command name; perhaps something else.
+
+        Attributes of the class should be:
+
+            - cmd -- what to type to invoke this
+            - shorthelp -- help for global printout
+            - invoke() -- method to invoke command
+    """
+
+
+def add_command(name, invoke, shorthelp, sortname=None):
+    cmd = OneCommand(sortname or name)
+    cmd.cmd = name
+    cmd.shorthelp = shorthelp
+    cmd.invoke = invoke
+    subcommands[name] = cmd
+
+
+def main(args=sys.argv):
+    # Make sure we've been imported properly
+    if __name__ == '__main__':
+        from .command_line import main
+        return main(args)
+
+    # Let all command line utilities register
+    for fname in os.listdir(os.path.dirname(__file__)):
+        if fname.endswith('.py') and not fname.startswith('__'):
+            __import__(fname[:-3], globals(), locals(), [], 1)
+
+    add_command('cprofile', do_profile, """
+        cprofile may be inserted in front of other commands in order
+        to do profiling.
+""", 'zzzz')
+
+    add_command('help', do_help, """
+        help may be used to provide help on sub-commands.
+""")
+
+    # Remove extraneous cruft from start of args for display
+    global fixed_path
+    if args is sys.argv and not fixed_path:
+        fixed_path = True
+        fullname = __name__
+        if fullname.replace('.', os.sep) in args[0]:
+            args[0] = fullname
+        else:
+            args[0] = fullname.rsplit('.', 1)[0]
+
+    # Find first possible command
+    cmds = set(subcommands) & set(args[1:])
+    if not cmds:
+        return main_help(args)
+
+    # Invoke the subcommand
+    cmd = sorted(((args.index(x), x) for x in cmds))[0][-1]
+    args.remove(cmd)
+    args[0] = '%s %s' % (args[0], cmd)
+    subcommands[cmd].invoke(args)
+
+
+def do_profile(args):
+    try:
+        import cProfile
+    except ImportError:
+        import profile as cProfile
+    cProfile.run('main()')
+
+
+def do_help(args):
+    if len(args) > 1:
+        args.append('-h')
+        args[0] = args[0].replace(' help', '')
+        return main(args)
+    main_help(args)
+
+
+def main_help(args):
+    from . import __doc__ as doc
+    name = args[0]
+    args = args[1:]
+    header, footer = normstring(doc).split('----\n')
+    print(header.rstrip('\n'))
+    if set(args) - set('-h --help'.split()):
+        print('\nUnknown command: %s' % ' '.join(args))
+
+    if subcommands:
+        print('\nAvailable subcommands:')
+        for cmd in sorted(subcommands.values()):
+            print('\n    %s %s:\n%s' % (name, cmd.cmd, cmd.shorthelp))
+    print(footer)
+
+
+def normstring(s):
+    return s.replace('\r\n', '\n').replace('\r', '\n')
+
+
+if __name__ == '__main__':
+    main()

astor/rtrip.py

@@ -9,7 +9,7 @@ Copyright (c) 2015 Patrick Maupin
 
 Usage:
 
-    python -m astor.rtrip [readonly] [<source>]
+    python -m astor.rtrip [-r] [<source>]
 
 
 This utility tests round-tripping of Python source to AST
@@ -17,7 +17,7 @@ and back to source.
 
     .. versionadded:: 0.6
 
-If readonly is specified, then the source will be tested,
+If -r is specified, then the source will be tested,
 but no files will be written.
 
 if the source is specified to be "stdin" (without quotes)
@@ -63,11 +63,8 @@ Note 1:
         canonical form.
 
 Note 2:
-        This tool WILL TRASH the tmp_rtrip directory (unless readonly
+        This tool WILL TRASH the tmp_rtrip directory (unless -r
         is specified) -- as far as it is concerned, it OWNS that directory.
-
-Note 3: Why is it "readonly" and not "-r"?  Because python -m slurps
-        all the thingies starting with the dash.
 """
 
 import sys
@@ -75,15 +72,23 @@ import os
 import ast
 import shutil
 import logging
+import optparse
 
-from astor.code_gen import to_source
-from astor.file_util import code_to_ast
-from astor.node_util import (allow_ast_comparison, dump_tree,
-                             strip_tree, fast_compare)
+from .code_gen import to_source
+from .file_util import code_to_ast
+from .node_util import (allow_ast_comparison, dump_tree,
+                        strip_tree, fast_compare)
+from .command_line import add_command
 
 
 dsttree = 'tmp_rtrip'
 
+# TODO:  Remove this workaround once we remove version 2 support
+
+
+def out_prep(s, pre_encoded=(sys.version_info[0] == 2)):
+    return s if pre_encoded else s.encode('utf-8')
+
 
 def convert(srctree, dsttree=dsttree, readonly=False, dumpall=False,
             ignore_exceptions=False, fullcomp=False):
@@ -147,8 +152,8 @@ def convert(srctree, dsttree=dsttree, readonly=False, dumpall=False,
         if not readonly:
             dstfname = os.path.join(dstpath, fname)
             try:
-                with open(dstfname, 'w') as f:
-                    f.write(dsttxt)
+                with open(dstfname, 'wb') as f:
+                    f.write(out_prep(dsttxt))
             except UnicodeEncodeError:
                 badfiles.add(dstfname)
 
@@ -174,13 +179,13 @@ def convert(srctree, dsttree=dsttree, readonly=False, dumpall=False,
             if dumpall or bad:
                 if not readonly:
                     try:
-                        with open(dstfname[:-3] + '.srcdmp', 'w') as f:
-                            f.write(srcdump)
+                        with open(dstfname[:-3] + '.srcdmp', 'wb') as f:
+                            f.write(out_prep(srcdump))
                     except UnicodeEncodeError:
                         badfiles.add(dstfname[:-3] + '.srcdmp')
                     try:
-                        with open(dstfname[:-3] + '.dstdmp', 'w') as f:
-                            f.write(dstdump)
+                        with open(dstfname[:-3] + '.dstdmp', 'wb') as f:
+                            f.write(out_prep(dstdump))
                     except UnicodeEncodeError:
                         badfiles.add(dstfname[:-3] + '.dstdmp')
                 elif dumpall:
@@ -208,80 +213,59 @@ def convert(srctree, dsttree=dsttree, readonly=False, dumpall=False,
     if bad_nodes:
         logging.error('\nERROR -- UNKNOWN NODES STRIPPED: %s' % bad_nodes)
     logging.info('\n')
+    return broken
 
 
-def usage(msg):
-    raise SystemExit(textwrap.dedent("""
-
-        Error: %s
-
-        Usage:
-
-            python -m astor.rtrip [readonly] [<source>]
+def main(args=None):
+    """ args are ignored"""
+    import textwrap
+    import time
 
+    # Configure the option parser
+    usage = """usage: %prog [options] [<source>]
 
         This utility tests round-tripping of Python source to AST
         and back to source.
 
-        If readonly is specified, then the source will be tested,
+        If -r is specified, then the source will be tested,
         but no files will be written.
 
-        if the source is specified to be "stdin" (without quotes)
-        then any source entered at the command line will be compiled
-        into an AST, converted back to text, and then compiled to
-        an AST again, and the results will be displayed to stdout.
-
-        If neither readonly nor stdin is specified, then rtrip
-        will create a mirror directory named tmp_rtrip and will
-        recursively round-trip all the Python source from the source
-        into the tmp_rtrip dir, after compiling it and then reconstituting
-        it through code_gen.to_source.
-
-        If the source is not specified, the entire Python library will be used.
-
-        """) % msg)
-
-
-if __name__ == '__main__':
-    import textwrap
-    import time
-
-    def go():
-        convert(fname, readonly=readonly or dumpall, dumpall=dumpall,
-            ignore_exceptions=ignoreexc, fullcomp=fullcomp and not dumpall)
-
-    args = sys.argv[1:]
+        If -r is not specified, then this will create a mirror directory
+        named tmp_rtrip and will recursively round-trip all the Python
+        source from the source into the tmp_rtrip dir, after compiling
+        it and then reconstituting it through code_gen.to_source.
 
-    flags = []
+        If the source is not specified, the entire Python library will
+        be used.
+        """
 
-    for special in 'readonly astunparse ignoreexc cprofile fullcomp'.split():
-        x = special in args
-        if x:
-            args.remove(special)
-            flags.append(special)
-        globals()[special] = x
+    parser = optparse.OptionParser(usage)
+    parser.add_option("-r", action="store_true", default=False,
+                      help='Read-only -- do not create temp dir')
+    (options, args) = parser.parse_args()
 
+    if len(args) > 1:
+        parser.error("only one source file or module allowed")
     if not args:
         args = [os.path.dirname(textwrap.__file__)]
 
-    if len(args) > 1:
-        usage("Too many arguments")
-
-    if astunparse:
-        from astunparse import unparse as to_source
-
-    if cprofile:
-        import cProfile
-        cprofile = cProfile.run
-
     fname, = args
-    dumpall = False
     if not os.path.exists(fname):
-        dumpall = fname == 'stdin' or usage("Cannot find directory %s" % fname)
+        parser.error("Cannot find directory %s" % fname)
 
     logging.basicConfig(format='%(msg)s', level=logging.INFO)
 
     start = time.time()
-    cprofile("go()") if cprofile else go()
+    convert(fname, readonly=options.r, dumpall=False,
+            ignore_exceptions=False, fullcomp=False)
     print('Finished in %0.2f seconds' % (time.time() - start))
-    print("Used flags %s" %(', '.join(flags) or None))
+
+
+add_command('rtrip', main, """
+        rtrip tests round-tripping of Python source to AST
+        and back to source.
+""")
+
+
+if __name__ == '__main__':
+    main()

astor/src_diff.py

+#! /usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+Part of the astor library for Python AST manipulation.
+
+Original code:  Copyright (c) Python Software Foundation
+
+https://docs.python.org/2/library/difflib.html
+
+Command line interface to difflib.py providing diffs in four formats:
+
+* ndiff:    lists every line and highlights interline changes.
+* context:  highlights clusters of changes in a before/after format.
+* unified:  highlights clusters of changes in an inline format.
+* html:     generates side by side comparison with change highlights.
+
+"""
+
+import sys
+import os
+import time
+import difflib
+import optparse
+
+from .command_line import add_command
+from .file_util import code_to_ast
+from . import to_source
+
+
+def main(args=None):
+    """ args are ignored """
+    # Configure the option parser
+    usage = """usage: %prog [options] fromfile tofile
+
+        %prog can convert two files to ASTs and then
+        back to Python for comparison purposes.  The reconstituted
+        Python is compared; this removes cosmetic differences such
+        as comments and PEP 8 reformatting before comparison.
+    """
+    parser = optparse.OptionParser(usage)
+    parser.add_option("-c", action="store_true", default=False,
+                      help='Produce a context format diff (default)')
+    parser.add_option("-u", action="store_true", default=False,
+                      help='Produce a unified format diff')
+    hlp = 'Produce HTML side by side diff (can use -c and -l in conjunction)'
+    parser.add_option("-m", action="store_true", default=False, help=hlp)
+    parser.add_option("-n", action="store_true", default=False,
+                      help='Produce a ndiff format diff')
+    parser.add_option("-l", "--lines", type="int", default=3,
+                      help='Set number of context lines (default 3)')
+    (options, args) = parser.parse_args()
+
+    if len(args) == 0:
+        parser.print_help()
+        sys.exit(1)
+    if len(args) != 2:
+        parser.error("need to specify both a fromfile and tofile")
+
+    n = options.lines
+    fromfile, tofile = args  # as specified in the usage string
+
+    # we're passing these as arguments to the diff function
+    fromdate = time.ctime(os.stat(fromfile).st_mtime)
+    todate = time.ctime(os.stat(tofile).st_mtime)
+    fromlines = to_source(code_to_ast.parse_file(fromfile)).splitlines(True)
+    tolines = to_source(code_to_ast.parse_file(tofile)).splitlines(True)
+
+    if options.u:
+        diff = difflib.unified_diff(fromlines, tolines, fromfile, tofile,
+                                    fromdate, todate, n=n)
+    elif options.n:
+        diff = difflib.ndiff(fromlines, tolines)
+    elif options.m:
+        diff = difflib.HtmlDiff().make_file(fromlines, tolines, fromfile,
+                                            tofile, context=options.c,
+                                            numlines=n)
+    else:
+        diff = difflib.context_diff(fromlines, tolines, fromfile, tofile,
+                                    fromdate, todate, n=n)
+
+    # we're using writelines because diff is a generator
+    sys.stdout.writelines(diff)
+
+
+add_command('diff', main, """
+        diff compares two python files or packages.
+""")

setup.py

@@ -27,9 +27,10 @@ setup(
         'Programming Language :: Python :: 2.6',
         'Programming Language :: Python :: 2.7',
         'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.2',
         'Programming Language :: Python :: 3.3',
         'Programming Language :: Python :: 3.4',
+        'Programming Language :: Python :: 3.5',
+        'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: Implementation',
         'Programming Language :: Python :: Implementation :: CPython',
         'Programming Language :: Python :: Implementation :: PyPy',
@@ -37,4 +38,7 @@ setup(
         'Topic :: Software Development :: Compilers',
     ],
     keywords='ast, codegen, PEP8',
+    entry_points = {
+        'console_scripts': ['astor=astor.command_line:main'],
+    },
 )